From 67c7ebde258d5302ee224b595521e175d85efbf5 Mon Sep 17 00:00:00 2001 From: Evan Williams Date: Tue, 10 Dec 2019 12:53:15 -0800 Subject: [PATCH] Add instrumentation timing markers for projection mask application As part of the latency instrumentation initiative, this commit adds timing markers for projection mask application code paths. RB=1745924 G=sf-reviewers R=dmessink,fcapponi,ssheng,bsoetarm,crzhang A=kbalasub,bsoetarm --- CHANGELOG | 3 +++ .../response/BatchCreateResponseBuilder.java | 19 +++++++++++++------ .../response/BatchFinderResponseBuilder.java | 12 +++++++++++- .../response/BatchGetResponseBuilder.java | 13 ++++++++++--- .../server/response/BatchResponseBuilder.java | 15 ++++++++++++--- .../response/CollectionResponseBuilder.java | 8 ++++++++ .../response/CreateResponseBuilder.java | 10 ++++++++++ .../server/response/GetResponseBuilder.java | 9 +++++++++ .../PartialUpdateResponseBuilder.java | 10 ++++++++++ .../TestBatchCreateResponseBuilder.java | 4 ++++ .../TestBatchFinderResponseBuilder.java | 2 ++ .../response/TestBatchGetResponseBuilder.java | 8 ++++++-- ...TestBatchPartialUpdateResponseBuilder.java | 2 ++ .../TestBatchUpdateResponseBuilder.java | 7 +++++-- .../TestCollectionResponseBuilder.java | 2 ++ .../response/TestCreateResponseBuilder.java | 3 +++ .../response/TestGetResponseBuilder.java | 2 ++ .../TestPartialUpdateResponseBuilder.java | 2 ++ 18 files changed, 114 insertions(+), 17 deletions(-) diff --git a/CHANGELOG b/CHANGELOG index 0047e7a3a3..fe88cc3861 100644 --- a/CHANGELOG +++ b/CHANGELOG @@ -1,5 +1,8 @@ 28.0.13 ------- +(RB=1745924) +Add instrumentation timing markers for projection mask application + (RB=1745918) Add support for framework latency instrumentation diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchCreateResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchCreateResponseBuilder.java index 8f2982b2f3..7107645f74 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchCreateResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchCreateResponseBuilder.java @@ -16,12 +16,13 @@ package com.linkedin.restli.internal.server.response; - import com.linkedin.data.DataMap; import com.linkedin.data.template.RecordTemplate; import com.linkedin.jersey.api.uri.UriBuilder; import com.linkedin.jersey.api.uri.UriComponent; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.BatchCreateIdResponse; import com.linkedin.restli.common.CreateIdEntityStatus; import com.linkedin.restli.common.CreateIdStatus; @@ -30,18 +31,16 @@ import com.linkedin.restli.common.RestConstants; import com.linkedin.restli.internal.common.ProtocolVersionUtil; import com.linkedin.restli.internal.common.URIParamUtils; -import com.linkedin.restli.internal.server.methods.AnyRecord; import com.linkedin.restli.internal.server.RoutingResult; +import com.linkedin.restli.internal.server.methods.AnyRecord; import com.linkedin.restli.internal.server.util.RestUtils; import com.linkedin.restli.server.BatchCreateKVResult; import com.linkedin.restli.server.BatchCreateResult; -import com.linkedin.restli.server.CreateResponse; import com.linkedin.restli.server.CreateKVResponse; +import com.linkedin.restli.server.CreateResponse; +import com.linkedin.restli.server.ResourceContext; import com.linkedin.restli.server.RestLiResponseData; import com.linkedin.restli.server.RestLiServiceException; -import com.linkedin.restli.server.ResourceContext; - - import java.net.HttpCookie; import java.util.ArrayList; import java.util.List; @@ -122,6 +121,9 @@ public RestLiResponseData buildRestLiResponseData(R } List collectionCreateList = new ArrayList<>(list.getResults().size()); + TimingContextUtil.beginTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + for (CreateKVResponse createKVResponse : list.getResults()) { if (createKVResponse == null) @@ -136,6 +138,7 @@ public RestLiResponseData buildRestLiResponseData(R if (createKVResponse.getError() == null) { DataMap entityData = createKVResponse.getEntity() != null ? createKVResponse.getEntity().data() : null; + final DataMap data = RestUtils.projectFields(entityData, resourceContext.getProjectionMode(), resourceContext.getProjectionMask()); @@ -156,6 +159,10 @@ public RestLiResponseData buildRestLiResponseData(R } } } + + TimingContextUtil.endTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + return new RestLiResponseDataImpl<>(new BatchCreateResponseEnvelope(HttpStatus.S_200_OK, collectionCreateList, true), headers, cookies); } else diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchFinderResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchFinderResponseBuilder.java index 66d4e14864..9dc792699a 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchFinderResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchFinderResponseBuilder.java @@ -21,6 +21,8 @@ import com.linkedin.data.collections.CheckedUtil; import com.linkedin.data.template.RecordTemplate; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.BatchCollectionResponse; import com.linkedin.restli.common.CollectionMetadata; import com.linkedin.restli.common.CollectionResponse; @@ -36,7 +38,6 @@ import com.linkedin.restli.server.ProjectionMode; import com.linkedin.restli.server.RestLiResponseData; import com.linkedin.restli.server.RestLiServiceException; - import java.net.HttpCookie; import java.net.URI; import java.util.ArrayList; @@ -100,6 +101,10 @@ public RestLiResponseData buildRestLiResponseData(R List collectionResponse = new ArrayList<>(criteriaParams.size()); final ResourceContextImpl resourceContext = (ResourceContextImpl) routingResult.getContext(); + + TimingContextUtil.beginTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + for (Object criteriaParam : criteriaParams.values()) { RecordTemplate criteria = new AnyRecord((DataMap) criteriaParam); @@ -114,8 +119,10 @@ public RestLiResponseData buildRestLiResponseData(R //Process paging final CollectionMetadata projectedPaging = buildPaginationMetaData(routingResult, criteria, resourceContext, request, cr); + //Process metadata final AnyRecord projectedCustomMetadata = buildMetaData(cr, resourceContext); + entry = new BatchFinderEntry(elements, projectedPaging, projectedCustomMetadata); } else if (result.getErrors().containsKey(criteria)) @@ -131,6 +138,9 @@ else if (result.getErrors().containsKey(criteria)) collectionResponse.add(entry); } + TimingContextUtil.endTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + return new RestLiResponseDataImpl<>(new BatchFinderResponseEnvelope(HttpStatus.S_200_OK, collectionResponse), headers, cookies); diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchGetResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchGetResponseBuilder.java index 4c766f28ce..cea03aa95b 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchGetResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchGetResponseBuilder.java @@ -16,25 +16,25 @@ package com.linkedin.restli.internal.server.response; - import com.linkedin.data.DataMap; import com.linkedin.data.collections.CheckedUtil; import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.template.SetMode; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.BatchResponse; import com.linkedin.restli.common.EntityResponse; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.common.ProtocolVersion; import com.linkedin.restli.internal.common.URIParamUtils; -import com.linkedin.restli.internal.server.response.BatchResponseEnvelope.BatchResponseEntry; import com.linkedin.restli.internal.server.RoutingResult; import com.linkedin.restli.internal.server.methods.AnyRecord; +import com.linkedin.restli.internal.server.response.BatchResponseEnvelope.BatchResponseEntry; import com.linkedin.restli.internal.server.util.RestUtils; import com.linkedin.restli.server.BatchResult; import com.linkedin.restli.server.RestLiResponseData; import com.linkedin.restli.server.RestLiServiceException; - import java.net.HttpCookie; import java.util.Collections; import java.util.HashMap; @@ -163,6 +163,9 @@ public RestLiResponseData buildRestLiResponseData(Requ // In this case it is OK to swallow this exception and proceed. } + TimingContextUtil.beginTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + Map batchResult = new HashMap<>(entities.size() + serviceErrors.size()); for (Map.Entry entity : entities.entrySet()) { @@ -177,10 +180,14 @@ public RestLiResponseData buildRestLiResponseData(Requ final DataMap projectedData = RestUtils.projectFields(entity.getValue().data(), routingResult.getContext().getProjectionMode(), routingResult.getContext().getProjectionMask()); + AnyRecord anyRecord = new AnyRecord(projectedData); batchResult.put(finalKey, new BatchResponseEntry(statuses.get(entity.getKey()), anyRecord)); } + TimingContextUtil.endTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + for (Map.Entry entity : serviceErrors.entrySet()) { if (entity.getKey() == null || entity.getValue() == null) diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchResponseBuilder.java index e53c6883c5..ef6aaadf29 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/BatchResponseBuilder.java @@ -16,29 +16,30 @@ package com.linkedin.restli.internal.server.response; - import com.linkedin.data.DataMap; import com.linkedin.data.collections.CheckedUtil; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.BatchResponse; -import com.linkedin.restli.internal.server.response.BatchResponseEnvelope.BatchResponseEntry; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.common.ProtocolVersion; import com.linkedin.restli.common.UpdateStatus; import com.linkedin.restli.internal.common.URIParamUtils; import com.linkedin.restli.internal.server.RoutingResult; import com.linkedin.restli.internal.server.methods.AnyRecord; +import com.linkedin.restli.internal.server.response.BatchResponseEnvelope.BatchResponseEntry; import com.linkedin.restli.server.BatchUpdateResult; import com.linkedin.restli.server.ResourceContext; import com.linkedin.restli.server.RestLiResponseData; import com.linkedin.restli.server.RestLiServiceException; import com.linkedin.restli.server.UpdateResponse; - import java.net.HttpCookie; import java.util.HashMap; import java.util.List; import java.util.Map; + /** * This is the base implementation for {@link RestLiResponseBuilder}s for BATCH_UPDATE, BATCH_PARTIAL_UPDATE, and * BATCH_DELETE responses. @@ -110,6 +111,9 @@ public D buildRestLiResponseData(Request request, + routingResult.getResourceMethod()); } + TimingContextUtil.beginTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + Map batchResponseMap = new HashMap<>(); for (Map.Entry entry : results.entrySet()) { @@ -123,11 +127,16 @@ public D buildRestLiResponseData(Request request, if (!serviceErrors.containsKey(entry.getKey())) { Object finalKey = ResponseUtils.translateCanonicalKeyToAlternativeKeyIfNeeded(entry.getKey(), routingResult); + UpdateStatus updateStatus = buildUpdateStatus(routingResult.getContext(), entry.getValue()); + batchResponseMap.put(finalKey, new BatchResponseEntry(entry.getValue().getStatus(), updateStatus)); } } + TimingContextUtil.endTiming(routingResult.getContext().getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + for (Map.Entry entry : serviceErrors.entrySet()) { if (entry.getKey() == null || entry.getValue() == null) diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CollectionResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CollectionResponseBuilder.java index 021a54e2ab..6d5472acc4 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CollectionResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CollectionResponseBuilder.java @@ -21,6 +21,8 @@ import com.linkedin.data.collections.CheckedUtil; import com.linkedin.data.template.RecordTemplate; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.CollectionMetadata; import com.linkedin.restli.common.CollectionResponse; import com.linkedin.restli.common.HttpStatus; @@ -122,6 +124,9 @@ private D buildRestLiResponseData(final Request request, RestUtils.buildMetadata(request.getURI(), resourceContext, routingResult.getResourceMethod(), elements, pageIncrement, totalResults); + TimingContextUtil.beginTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + //PagingMetadata cannot be null at this point so we skip the null check. Notice here that we are using automatic //intentionally since resource methods cannot explicitly project paging. However, it should be noted that client //resource methods have the option of selectively setting the total to null. This happens if a client decides @@ -165,6 +170,9 @@ private D buildRestLiResponseData(final Request request, projectedCustomMetadata = null; } + TimingContextUtil.endTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + return buildResponseData(HttpStatus.S_200_OK, processedElements, projectedPaging, projectedCustomMetadata, headers, cookies); } diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CreateResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CreateResponseBuilder.java index 8f4c6f4d4d..a66be1c5ca 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CreateResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/CreateResponseBuilder.java @@ -22,6 +22,8 @@ import com.linkedin.jersey.api.uri.UriBuilder; import com.linkedin.jersey.api.uri.UriComponent; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.common.IdResponse; import com.linkedin.restli.common.ProtocolVersion; @@ -117,7 +119,15 @@ public RestLiResponseData buildRestLiResponseData(Reques } DataMap entityData = entity.data(); + + TimingContextUtil.beginTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + final DataMap data = RestUtils.projectFields(entityData, resourceContext.getProjectionMode(), resourceContext.getProjectionMask()); + + TimingContextUtil.endTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + idResponse = new AnyRecord(data); // Ideally, we should set an IdEntityResponse to the envelope. But we are keeping AnyRecord // to make sure the runtime object is backwards compatible. diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/GetResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/GetResponseBuilder.java index a102555ac8..92a22dbf96 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/GetResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/GetResponseBuilder.java @@ -20,6 +20,8 @@ import com.linkedin.data.DataMap; import com.linkedin.data.template.RecordTemplate; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.internal.server.RoutingResult; import com.linkedin.restli.internal.server.methods.AnyRecord; @@ -73,9 +75,16 @@ record = (RecordTemplate) result; status = HttpStatus.S_200_OK; } final ResourceContext resourceContext = routingResult.getContext(); + + TimingContextUtil.beginTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + final DataMap data = RestUtils.projectFields(record.data(), resourceContext.getProjectionMode(), resourceContext.getProjectionMask()); + TimingContextUtil.endTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + return new RestLiResponseDataImpl<>(new GetResponseEnvelope(status, new AnyRecord(data)), headers, cookies); } } \ No newline at end of file diff --git a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/PartialUpdateResponseBuilder.java b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/PartialUpdateResponseBuilder.java index 161f5d5a11..a3889ce8ff 100644 --- a/restli-server/src/main/java/com/linkedin/restli/internal/server/response/PartialUpdateResponseBuilder.java +++ b/restli-server/src/main/java/com/linkedin/restli/internal/server/response/PartialUpdateResponseBuilder.java @@ -19,6 +19,8 @@ import com.linkedin.data.DataMap; import com.linkedin.data.template.RecordTemplate; import com.linkedin.r2.message.Request; +import com.linkedin.r2.message.timing.FrameworkTimingKeys; +import com.linkedin.r2.message.timing.TimingContextUtil; import com.linkedin.restli.common.EntityResponse; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.internal.server.ResponseType; @@ -81,7 +83,15 @@ public RestLiResponseData buildRestLiResponseData if (updateEntityResponse.hasEntity()) { DataMap entityData = updateEntityResponse.getEntity().data(); + + TimingContextUtil.beginTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + final DataMap data = RestUtils.projectFields(entityData, resourceContext.getProjectionMode(), resourceContext.getProjectionMask()); + + TimingContextUtil.endTiming(resourceContext.getRawRequestContext(), + FrameworkTimingKeys.SERVER_RESPONSE_RESTLI_PROJECTION_APPLY.key()); + // Returned entity is to be added to the response envelope entityResponse = new EntityResponse<>(data, updateEntityResponse.getEntity().getClass()); } diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchCreateResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchCreateResponseBuilder.java index b7219933a7..4485149179 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchCreateResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchCreateResponseBuilder.java @@ -27,6 +27,7 @@ import com.linkedin.data.transform.filter.request.MaskTree; import com.linkedin.pegasus.generator.examples.Foo; import com.linkedin.pegasus.generator.examples.Fruits; +import com.linkedin.r2.message.RequestContext; import com.linkedin.r2.message.rest.RestRequest; import com.linkedin.r2.message.rest.RestRequestBuilder; import com.linkedin.restli.common.BatchCreateIdResponse; @@ -275,6 +276,7 @@ public void testReturnEntityInBuildRestLiResponseData(Object batchCreateResult, EasyMock.expect(mockContext.isReturnEntityRequested()).andReturn(isReturnEntityRequested); EasyMock.expect(mockContext.getProjectionMode()).andReturn(ProjectionMode.AUTOMATIC); EasyMock.expect(mockContext.getProjectionMask()).andReturn(null); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); ResourceMethodDescriptor mockDescriptor = getMockResourceMethodDescriptor(null); @@ -321,6 +323,7 @@ public void testProjectionInBuildRestLiResponseData() throws URISyntaxException EasyMock.expect(mockContext.isReturnEntityRequested()).andReturn(true); EasyMock.expect(mockContext.getProjectionMode()).andReturn(ProjectionMode.AUTOMATIC); EasyMock.expect(mockContext.getProjectionMask()).andReturn(maskTree); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); ResourceMethodDescriptor mockDescriptor = getMockResourceMethodDescriptor(null); @@ -398,6 +401,7 @@ private static ServerResourceContext getMockKVResourceContext(String altKeyName) EasyMock.expect(mockContext.getProjectionMask()).andReturn(null).atLeastOnce(); Map protocolVersionOnlyHeaders = Collections.singletonMap(RestConstants.HEADER_RESTLI_PROTOCOL_VERSION, AllProtocolVersions.RESTLI_PROTOCOL_2_0_0.getProtocolVersion().toString()); EasyMock.expect(mockContext.getRequestHeaders()).andReturn(protocolVersionOnlyHeaders).atLeastOnce(); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); return mockContext; diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchFinderResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchFinderResponseBuilder.java index 592f27e1c1..14d0378124 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchFinderResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchFinderResponseBuilder.java @@ -8,6 +8,7 @@ import com.linkedin.jersey.api.uri.UriComponent; import com.linkedin.pegasus.generator.examples.Foo; import com.linkedin.pegasus.generator.examples.Fruits; +import com.linkedin.r2.message.RequestContext; import com.linkedin.r2.message.rest.RestRequest; import com.linkedin.r2.message.rest.RestRequestBuilder; import com.linkedin.restli.common.CollectionMetadata; @@ -213,6 +214,7 @@ private static RoutingResult getMockRoutingResult(List criteria, Proto EasyMock.expect(mockContext.getProjectionMask()).andStubReturn(mockMask); EasyMock.expect(mockContext.getMetadataProjectionMask()).andStubReturn(mockMask); EasyMock.expect(mockContext.getMetadataProjectionMode()).andStubReturn(ProjectionMode.MANUAL); + EasyMock.expect(mockContext.getRawRequestContext()).andStubReturn(new RequestContext()); EasyMock.replay(mockContext); diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchGetResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchGetResponseBuilder.java index ee1b5a6efa..3262396766 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchGetResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchGetResponseBuilder.java @@ -27,6 +27,7 @@ import com.linkedin.data.transform.filter.request.MaskTree; import com.linkedin.pegasus.generator.examples.Foo; import com.linkedin.pegasus.generator.examples.Fruits; +import com.linkedin.r2.message.RequestContext; import com.linkedin.restli.common.BatchResponse; import com.linkedin.restli.common.CompoundKey; import com.linkedin.restli.common.ErrorResponse; @@ -202,8 +203,9 @@ public void testContextErrors() Map errors = new HashMap<>(); RestLiServiceException exception = new RestLiServiceException(HttpStatus.S_402_PAYMENT_REQUIRED); errors.put("foo", exception); - EasyMock.expect(context.hasParameter("altkey")).andReturn(false); - EasyMock.expect(context.getBatchKeyErrors()).andReturn(errors); + EasyMock.expect(context.hasParameter("altkey")).andReturn(false).anyTimes(); + EasyMock.expect(context.getBatchKeyErrors()).andReturn(errors).anyTimes(); + EasyMock.expect(context.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(context); RoutingResult routingResult = new RoutingResult(context, null); RestLiResponseData responseData = builder.buildRestLiResponseData(null, @@ -387,6 +389,7 @@ public void testProjectionInBuildRestliResponseData() EasyMock.expect(mockContext.getProjectionMode()).andReturn(ProjectionMode.AUTOMATIC); EasyMock.expect(mockContext.getProjectionMask()).andReturn(maskTree); EasyMock.expect(mockContext.getBatchKeyErrors()).andReturn(Collections.emptyMap()).once(); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); ResourceMethodDescriptor mockDescriptor = getMockResourceMethodDescriptor(null); @@ -425,6 +428,7 @@ private static ServerResourceContext getMockResourceContext(ProtocolVersion prot { EasyMock.expect(mockContext.getParameter(RestConstants.ALT_KEY_PARAM)).andReturn(altKeyName).atLeastOnce(); } + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); return mockContext; } diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchPartialUpdateResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchPartialUpdateResponseBuilder.java index 6689204c62..6fee68ce61 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchPartialUpdateResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchPartialUpdateResponseBuilder.java @@ -20,6 +20,7 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.transform.filter.request.MaskOperation; import com.linkedin.data.transform.filter.request.MaskTree; +import com.linkedin.r2.message.RequestContext; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.common.ResourceMethod; import com.linkedin.restli.common.UpdateEntityStatus; @@ -210,6 +211,7 @@ private static RoutingResult getMockRoutingResult(boolean isReturnEntityRequeste when(mockServerResourceContext.getProjectionMask()).thenReturn(projectionMask); when(mockServerResourceContext.isReturnEntityRequested()).thenReturn(isReturnEntityRequested); when(mockServerResourceContext.getRestliProtocolVersion()).thenReturn(AllProtocolVersions.RESTLI_PROTOCOL_2_0_0.getProtocolVersion()); + when(mockServerResourceContext.getRawRequestContext()).thenReturn(new RequestContext()); ResourceMethodDescriptor mockResourceMethodDescriptor = mock(ResourceMethodDescriptor.class); return new RoutingResult(mockServerResourceContext, mockResourceMethodDescriptor); } diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchUpdateResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchUpdateResponseBuilder.java index afa5416f55..4716aaf494 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchUpdateResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestBatchUpdateResponseBuilder.java @@ -23,6 +23,7 @@ import com.linkedin.data.template.KeyCoercer; import com.linkedin.data.template.RecordTemplate; import com.linkedin.pegasus.generator.examples.Foo; +import com.linkedin.r2.message.RequestContext; import com.linkedin.restli.common.BatchResponse; import com.linkedin.restli.common.CompoundKey; import com.linkedin.restli.common.ErrorResponse; @@ -178,8 +179,9 @@ public void testContextErrors() Map errors = new HashMap<>(); RestLiServiceException exception = new RestLiServiceException(HttpStatus.S_402_PAYMENT_REQUIRED); errors.put("foo", exception); - EasyMock.expect(context.hasParameter("altkey")).andReturn(false); - EasyMock.expect(context.getBatchKeyErrors()).andReturn(errors); + EasyMock.expect(context.hasParameter("altkey")).andReturn(false).anyTimes(); + EasyMock.expect(context.getBatchKeyErrors()).andReturn(errors).anyTimes(); + EasyMock.expect(context.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(context); RoutingResult routingResult = new RoutingResult(context, getMockResourceMethodDescriptor(null)); RestLiResponseData responseData = builder.buildRestLiResponseData(null, @@ -327,6 +329,7 @@ private static ServerResourceContext getMockResourceContext(ProtocolVersion prot { EasyMock.expect(mockContext.getParameter(RestConstants.ALT_KEY_PARAM)).andReturn(altKeyName).atLeastOnce(); } + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); return mockContext; } diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCollectionResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCollectionResponseBuilder.java index b299ae3e97..6259ed7e33 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCollectionResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCollectionResponseBuilder.java @@ -24,6 +24,7 @@ import com.linkedin.data.transform.filter.request.MaskTree; import com.linkedin.pegasus.generator.examples.Foo; import com.linkedin.pegasus.generator.examples.Fruits; +import com.linkedin.r2.message.RequestContext; import com.linkedin.r2.message.rest.RestRequest; import com.linkedin.r2.message.rest.RestRequestBuilder; import com.linkedin.restli.common.CollectionMetadata; @@ -298,6 +299,7 @@ private static ServerResourceContext getMockResourceContext(MaskTree dataMaskTre ServerResourceContext mockContext = EasyMock.createMock(ServerResourceContext.class); EasyMock.expect(mockContext.getParameter(EasyMock.anyObject())).andReturn(null).times(2); EasyMock.expect(mockContext.getRequestHeaders()).andReturn(ResponseBuilderUtil.getHeaders()).once(); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); //Field Projection EasyMock.expect(mockContext.getProjectionMode()).andReturn(dataProjectionMode).times(generateTestList().size()); diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCreateResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCreateResponseBuilder.java index 280cf47fa7..40f82aac66 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCreateResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestCreateResponseBuilder.java @@ -27,6 +27,7 @@ import com.linkedin.data.transform.filter.request.MaskTree; import com.linkedin.pegasus.generator.examples.Foo; import com.linkedin.pegasus.generator.examples.Fruits; +import com.linkedin.r2.message.RequestContext; import com.linkedin.r2.message.rest.RestRequest; import com.linkedin.r2.message.rest.RestRequestBuilder; import com.linkedin.restli.common.CompoundKey; @@ -189,6 +190,7 @@ public void testReturnEntityInBuildRestLiResponseData(CreateResponse createRespo EasyMock.expect(mockContext.isReturnEntityRequested()).andReturn(isReturnEntityRequested); EasyMock.expect(mockContext.getProjectionMask()).andReturn(null); EasyMock.expect(mockContext.getProjectionMode()).andReturn(ProjectionMode.AUTOMATIC); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); RoutingResult routingResult = new RoutingResult(mockContext, null); @@ -256,6 +258,7 @@ public void testProjectionInBuildRestLiResponseData() throws URISyntaxException EasyMock.expect(mockContext.isReturnEntityRequested()).andReturn(true); EasyMock.expect(mockContext.getProjectionMask()).andReturn(maskTree); EasyMock.expect(mockContext.getProjectionMode()).andReturn(ProjectionMode.AUTOMATIC); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.replay(mockContext); RoutingResult routingResult = new RoutingResult(mockContext, null); diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestGetResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestGetResponseBuilder.java index 6b0f1033a0..027dfc076d 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestGetResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestGetResponseBuilder.java @@ -24,6 +24,7 @@ import com.linkedin.data.transform.filter.request.MaskTree; import com.linkedin.pegasus.generator.examples.Foo; import com.linkedin.pegasus.generator.examples.Fruits; +import com.linkedin.r2.message.RequestContext; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.internal.server.RoutingResult; import com.linkedin.restli.internal.server.ServerResourceContext; @@ -145,6 +146,7 @@ private static ResourceMethodDescriptor getMockResourceMethodDescriptor() private static ServerResourceContext getMockResourceContext(MaskTree maskTree, ProjectionMode projectionMode) { ServerResourceContext mockContext = EasyMock.createMock(ServerResourceContext.class); + EasyMock.expect(mockContext.getRawRequestContext()).andReturn(new RequestContext()).anyTimes(); EasyMock.expect(mockContext.getProjectionMode()).andReturn(projectionMode).once(); EasyMock.expect(mockContext.getProjectionMask()).andReturn(maskTree).once(); EasyMock.replay(mockContext); diff --git a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestPartialUpdateResponseBuilder.java b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestPartialUpdateResponseBuilder.java index ce755380df..f57b6511bd 100644 --- a/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestPartialUpdateResponseBuilder.java +++ b/restli-server/src/test/java/com/linkedin/restli/internal/server/response/TestPartialUpdateResponseBuilder.java @@ -20,6 +20,7 @@ import com.linkedin.data.template.RecordTemplate; import com.linkedin.data.transform.filter.request.MaskOperation; import com.linkedin.data.transform.filter.request.MaskTree; +import com.linkedin.r2.message.RequestContext; import com.linkedin.restli.common.HttpStatus; import com.linkedin.restli.common.ResourceMethod; import com.linkedin.restli.internal.server.ResponseType; @@ -171,6 +172,7 @@ private static RoutingResult getMockRoutingResult(boolean isReturnEntityRequeste when(mockServerResourceContext.getProjectionMode()).thenReturn(ProjectionMode.AUTOMATIC); when(mockServerResourceContext.getProjectionMask()).thenReturn(projectionMask); when(mockServerResourceContext.isReturnEntityRequested()).thenReturn(isReturnEntityRequested); + when(mockServerResourceContext.getRawRequestContext()).thenReturn(new RequestContext()); ResourceMethodDescriptor mockResourceMethodDescriptor = mock(ResourceMethodDescriptor.class); return new RoutingResult(mockServerResourceContext, mockResourceMethodDescriptor); }