From 115d7c1342a4ea59cce2c624229631c825ce62cf Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 8 Jul 2024 15:08:52 +0000 Subject: [PATCH 1/9] Added metrics and tracing for ReadMany operations --- .../com/azure/cosmos/ClientMetricsTest.java | 105 ++++++++++ .../com/azure/cosmos/CosmosTracerTest.java | 29 ++- .../azure/cosmos/CosmosAsyncContainer.java | 25 ++- .../implementation/DiagnosticsProvider.java | 182 +++++++++++++++--- .../azure/cosmos/util/CosmosPagedFlux.java | 71 ++----- 5 files changed, 336 insertions(+), 76 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java index 873bf2c7e017..6bac2386701e 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java @@ -31,6 +31,7 @@ import com.azure.cosmos.models.CosmosBulkExecutionOptions; import com.azure.cosmos.models.CosmosBulkOperations; import com.azure.cosmos.models.CosmosClientTelemetryConfig; +import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosItemOperation; import com.azure.cosmos.models.CosmosItemRequestOptions; import com.azure.cosmos.models.CosmosItemResponse; @@ -41,6 +42,8 @@ import com.azure.cosmos.models.CosmosMicrometerMetricsOptions; import com.azure.cosmos.models.CosmosPatchItemRequestOptions; import com.azure.cosmos.models.CosmosQueryRequestOptions; +import com.azure.cosmos.models.CosmosReadManyRequestOptions; +import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.ModelBridgeInternal; import com.azure.cosmos.models.PartitionKey; import com.azure.cosmos.models.SqlQuerySpec; @@ -408,6 +411,95 @@ public void readItem() throws Exception { } } + @Test(groups = { "fast" }, timeOut = TIMEOUT) + public void readManySingleItem() throws Exception { + this.beforeTest(CosmosMetricCategory.DEFAULT); + try { + InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString()); + container.createItem(properties); + + List tuplesToBeRead = new ArrayList<>(); + tuplesToBeRead.add(new CosmosItemIdentity( + new PartitionKey(properties.get("mypk")), + properties.getId() + )); + + FeedResponse readManyResponse = container.readMany( + tuplesToBeRead, + new CosmosReadManyRequestOptions(), + InternalObjectNode.class); + validateReadManyFeedResponse(List.of(properties), readManyResponse); + + this.validateMetrics( + Tag.of(TagName.OperationStatusCode.toString(), "200"), + Tag.of(TagName.RequestStatusCode.toString(), "200/0"), + 0, + 500 + ); + + this.validateMetrics( + Tag.of( + TagName.Operation.toString(), "Document/Query/readMany"), + Tag.of(TagName.RequestOperationType.toString(), "Document/Read"), + 0, + 500 + ); + + Tag queryPlanTag = Tag.of(TagName.RequestOperationType.toString(), "DocumentCollection_QueryPlan"); + this.assertMetrics("cosmos.client.req.gw", false, queryPlanTag); + this.assertMetrics("cosmos.client.req.rntbd", false, queryPlanTag); + } finally { + this.afterTest(); + } + } + + @Test(groups = { "fast" }, timeOut = TIMEOUT) + public void readManyMultipleItems() throws Exception { + this.beforeTest(CosmosMetricCategory.DEFAULT); + + List createdDocs = new ArrayList<>(); + List tuplesToBeRead = new ArrayList<>(); + try { + for (int i = 0; i < 20; i++) { + InternalObjectNode properties = getDocumentDefinition(UUID.randomUUID().toString()); + container.createItem(properties); + createdDocs.add(properties); + tuplesToBeRead.add(new CosmosItemIdentity( + new PartitionKey(properties.get("mypk")), + properties.getId() + )); + } + + + FeedResponse readManyResponse = container.readMany( + tuplesToBeRead, + new CosmosReadManyRequestOptions(), + InternalObjectNode.class); + validateReadManyFeedResponse(createdDocs, readManyResponse); + + this.validateMetrics( + Tag.of(TagName.OperationStatusCode.toString(), "200"), + Tag.of(TagName.RequestStatusCode.toString(), "200/0"), + 0, + 500 + ); + + this.validateMetrics( + Tag.of( + TagName.Operation.toString(), "Document/Query/readMany"), + Tag.of(TagName.RequestOperationType.toString(), "Document/Query"), + 0, + 500 + ); + + Tag queryPlanTag = Tag.of(TagName.RequestOperationType.toString(), "DocumentCollection_QueryPlan"); + this.assertMetrics("cosmos.client.req.gw", false, queryPlanTag); + this.assertMetrics("cosmos.client.req.rntbd", false, queryPlanTag); + } finally { + this.afterTest(); + } + } + private void runReadItemTestWithThresholds( CosmosDiagnosticsThresholds thresholds, boolean expectRequestMetrics @@ -1351,6 +1443,19 @@ private void validateItemResponse(InternalObjectNode containerProperties, .isEqualTo(containerProperties.getId()); } + private void validateReadManyFeedResponse( + List createdDocs, + FeedResponse readManyResponse) { + // Basic validation + assertThat(readManyResponse).isNotNull(); + assertThat(readManyResponse.getResults()).isNotNull(); + List docsFromResponse = readManyResponse.getResults(); + assertThat(docsFromResponse).hasSize(createdDocs.size()); + for (InternalObjectNode doc: createdDocs) { + assertThat(docsFromResponse.stream().anyMatch(r -> r.getId() != null && r.getId().equals(doc.getId()))); + } + } + private void validateItemCountMetrics(Tag expectedOperationTag) { if (this.getEffectiveMetricCategories().contains(MetricCategory.OperationDetails)) { this.assertMetrics("cosmos.client.op.maxItemCount", true, expectedOperationTag); diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 6992716f3078..30fb8418f3ec 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -15,6 +15,7 @@ import com.azure.cosmos.implementation.SerializationDiagnosticsContext; import com.azure.cosmos.implementation.TestConfigurations; import com.azure.cosmos.implementation.Utils; +import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.ShowQueryMode; import com.azure.cosmos.implementation.directconnectivity.ReflectionUtils; import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics; @@ -74,6 +75,7 @@ import java.time.Instant; import java.time.OffsetDateTime; import java.time.ZoneOffset; +import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.HashMap; @@ -616,10 +618,13 @@ public void cosmosAsyncContainer( samplingRate); mockTracer.reset(); + List createdDocs = new ArrayList<>(); for (int i = 0; i < 30; i++) { // inserting high enough number of documents to make sure we have at least 1 doc on each // of the two partitions - item = getDocumentDefinition(ITEM_ID + "_" + i); + String id = ITEM_ID + "_" + i; + createdDocs.add(new CosmosItemIdentity(new PartitionKey(id), id)); + item = getDocumentDefinition(id); requestOptions = new CosmosItemRequestOptions(); cosmosItemResponse = cosmosAsyncContainer .createItem(item, requestOptions) @@ -883,7 +888,27 @@ public void cosmosAsyncContainer( forceThresholdViolations, null, samplingRate); - mockTracer.reset(); + mockTracer.reset(); + + // read many single item + feedItemResponse = cosmosAsyncContainer.readMany(List.of(createdDocs.get(0)), ObjectNode.class) + .block(); + assertThat(feedItemResponse).isNotNull(); + verifyTracerAttributes( + mockTracer, + "readManyItems." + cosmosAsyncContainer.getId(), + cosmosAsyncDatabase.getId(), + cosmosAsyncContainer.getId(), + feedItemResponse.getCosmosDiagnostics(), + null, + useLegacyTracing, + enableRequestLevelTracing, + showQueryMode, + query, + forceThresholdViolations, + "readMany", + samplingRate); + mockTracer.reset(); } @Test(groups = { "fast", "simple" }, dataProvider = "traceTestCaseProvider", timeOut = TIMEOUT) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java index d110a0bc5986..21fcc64bfa5a 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncContainer.java @@ -1526,6 +1526,15 @@ public Mono> readMany( CosmosReadManyRequestOptions requestOptions, Class classType) { + return withContext(context -> this.readManyInternal(itemIdentityList, requestOptions, classType, context)); + } + + private Mono> readManyInternal( + List itemIdentityList, + CosmosReadManyRequestOptions requestOptions, + Class classType, + Context context) { + CosmosQueryRequestOptions queryRequestOptions = requestOptions == null ? new CosmosQueryRequestOptions() : queryOptionsAccessor.clone(readManyOptionsAccessor.getImpl(requestOptions)); @@ -1539,7 +1548,7 @@ public Mono> readMany( fluxOptions.setMaxItemCount(itemIdentityList != null ? itemIdentityList.size() : 0); QueryFeedOperationState state = new QueryFeedOperationState( client, - this.readAllItemsSpanName, + this.readManyItemsSpanName, database.getId(), this.getId(), ResourceType.Document, @@ -1549,9 +1558,21 @@ public Mono> readMany( fluxOptions ); - return CosmosBridgeInternal + Mono> responseMono = CosmosBridgeInternal .getAsyncDocumentClient(this.getDatabase()) .readMany(itemIdentityList, BridgeInternal.getLink(this), state, classType); + + RequestOptions options = queryOptionsAccessor.toRequestOptions(queryRequestOptions); + + return client + .getDiagnosticsProvider() + .traceEnabledReadManyResponsePublisher( + itemIdentityList, + state, + responseMono, + client, + options, + context); } /** diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java index 698f3b918cdd..67cb19b6ae6b 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java @@ -16,14 +16,17 @@ import com.azure.cosmos.CosmosDiagnosticsHandler; import com.azure.cosmos.CosmosDiagnosticsThresholds; import com.azure.cosmos.CosmosException; +import com.azure.cosmos.CosmosItemSerializer; import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics; import com.azure.cosmos.implementation.directconnectivity.StoreResultDiagnostics; import com.azure.cosmos.implementation.guava25.base.Splitter; import com.azure.cosmos.implementation.query.QueryInfo; import com.azure.cosmos.models.CosmosBatchResponse; import com.azure.cosmos.models.CosmosClientTelemetryConfig; +import com.azure.cosmos.models.CosmosItemIdentity; import com.azure.cosmos.models.CosmosItemResponse; import com.azure.cosmos.models.CosmosResponse; +import com.azure.cosmos.models.FeedResponse; import com.azure.cosmos.models.ShowQueryMode; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; @@ -41,6 +44,7 @@ import java.io.PrintWriter; import java.io.StringWriter; import java.time.Duration; +import java.time.Instant; import java.time.OffsetDateTime; import java.time.ZoneOffset; import java.util.ArrayList; @@ -54,8 +58,11 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.atomic.AtomicLong; import java.util.function.BiFunction; +import java.util.function.Consumer; import java.util.function.Function; +import java.util.function.Supplier; import static com.azure.cosmos.implementation.RequestTimeline.EventName.CREATED; import static com.azure.cosmos.implementation.guava25.base.Preconditions.checkArgument; @@ -97,6 +104,8 @@ public final class DiagnosticsProvider { private final CosmosClientTelemetryConfig telemetryConfig; private final boolean shouldSystemExitOnError; + final Supplier samplingRateSnapshotSupplier; + public DiagnosticsProvider( CosmosClientTelemetryConfig clientTelemetryConfig, @@ -110,6 +119,11 @@ public DiagnosticsProvider( checkNotNull(connectionMode, "Argument 'connectionMode' must not be null."); this.telemetryConfig = clientTelemetryConfig; + + this.samplingRateSnapshotSupplier = () -> isEnabled() + ? clientTelemetryConfigAccessor.getSamplingRate(this.telemetryConfig) + : 0; + this.diagnosticHandlers = new ArrayList<>( clientTelemetryConfigAccessor.getDiagnosticHandlers(clientTelemetryConfig)); Tracer tracerCandidate = clientTelemetryConfigAccessor.getOrCreateTracer(clientTelemetryConfig); @@ -541,7 +555,8 @@ public > Mono traceEnabledCosmosResponsePublisher return diagnostics; }, - requestOptions); + requestOptions, + null); } public Mono traceEnabledBatchResponsePublisher( @@ -585,7 +600,8 @@ public Mono traceEnabledBatchResponsePublishe return diagnostics; }, - requestOptions); + requestOptions, + null); } public Mono> traceEnabledCosmosItemResponsePublisher( @@ -630,7 +646,126 @@ public Mono> traceEnabledCosmosItemResponsePublisher( return diagnostics; }, - requestOptions); + requestOptions, + null); + } + + private Mono> wrapReadManyFeedResponseWithTracingIfEnabled( + CosmosAsyncClient client, + List itemIdentityList, + FeedOperationState state, + Mono> publisher, + RequestOptions requestOptions, + Context context) { + + final double samplingRateSnapshot = this.samplingRateSnapshotSupplier.get(); + final boolean isSampledOut = this.shouldSampleOutOperation(samplingRateSnapshot); + final CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot(); + + if (ctx == null || isSampledOut) { + return publisher; + } + + return publisherWithDiagnostics( + publisher, + context, + state.getSpanName(), + ctx.getContainerName(), + ctx.getDatabaseName(), + ctx.getAccountName(), + client, + ctx.getEffectiveConsistencyLevel(), + ctxAccessor.getOperationType(ctx), + ctxAccessor.getResourceType(ctx), + null, + itemIdentityList.size(), + (r) -> HttpConstants.StatusCodes.OK, // FeedResponse would only ever be created in success case + (r) -> r.getResults().size(), + (r) -> r.getRequestCharge(), + (r, samplingRate) -> { + CosmosDiagnostics diagnostics = r.getCosmosDiagnostics(); + if (diagnostics != null) { + diagnosticsAccessor.setSamplingRateSnapshot(diagnostics, samplingRate); + } + + return diagnostics; + }, + requestOptions, + ctx + ); + } + + private static boolean isTracerEnabled(DiagnosticsProvider tracerProvider) { + return tracerProvider != null; + } + + public static void recordFeedResponse( + Consumer> feedResponseConsumer, + FeedOperationState state, + Supplier samplingRateSnapshotSupplier, + DiagnosticsProvider tracerProvider, + FeedResponse response, + AtomicLong feedResponseConsumerLatencyInNanos) { + + CosmosDiagnostics diagnostics = response != null ? response.getCosmosDiagnostics() : null; + + Integer actualItemCount = response != null && response.getResults() != null ? + response.getResults().size() : null; + + if (diagnostics != null && + diagnosticsAccessor + .isDiagnosticsCapturedInPagedFlux(diagnostics) + .compareAndSet(false, true)) { + + Double samplingRateSnapshot = samplingRateSnapshotSupplier.get(); + if (samplingRateSnapshot != null && samplingRateSnapshot < 1) { + diagnosticsAccessor + .setSamplingRateSnapshot(diagnostics, samplingRateSnapshot); + } + + if (isTracerEnabled(tracerProvider)) { + tracerProvider.recordPage( + state.getDiagnosticsContextSnapshot(), + diagnostics, + actualItemCount, + response.getRequestCharge()); + } + + // If the user has passed feedResponseConsumer, then call it with each feedResponse + if (feedResponseConsumer != null) { + // NOTE this call is happening in a span counted against client telemetry / metric latency + // So, the latency of the user's callback is accumulated here to correct the latency + // reported to client telemetry and client metrics + Instant feedResponseConsumerStart = Instant.now(); + feedResponseConsumer.accept(response); + feedResponseConsumerLatencyInNanos.addAndGet( + Duration.between(Instant.now(), feedResponseConsumerStart).toNanos()); + } + } + } + + public Mono> traceEnabledReadManyResponsePublisher( + List identities, + FeedOperationState state, + Mono> resultPublisher, + CosmosAsyncClient client, + RequestOptions requestOptions, + Context context) { + + checkNotNull(resultPublisher, "Argument 'resultPublisher' must not be null."); + checkNotNull(state, "Argument 'state' must not be null."); + checkNotNull(requestOptions, "Argument 'requestOptions' must not be null."); + checkNotNull(client, "Argument 'client' must not be null."); + + String accountName = clientAccessor.getAccountTagValue(client); + + return wrapReadManyFeedResponseWithTracingIfEnabled( + client, + identities, + state, + resultPublisher, + requestOptions, + context); } /** @@ -740,30 +875,33 @@ private Mono publisherWithDiagnostics(Mono resultPublisher, Function actualItemCountFunc, Function requestChargeFunc, BiFunction diagnosticFunc, - RequestOptions requestOptions) { + RequestOptions requestOptions, + CosmosDiagnosticsContext cosmosCtxFromUpstream) { CosmosDiagnosticsThresholds thresholds = requestOptions != null ? clientAccessor.getEffectiveDiagnosticsThresholds(client, requestOptions.getDiagnosticsThresholds()) : clientAccessor.getEffectiveDiagnosticsThresholds(client, null); - CosmosDiagnosticsContext cosmosCtx = ctxAccessor.create( - spanName, - accountName, - BridgeInternal.getServiceEndpoint(client), - databaseId, - containerId, - resourceType, - operationType, - null, - clientAccessor.getEffectiveConsistencyLevel(client, operationType, consistencyLevel), - maxItemCount, - thresholds, - trackingId, - clientAccessor.getConnectionMode(client), - clientAccessor.getUserAgent(client), - null, - null, - requestOptions); + CosmosDiagnosticsContext cosmosCtx = cosmosCtxFromUpstream != null + ? cosmosCtxFromUpstream + : ctxAccessor.create( + spanName, + accountName, + BridgeInternal.getServiceEndpoint(client), + databaseId, + containerId, + resourceType, + operationType, + null, + clientAccessor.getEffectiveConsistencyLevel(client, operationType, consistencyLevel), + maxItemCount, + thresholds, + trackingId, + clientAccessor.getConnectionMode(client), + clientAccessor.getUserAgent(client), + null, + null, + requestOptions); if (requestOptions != null) { requestOptions.setDiagnosticsContextSupplier(() -> cosmosCtx); diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java index 968b3bd07d6a..43a41bfab8da 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/util/CosmosPagedFlux.java @@ -7,7 +7,6 @@ import com.azure.core.util.FluxUtil; import com.azure.core.util.IterableStream; import com.azure.core.util.paging.ContinuablePagedFlux; -import com.azure.cosmos.CosmosDiagnostics; import com.azure.cosmos.CosmosDiagnosticsContext; import com.azure.cosmos.implementation.CosmosPagedFluxOptions; import com.azure.cosmos.implementation.DiagnosticsProvider; @@ -41,8 +40,6 @@ * @see FeedResponse */ public final class CosmosPagedFlux extends ContinuablePagedFlux> { - private final static ImplementationBridgeHelpers.CosmosDiagnosticsHelper.CosmosDiagnosticsAccessor cosmosDiagnosticsAccessor = - ImplementationBridgeHelpers.CosmosDiagnosticsHelper.getCosmosDiagnosticsAccessor(); private static final ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.CosmosDiagnosticsContextAccessor ctxAccessor = ImplementationBridgeHelpers.CosmosDiagnosticsContextHelper.getCosmosDiagnosticsContextAccessor(); @@ -159,7 +156,13 @@ private Flux> wrapWithTracingIfEnabled(CosmosPagedFluxOptions p switch (signal.getType()) { case ON_COMPLETE: case ON_NEXT: - this.recordFeedResponse(pagedFluxOptions, tracerProvider, response, feedResponseConsumerLatencyInNanos); + DiagnosticsProvider.recordFeedResponse( + feedResponseConsumer, + pagedFluxOptions.getFeedOperationState(), + () ->pagedFluxOptions.getSamplingRateSnapshot(), + tracerProvider, + response, + feedResponseConsumerLatencyInNanos); break; default: break; @@ -185,7 +188,13 @@ private Flux> wrapWithTracingIfEnabled(CosmosPagedFluxOptions p switch (signal.getType()) { case ON_COMPLETE: if (response != null) { - this.recordFeedResponse(pagedFluxOptions, tracerProvider, response, feedResponseConsumerLatencyInNanos); + DiagnosticsProvider.recordFeedResponse( + feedResponseConsumer, + pagedFluxOptions.getFeedOperationState(), + () ->pagedFluxOptions.getSamplingRateSnapshot(), + tracerProvider, + response, + feedResponseConsumerLatencyInNanos); } state.mergeDiagnosticsContext(); @@ -203,7 +212,13 @@ private Flux> wrapWithTracingIfEnabled(CosmosPagedFluxOptions p break; case ON_NEXT: - this.recordFeedResponse(pagedFluxOptions, tracerProvider, response, feedResponseConsumerLatencyInNanos); + DiagnosticsProvider.recordFeedResponse( + feedResponseConsumer, + pagedFluxOptions.getFeedOperationState(), + () ->pagedFluxOptions.getSamplingRateSnapshot(), + tracerProvider, + response, + feedResponseConsumerLatencyInNanos); state.mergeDiagnosticsContext(); CosmosDiagnosticsContext ctxSnapshotOnNext = state.getDiagnosticsContextSnapshot(); ctxAccessor @@ -279,47 +294,7 @@ private Flux> wrapWithTracingIfEnabled(CosmosPagedFluxOptions p )); } - private void recordFeedResponse( - CosmosPagedFluxOptions pagedFluxOptions, - DiagnosticsProvider tracerProvider, - FeedResponse response, - AtomicLong feedResponseConsumerLatencyInNanos) { - CosmosDiagnostics diagnostics = response != null ? response.getCosmosDiagnostics() : null; - - Integer actualItemCount = response != null && response.getResults() != null ? - response.getResults().size() : null; - - if (diagnostics != null && - cosmosDiagnosticsAccessor - .isDiagnosticsCapturedInPagedFlux(diagnostics) - .compareAndSet(false, true)) { - - if (pagedFluxOptions.getSamplingRateSnapshot() < 1) { - cosmosDiagnosticsAccessor - .setSamplingRateSnapshot(diagnostics, pagedFluxOptions.getSamplingRateSnapshot()); - } - - if (isTracerEnabled(tracerProvider)) { - tracerProvider.recordPage( - pagedFluxOptions.getFeedOperationState().getDiagnosticsContextSnapshot(), - diagnostics, - actualItemCount, - response.getRequestCharge()); - } - - // If the user has passed feedResponseConsumer, then call it with each feedResponse - if (feedResponseConsumer != null) { - // NOTE this call is happening in a span counted against client telemetry / metric latency - // So, the latency of the user's callback is accumulated here to correct the latency - // reported to client telemetry and client metrics - Instant feedResponseConsumerStart = Instant.now(); - feedResponseConsumer.accept(response); - feedResponseConsumerLatencyInNanos.addAndGet( - Duration.between(Instant.now(), feedResponseConsumerStart).toNanos()); - } - } - } private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Context context) { AtomicReference startTime = new AtomicReference<>(); @@ -339,10 +314,6 @@ private Flux> byPage(CosmosPagedFluxOptions pagedFluxOptions, Co return result; } - private boolean isTracerEnabled(DiagnosticsProvider tracerProvider) { - return tracerProvider != null; - } - /////////////////////////////////////////////////////////////////////////////////////////// // the following helper/accessor only helps to access this class outside of this package.// /////////////////////////////////////////////////////////////////////////////////////////// From 6257f934d0f4ed7d73c86ee43ede6db7b29375b4 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 8 Jul 2024 15:26:18 +0000 Subject: [PATCH 2/9] Update CHANGELOG.md --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 6eb61c0546d2..bfb7a19cca5f 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -9,6 +9,7 @@ #### Bugs Fixed #### Other Changes +* Added metrics and tracing for ReadMany operations. - See [PR 41042](https://github.com/Azure/azure-sdk-for-java/pull/41042) ### 4.62.0 (2024-07-02) From 59fb212a099b8429f565a5fa6608371894d57220 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 8 Jul 2024 18:47:26 +0000 Subject: [PATCH 3/9] Removing List.of (not compatible with Java 8) --- .../src/test/java/com/azure/cosmos/ClientMetricsTest.java | 3 ++- .../src/test/java/com/azure/cosmos/CosmosTracerTest.java | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java index 6bac2386701e..27a1e4456001 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/ClientMetricsTest.java @@ -63,6 +63,7 @@ import java.net.URI; import java.time.Duration; import java.util.ArrayList; +import java.util.Arrays; import java.util.Collections; import java.util.EnumSet; import java.util.List; @@ -428,7 +429,7 @@ public void readManySingleItem() throws Exception { tuplesToBeRead, new CosmosReadManyRequestOptions(), InternalObjectNode.class); - validateReadManyFeedResponse(List.of(properties), readManyResponse); + validateReadManyFeedResponse(Arrays.asList(properties), readManyResponse); this.validateMetrics( Tag.of(TagName.OperationStatusCode.toString(), "200"), diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 30fb8418f3ec..07951e1c47cb 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -891,7 +891,7 @@ public void cosmosAsyncContainer( mockTracer.reset(); // read many single item - feedItemResponse = cosmosAsyncContainer.readMany(List.of(createdDocs.get(0)), ObjectNode.class) + feedItemResponse = cosmosAsyncContainer.readMany(Arrays.asList(createdDocs.get(0)), ObjectNode.class) .block(); assertThat(feedItemResponse).isNotNull(); verifyTracerAttributes( From 1a7558fd8f9ef1bef340a355254c2a06f9bc76d4 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 8 Jul 2024 23:04:30 +0200 Subject: [PATCH 4/9] Update CHANGELOG.md --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index bfb7a19cca5f..7dd17df5375a 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -1,6 +1,6 @@ ## Release History -### 4.63.0-beta.1 (Unreleased) +### 4.63.0-beta.1 (Unreleased) #### Features Added From dfa69d5412f5e525e710fb90f5874deb2a9d7192 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Mon, 8 Jul 2024 23:04:44 +0200 Subject: [PATCH 5/9] Update CHANGELOG.md --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 7dd17df5375a..bfb7a19cca5f 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -1,6 +1,6 @@ ## Release History -### 4.63.0-beta.1 (Unreleased) +### 4.63.0-beta.1 (Unreleased) #### Features Added From 999f9c9ec5cb0ccc42f902985be2fddd3a2e06c2 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 9 Jul 2024 08:47:50 +0000 Subject: [PATCH 6/9] Update NonStreamingOrderByQueryVectorSearchTest.java --- .../cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java index 13187b97c0d1..07058ee66ee9 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java @@ -123,14 +123,14 @@ public void flatIndexVectorSearch() { .collectList().block(); validateOrdering(6, resultDocs, false); - String euclideanSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], {'distanceFunction': 'euclidean'}) AS " + + String euclideanSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'euclidean'}) AS " + "score FROM c ORDER BY VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'euclidean'})", queryVector, queryVector); resultDocs = flatIndexContainer.queryItems(euclideanSpecsQuery, new CosmosQueryRequestOptions(), Document.class).byPage() .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) .collectList().block(); validateOrdering(6, resultDocs, false); - String dotproductSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], {'distanceFunction': 'dotproduct'}) AS " + + String dotproductSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'}) AS " + "score FROM c ORDER BY VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'})", queryVector, queryVector); resultDocs = flatIndexContainer.queryItems(dotproductSpecsQuery, new CosmosQueryRequestOptions(), Document.class).byPage() .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) @@ -156,7 +156,7 @@ public void quantizedIndexVectorSearch() { .collectList().block(); validateOrdering(6, resultDocs, true); - String dotproduct_specs_query = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], {'distanceFunction': 'dotproduct'}) AS " + + String dotproduct_specs_query = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'}) AS " + "score FROM c ORDER BY VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'})", queryVector, queryVector); resultDocs = quantizedIndexContainer.queryItems(dotproduct_specs_query, new CosmosQueryRequestOptions(), Document.class).byPage() .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) From b8d9f1dd973c2ebcb40ff13df5cedbcbe6088b31 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 9 Jul 2024 11:09:47 +0000 Subject: [PATCH 7/9] Update NonStreamingOrderByQueryVectorSearchTest.java --- .../cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java index 07058ee66ee9..082393aae831 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/rx/NonStreamingOrderByQueryVectorSearchTest.java @@ -128,7 +128,7 @@ public void flatIndexVectorSearch() { resultDocs = flatIndexContainer.queryItems(euclideanSpecsQuery, new CosmosQueryRequestOptions(), Document.class).byPage() .flatMap(feedResponse -> Flux.fromIterable(feedResponse.getResults())) .collectList().block(); - validateOrdering(6, resultDocs, false); + validateOrdering(6, resultDocs, true); String dotproductSpecsQuery = String.format("SELECT DISTINCT TOP 6 c.text, VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'}) AS " + "score FROM c ORDER BY VectorDistance(c.embedding, [%s], false, {'distanceFunction': 'dotproduct'})", queryVector, queryVector); From f0167225cf5d4e42ba7e68c06d2fe6518e241a92 Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 9 Jul 2024 12:08:11 +0000 Subject: [PATCH 8/9] Update CosmosTracerTest.java --- .../src/test/java/com/azure/cosmos/CosmosTracerTest.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java index 07951e1c47cb..b36d8d888664 100644 --- a/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java +++ b/sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosTracerTest.java @@ -1470,10 +1470,11 @@ private void verifyOTelTracerAttributes( String dbStatement = (String) attributes.get("db.statement"); - if (ShowQueryMode.ALL.equals(showQueryMode)) { + boolean isReadMany = "readMany".equals(cosmosDiagnostics.getDiagnosticsContext().getOperationId()); + if (!isReadMany && ShowQueryMode.ALL.equals(showQueryMode)) { assertThat(attributes.get("db.statement")).isNotNull(); assertThat(attributes.get("db.statement")).isEqualTo(queryStatement); - } else if (ShowQueryMode.PARAMETERIZED_ONLY.equals(showQueryMode) + } else if (!isReadMany && ShowQueryMode.PARAMETERIZED_ONLY.equals(showQueryMode) && null != dbStatement && dbStatement.contains("@")) { assertThat(attributes.get("db.statement")).isNotNull(); assertThat(attributes.get("db.statement")).isEqualTo(queryStatement); From 9e9c6998ca0aa7f2c48396b399119b770fad352a Mon Sep 17 00:00:00 2001 From: Fabian Meiswinkel Date: Tue, 9 Jul 2024 12:53:52 +0000 Subject: [PATCH 9/9] Update DiagnosticsProvider.java --- .../azure/cosmos/implementation/DiagnosticsProvider.java | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java index 67cb19b6ae6b..f45e210edaf2 100644 --- a/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java +++ b/sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/implementation/DiagnosticsProvider.java @@ -661,9 +661,16 @@ private Mono> wrapReadManyFeedResponseWithTracingIfEnabled( final double samplingRateSnapshot = this.samplingRateSnapshotSupplier.get(); final boolean isSampledOut = this.shouldSampleOutOperation(samplingRateSnapshot); final CosmosDiagnosticsContext ctx = state.getDiagnosticsContextSnapshot(); + ctxAccessor.setSamplingRateSnapshot(ctx, samplingRateSnapshot, isSampledOut); if (ctx == null || isSampledOut) { - return publisher; + return publisher.map(r -> { + CosmosDiagnostics diagnostics = r.getCosmosDiagnostics(); + if (diagnostics != null) { + diagnosticsAccessor.setSamplingRateSnapshot(diagnostics, samplingRateSnapshot); + } + return r; + }); } return publisherWithDiagnostics(