Skip to content

Commit

Permalink
Fixed issue where CosmosDiagnosticsContext is null when diagnostics a…
Browse files Browse the repository at this point in the history 
…re sampled out. (#39352)

* Fixed issue where CosmosDiagnosticsContext is null when diagnostics are sampled out.

* Update CHANGELOG.md

* Update DiagnosticsProvider.java

* Fixed sampling for tracing

* Making sure Ctx exists when Tracer is disabled
  • Loading branch information
@FabianMeiswinkel
FabianMeiswinkel authored Mar 23, 2024
1 parent c71ae66 commit b445345
Show file tree
Hide file tree
Showing 8 changed files with 189 additions and 70 deletions.
102 changes: 89 additions & 13 deletions sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosDiagnosticsE2ETest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
import java.util.List;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;

import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
Expand Down Expand Up @@ -132,6 +133,7 @@ public void onlyDefaultLogger() {
CosmosContainer container = this.getContainer(builder);
executeTestCase(container);


// no assertions here - invocations for diagnostics handler are validated above
// log4j event logging isn't validated in general in unit tests because it is too brittle to do so
// with custom appender
Expand All @@ -152,8 +154,33 @@ public void onlyLoggerWithCustomConfig() {
.diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER)
);
CosmosContainer container = this.getContainer(builder);
executeTestCase(container);
CosmosItemResponse<ObjectNode> response = executeTestCase(container);
assertThat(response.getDiagnostics()).isNotNull();
assertThat(response.getDiagnostics().getDiagnosticsContext()).isNotNull();
// no assertions here - invocations for diagnostics handler are validated above
// log4j event logging isn't validated in general in unit tests because it is too brittle to do so
// with custom appender
}

@Test(groups = { "fast", "emulator" }, timeOut = TIMEOUT)
public void onlyLoggerAlwaysSampledOut() {
CosmosClientBuilder builder = this
.getClientBuilder()
.clientTelemetryConfig(
new CosmosClientTelemetryConfig()
.diagnosticsThresholds(
new CosmosDiagnosticsThresholds()
.setPointOperationLatencyThreshold(Duration.ofMillis(100))
.setNonPointOperationLatencyThreshold(Duration.ofMillis(2000))
.setRequestChargeThreshold(100)
)
.diagnosticsHandler(CosmosDiagnosticsHandler.DEFAULT_LOGGING_HANDLER)
.sampleDiagnostics(0)
);
CosmosContainer container = this.getContainer(builder);
CosmosItemResponse<ObjectNode> response = executeTestCase(container);
assertThat(response.getDiagnostics()).isNotNull();
assertThat(response.getDiagnostics().getDiagnosticsContext()).isNotNull();
// no assertions here - invocations for diagnostics handler are validated above
// log4j event logging isn't validated in general in unit tests because it is too brittle to do so
// with custom appender
Expand All @@ -177,10 +204,46 @@ public void onlyCustomLoggerWithCustomConfig() {
.diagnosticsHandler(capturingLogger)
);
CosmosContainer container = this.getContainer(builder);
executeTestCase(container);
String id = UUID.randomUUID().toString();
ObjectNode doc = getDocumentDefinition(id);

CosmosDiagnostics diagnostics = executeDocumentOperation(container, OperationType.Create, id, doc);
assertThat(diagnostics).isNotNull();
assertThat(diagnostics.getDiagnosticsContext()).isNotNull();

diagnostics = executeDocumentOperation(container, OperationType.Query, id, doc);
assertThat(diagnostics).isNotNull();
assertThat(diagnostics.getDiagnosticsContext()).isNotNull();

assertThat(capturingLogger.getLoggedMessages()).isNotNull();
assertThat(capturingLogger.getLoggedMessages()).hasSize(1);
assertThat(capturingLogger.getLoggedMessages()).hasSize(2);
}

@Test(groups = { "fast", "emulator" }, timeOut = TIMEOUT)
public void onlyCustomLoggerAlwaysSampledOut() {

CapturingLogger capturingLogger = new CapturingLogger();

CosmosClientBuilder builder = this
.getClientBuilder()
.clientTelemetryConfig(
new CosmosClientTelemetryConfig()
.diagnosticsThresholds(
new CosmosDiagnosticsThresholds()
.setPointOperationLatencyThreshold(Duration.ofMillis(100))
.setNonPointOperationLatencyThreshold(Duration.ofMillis(2000))
.setRequestChargeThreshold(100)
)
.diagnosticsHandler(capturingLogger)
.sampleDiagnostics(0)
);
CosmosContainer container = this.getContainer(builder);
CosmosItemResponse<ObjectNode> response = executeTestCase(container);

assertThat(response.getDiagnostics()).isNotNull();
assertThat(response.getDiagnostics().getDiagnosticsContext()).isNotNull();
assertThat(capturingLogger.getLoggedMessages()).isNotNull();
assertThat(capturingLogger.getLoggedMessages()).hasSize(0);
}

@Test(groups = { "fast", "emulator" }, timeOut = TIMEOUT)
Expand Down Expand Up @@ -368,14 +431,16 @@ public void run() {
assertThat(cosmosDiagnosticsNode.get("jvmFatalErrorMapperExecutionCount").asLong()).isGreaterThan(0);
}

private void executeTestCase(CosmosContainer container) {
private CosmosItemResponse<ObjectNode> executeTestCase(CosmosContainer container) {
String id = UUID.randomUUID().toString();
CosmosItemResponse<ObjectNode> response = container.createItem(
getDocumentDefinition(id),
new PartitionKey(id),
null);
new CosmosItemRequestOptions());
assertThat(response).isNotNull();
assertThat(response.getStatusCode()).isEqualTo(201);

return response;
}

private ObjectNode getDocumentDefinition(String documentId) {
Expand Down Expand Up @@ -407,47 +472,58 @@ private CosmosContainer getContainer(CosmosClientBuilder builder) {
return this.client.getDatabase(asyncContainer.getDatabase().getId()).getContainer(asyncContainer.getId());
}

private void executeDocumentOperation(
private CosmosDiagnostics executeDocumentOperation(
CosmosContainer cosmosContainer,
OperationType operationType,
String createdItemId,
ObjectNode createdItem) {

final AtomicReference<CosmosDiagnostics> diagnostics = new AtomicReference<>(null);

switch (operationType) {
case Query:
String query = String.format("SELECT * from c");
CosmosQueryRequestOptions queryRequestOptions = new CosmosQueryRequestOptions();
queryRequestOptions.setFeedRange(FeedRange.forLogicalPartition(new PartitionKey(createdItemId)));
Iterable<FeedResponse<JsonNode>> results = cosmosContainer.queryItems(query, queryRequestOptions, JsonNode.class).iterableByPage();
results.forEach(t -> {});
Iterable<FeedResponse<JsonNode>> queryResults = cosmosContainer.queryItems(query, queryRequestOptions, JsonNode.class).iterableByPage();
queryResults.forEach(t -> diagnostics.set(t.getCosmosDiagnostics()));
break;
case ReadFeed:
CosmosChangeFeedRequestOptions changeFeedRequestOptions = CosmosChangeFeedRequestOptions
.createForProcessingFromBeginning(FeedRange.forFullRange());
cosmosContainer.queryChangeFeed(changeFeedRequestOptions, JsonNode.class).iterableByPage();
Iterable<FeedResponse<JsonNode>> changeFeedResults =
cosmosContainer.queryChangeFeed(changeFeedRequestOptions, JsonNode.class).iterableByPage();
changeFeedResults.forEach(t -> diagnostics.set(t.getCosmosDiagnostics()));
break;
case Read:
cosmosContainer
CosmosItemResponse<JsonNode> readResponse = cosmosContainer
.readItem(createdItemId, new PartitionKey(createdItemId), JsonNode.class);
diagnostics.set(readResponse.getDiagnostics());
break;
case Replace:
cosmosContainer
CosmosItemResponse<JsonNode> replaceResponse = cosmosContainer
.replaceItem(createdItem, createdItemId, new PartitionKey(createdItemId), new CosmosItemRequestOptions());
diagnostics.set(replaceResponse.getDiagnostics());
break;
case Delete:
try {
cosmosContainer.deleteItem(getDocumentDefinition(UUID.randomUUID().toString()), new CosmosItemRequestOptions());
CosmosItemResponse<Object> deleteResponse = cosmosContainer.deleteItem(getDocumentDefinition(UUID.randomUUID().toString()), new CosmosItemRequestOptions());
diagnostics.set(deleteResponse.getDiagnostics());
} catch (CosmosException e) {
if (!Exceptions.isNotFound(e)) {
throw e;
}
}
break;
case Create:
cosmosContainer.createItem(getDocumentDefinition(UUID.randomUUID().toString()));
CosmosItemResponse<JsonNode> createResponse = cosmosContainer.createItem(getDocumentDefinition(UUID.randomUUID().toString()));
diagnostics.set(createResponse.getDiagnostics());
break;
default:
throw new IllegalArgumentException("The operation type is not supported");
}

return diagnostics.get();
}

private static class CapturingDiagnosticsHandler implements CosmosDiagnosticsHandler {
Expand Down
5 changes: 3 additions & 2 deletions sdk/cosmos/azure-cosmos/CHANGELOG.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,9 @@

#### Bugs Fixed
* Suppress exceptions when calling diagnostics handlers. - See [PR 39077](https://github.com/Azure/azure-sdk-for-java/pull/39077)
* Fixed an issue where no cross region retry for write operations due to channel acquisition timeout - See [PR 39255](https://github.com/Azure/azure-sdk-for-java/pull/39255)

* Fixed an issue where no cross region retry for write operations due to channel acquisition timeout. - See [PR 39255](https://github.com/Azure/azure-sdk-for-java/pull/39255)
* Fixed incorrect container tag value in metrics. - See [PR 39322](https://github.com/Azure/azure-sdk-for-java/pull/39322)
* Fixed issue where CosmosDiagnosticsContext is null when diagnostics are sampled out. - See [PR 39352](https://github.com/Azure/azure-sdk-for-java/pull/39352)
#### Other Changes
* Only call System.exit in `DiagnosticsProvider` for `Error` scenario. Also add `System.err` for `Error` cases. - See [PR 39077](https://github.com/Azure/azure-sdk-for-java/pull/39077)

Expand Down
14 changes: 11 additions & 3 deletions sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosDiagnosticsContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,8 @@ public final class CosmosDiagnosticsContext {

private Double samplingRateSnapshot;

private boolean isSampledOut;

private ArrayList<CosmosDiagnosticsRequestInfo> requestInfo = null;

private final Integer sequenceNumber;
Expand Down Expand Up @@ -128,6 +130,7 @@ public final class CosmosDiagnosticsContext {
this.userAgent = userAgent;
this.connectionMode = connectionMode;
this.sequenceNumber = sequenceNumber;
this.isSampledOut = false;
}

/**
Expand Down Expand Up @@ -556,15 +559,20 @@ void recordOperation(int statusCode,
}
}

void setSamplingRateSnapshot(double samplingRate) {
void setSamplingRateSnapshot(double samplingRate, boolean isSampledOut) {
synchronized (this.spanName) {
this.samplingRateSnapshot = samplingRate;
this.isSampledOut = isSampledOut;
for (CosmosDiagnostics d : this.diagnostics) {
diagAccessor.setSamplingRateSnapshot(d, samplingRate);
}
}
}

boolean isSampledOut() {
return this.isSampledOut;
}

String getRequestDiagnostics() {
ObjectNode ctxNode = mapper.createObjectNode();

Expand Down Expand Up @@ -1014,9 +1022,9 @@ public String getSpanName(CosmosDiagnosticsContext ctx) {
}

@Override
public void setSamplingRateSnapshot(CosmosDiagnosticsContext ctx, double samplingRate) {
public void setSamplingRateSnapshot(CosmosDiagnosticsContext ctx, double samplingRate, boolean isSampledOut) {
checkNotNull(ctx, "Argument 'ctx' must not be null.");
ctx.setSamplingRateSnapshot(samplingRate);
ctx.setSamplingRateSnapshot(samplingRate, isSampledOut);
}

@Override
Expand Down
4 changes: 2 additions & 2 deletions ...os/azure-cosmos/src/main/java/com/azure/cosmos/implementation/CosmosPagedFluxOptions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,9 +102,9 @@ public double getSamplingRateSnapshot() {
return samplingRateSnapshot;
}

public void setSamplingRateSnapshot(double samplingRateSnapshot) {
public void setSamplingRateSnapshot(double samplingRateSnapshot, boolean isSampledOut) {
if (this.operationState != null) {
this.operationState.setSamplingRateSnapshot(samplingRateSnapshot);
this.operationState.setSamplingRateSnapshot(samplingRateSnapshot, isSampledOut);
}
}
}
Loading

0 comments on commit b445345

Please sign in to comment.