Merging back Tracer PR without conditional check on withContext()

sdk/cosmos/azure-cosmos/pom.xml

@@ -128,7 +128,18 @@ Licensed under the MIT License.
         </exclusion>
       </exclusions>
     </dependency>
-
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-core-tracing-opentelemetry</artifactId>
+      <scope>test</scope>
+      <version>1.0.0-beta.5</version> <!-- {x-version-update;com.azure:azure-core-tracing-opentelemetry;dependency} -->
+      <exclusions>
+        <exclusion>
+          <groupId>io.netty</groupId>
+          <artifactId>netty-tcnative-boringssl-static</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
     <dependency>
       <groupId>com.fasterxml.jackson.module</groupId>
       <artifactId>jackson-module-afterburner</artifactId>

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncClient.java

@@ -4,16 +4,19 @@
 
 import com.azure.core.annotation.ServiceClient;
 import com.azure.core.credential.AzureKeyCredential;
+import com.azure.core.util.Context;
+import com.azure.core.util.tracing.Tracer;
 import com.azure.cosmos.implementation.AsyncDocumentClient;
 import com.azure.cosmos.implementation.Configs;
 import com.azure.cosmos.implementation.ConnectionPolicy;
 import com.azure.cosmos.implementation.CosmosAuthorizationTokenResolver;
 import com.azure.cosmos.implementation.Database;
 import com.azure.cosmos.implementation.HttpConstants;
+import com.azure.cosmos.implementation.TracerProvider;
 import com.azure.cosmos.implementation.directconnectivity.rntbd.RntbdMetrics;
-import com.azure.cosmos.models.CosmosDatabaseResponse;
 import com.azure.cosmos.models.CosmosDatabaseProperties;
 import com.azure.cosmos.models.CosmosDatabaseRequestOptions;
+import com.azure.cosmos.models.CosmosDatabaseResponse;
 import com.azure.cosmos.models.CosmosPermissionProperties;
 import com.azure.cosmos.models.CosmosQueryRequestOptions;
 import com.azure.cosmos.models.ModelBridgeInternal;
@@ -27,7 +30,9 @@
 
 import java.io.Closeable;
 import java.util.List;
+import java.util.ServiceLoader;
 
+import static com.azure.core.util.FluxUtil.withContext;
 import static com.azure.cosmos.implementation.Utils.setContinuationTokenAndMaxItemCount;
 
 /**
@@ -51,6 +56,7 @@ public final class CosmosAsyncClient implements Closeable {
     private final AzureKeyCredential credential;
     private final boolean sessionCapturingOverride;
     private final boolean enableTransportClientSharing;
+    private final TracerProvider tracerProvider;
     private final boolean contentResponseOnWriteEnabled;
 
     CosmosAsyncClient(CosmosClientBuilder builder) {
@@ -65,6 +71,7 @@ public final class CosmosAsyncClient implements Closeable {
         this.sessionCapturingOverride = builder.isSessionCapturingOverrideEnabled();
         this.enableTransportClientSharing = builder.isConnectionSharingAcrossClientsEnabled();
         this.contentResponseOnWriteEnabled = builder.isContentResponseOnWriteEnabled();
+        this.tracerProvider = new TracerProvider(ServiceLoader.load(Tracer.class));
         this.asyncDocumentClient = new AsyncDocumentClient.Builder()
                                        .withServiceEndpoint(this.serviceEndpoint)
                                        .withMasterKeyOrResourceToken(this.keyOrResourceToken)
@@ -194,8 +201,9 @@ boolean isContentResponseOnWriteEnabled() {
      * @return a {@link Mono} containing the cosmos database response with the created or existing database or
      * an error.
      */
-    Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProperties databaseProperties) {
-        return createDatabaseIfNotExistsInternal(getDatabase(databaseProperties.getId()));
+    public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProperties databaseProperties) {
+        return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(databaseProperties.getId()),
+            null, context));
     }
 
     /**
@@ -209,21 +217,7 @@ Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(CosmosDatabaseProperties
      * an error.
      */
     public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(String id) {
-        return createDatabaseIfNotExistsInternal(getDatabase(id));
-    }
-
-    private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database) {
-        return database.read().onErrorResume(exception -> {
-            final Throwable unwrappedException = Exceptions.unwrap(exception);
-            if (unwrappedException instanceof CosmosException) {
-                final CosmosException cosmosException = (CosmosException) unwrappedException;
-                if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
-                    return createDatabase(new CosmosDatabaseProperties(database.getId()),
-                        new CosmosDatabaseRequestOptions());
-                }
-            }
-            return Mono.error(unwrappedException);
-        });
+        return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id), null, context));
     }
 
     /**
@@ -240,19 +234,8 @@ private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosAsy
      * @return the mono.
      */
     public Mono<CosmosDatabaseResponse> createDatabaseIfNotExists(String id, ThroughputProperties throughputProperties) {
-        return this.getDatabase(id).read().onErrorResume(exception -> {
-            final Throwable unwrappedException = Exceptions.unwrap(exception);
-            if (unwrappedException instanceof CosmosException) {
-                final CosmosException cosmosException = (CosmosException) unwrappedException;
-                if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
-                    CosmosDatabaseRequestOptions options = new CosmosDatabaseRequestOptions();
-                    ModelBridgeInternal.setThroughputProperties(options, throughputProperties);
-                    return createDatabase(new CosmosDatabaseProperties(id),
-                        options);
-                }
-            }
-            return Mono.error(unwrappedException);
-        });
+        return withContext(context -> createDatabaseIfNotExistsInternal(getDatabase(id),
+            throughputProperties, context));
     }
 
     /**
@@ -274,9 +257,8 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
         }
         Database wrappedDatabase = new Database();
         wrappedDatabase.setId(databaseProperties.getId());
-        return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options))
-                   .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse))
-                   .single();
+        final CosmosDatabaseRequestOptions requestOptions = options;
+        return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context));
     }
 
     /**
@@ -331,9 +313,8 @@ public Mono<CosmosDatabaseResponse> createDatabase(CosmosDatabaseProperties data
         ModelBridgeInternal.setThroughputProperties(options, throughputProperties);
         Database wrappedDatabase = new Database();
         wrappedDatabase.setId(databaseProperties.getId());
-        return asyncDocumentClient.createDatabase(wrappedDatabase, ModelBridgeInternal.toRequestOptions(options))
-                   .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse))
-                   .single();
+        final CosmosDatabaseRequestOptions requestOptions = options;
+        return withContext(context -> createDatabaseInternal(wrappedDatabase, requestOptions, context));
     }
 
     /**
@@ -397,12 +378,14 @@ public Mono<CosmosDatabaseResponse> createDatabase(String id, ThroughputProperti
      */
     CosmosPagedFlux<CosmosDatabaseProperties> readAllDatabases(CosmosQueryRequestOptions options) {
         return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> {
+            String spanName = "readAllDatabases";
+            pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null);
             setContinuationTokenAndMaxItemCount(pagedFluxOptions, options);
             return getDocClientWrapper().readDatabases(options)
-                                        .map(response ->
-                                            BridgeInternal.createFeedResponse(
-                                                ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()),
-                                                response.getResponseHeaders()));
+                .map(response ->
+                    BridgeInternal.createFeedResponse(
+                        ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()),
+                        response.getResponseHeaders()));
         });
     }
 
@@ -432,7 +415,7 @@ public CosmosPagedFlux<CosmosDatabaseProperties> readAllDatabases() {
      * @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error.
      */
     public CosmosPagedFlux<CosmosDatabaseProperties> queryDatabases(String query, CosmosQueryRequestOptions options) {
-        return queryDatabases(new SqlQuerySpec(query), options);
+        return queryDatabasesInternal(new SqlQuerySpec(query), options);
     }
 
     /**
@@ -447,13 +430,7 @@ public CosmosPagedFlux<CosmosDatabaseProperties> queryDatabases(String query, Co
      * @return a {@link CosmosPagedFlux} containing one or several feed response pages of read databases or an error.
      */
     public CosmosPagedFlux<CosmosDatabaseProperties> queryDatabases(SqlQuerySpec querySpec, CosmosQueryRequestOptions options) {
-        return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> {
-            setContinuationTokenAndMaxItemCount(pagedFluxOptions, options);
-            return getDocClientWrapper().queryDatabases(querySpec, options)
-                                        .map(response -> BridgeInternal.createFeedResponse(
-                                            ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()),
-                                            response.getResponseHeaders()));
-        });
+        return queryDatabasesInternal(querySpec, options);
     }
 
     /**
@@ -473,4 +450,77 @@ public CosmosAsyncDatabase getDatabase(String id) {
     public void close() {
         asyncDocumentClient.close();
     }
+
+    TracerProvider getTracerProvider(){
+        return this.tracerProvider;
+    }
+
+    private CosmosPagedFlux<CosmosDatabaseProperties> queryDatabasesInternal(SqlQuerySpec querySpec, CosmosQueryRequestOptions options){
+        return UtilBridgeInternal.createCosmosPagedFlux(pagedFluxOptions -> {
+            String spanName = "queryDatabases";
+            pagedFluxOptions.setTracerInformation(this.tracerProvider, spanName, this.serviceEndpoint, null);
+            setContinuationTokenAndMaxItemCount(pagedFluxOptions, options);
+            return getDocClientWrapper().queryDatabases(querySpec, options)
+                .map(response -> BridgeInternal.createFeedResponse(
+                    ModelBridgeInternal.getCosmosDatabasePropertiesFromV2Results(response.getResults()),
+                    response.getResponseHeaders()));
+        });
+    }
+
+
+    private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(CosmosAsyncDatabase database,
+                                                                           ThroughputProperties throughputProperties, Context context) {
+        String spanName = "createDatabaseIfNotExists." + database.getId();
+        Context nestedContext = context.addData(TracerProvider.COSMOS_CALL_DEPTH, TracerProvider.COSMOS_CALL_DEPTH_VAL);
+        Mono<CosmosDatabaseResponse> responseMono = createDatabaseIfNotExistsInternal(database.readInternal(new CosmosDatabaseRequestOptions(), nestedContext), database, throughputProperties, nestedContext);
+        return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono,
+            context,
+            spanName,
+            database.getId(),
+            this.serviceEndpoint);
+    }
+
+    private Mono<CosmosDatabaseResponse> createDatabaseIfNotExistsInternal(Mono<CosmosDatabaseResponse> responseMono, CosmosAsyncDatabase database, ThroughputProperties throughputProperties, Context context) {
+        return responseMono.onErrorResume(exception -> {
+            final Throwable unwrappedException = Exceptions.unwrap(exception);
+            if (unwrappedException instanceof CosmosException) {
+                final CosmosException cosmosException = (CosmosException) unwrappedException;
+                if (cosmosException.getStatusCode() == HttpConstants.StatusCodes.NOTFOUND) {
+                    CosmosDatabaseRequestOptions requestOptions = new CosmosDatabaseRequestOptions();
+                    if(throughputProperties != null) {
+                        ModelBridgeInternal.setThroughputProperties(requestOptions, throughputProperties);
+                    }
+
+                    if (context != null) {
+                        Database wrappedDatabase = new Database();
+                        wrappedDatabase.setId(database.getId());
+                        return createDatabaseInternal(wrappedDatabase,
+                            requestOptions, context);
+                    }
+
+                    return createDatabase(new CosmosDatabaseProperties(database.getId()),
+                        requestOptions);
+                }
+            }
+            return Mono.error(unwrappedException);
+        });
+    }
+
+
+    private Mono<CosmosDatabaseResponse> createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options,
+                                                             Context context) {
+        String spanName = "createDatabase." + database.getId();
+        Mono<CosmosDatabaseResponse> responseMono = createDatabaseInternal(database, options);
+        return tracerProvider.traceEnabledCosmosResponsePublisher(responseMono,
+            context,
+            spanName,
+            database.getId(),
+            this.serviceEndpoint);
+    }
+
+    private Mono<CosmosDatabaseResponse> createDatabaseInternal(Database database, CosmosDatabaseRequestOptions options) {
+        return asyncDocumentClient.createDatabase(database, ModelBridgeInternal.toRequestOptions(options))
+            .map(databaseResourceResponse -> ModelBridgeInternal.createCosmosDatabaseResponse(databaseResourceResponse))
+            .single();
+    }
 }

sdk/cosmos/azure-cosmos/src/main/java/com/azure/cosmos/CosmosAsyncConflict.java

@@ -2,13 +2,16 @@
 // Licensed under the MIT License.
 package com.azure.cosmos;
 
+import com.azure.core.util.Context;
 import com.azure.cosmos.implementation.Paths;
 import com.azure.cosmos.implementation.RequestOptions;
 import com.azure.cosmos.models.CosmosConflictResponse;
 import com.azure.cosmos.models.CosmosConflictRequestOptions;
 import com.azure.cosmos.models.ModelBridgeInternal;
 import reactor.core.publisher.Mono;
 
+import static com.azure.core.util.FluxUtil.withContext;
+
 /**
  * Read and delete conflicts
  */
@@ -64,9 +67,7 @@ public Mono<CosmosConflictResponse> read(CosmosConflictRequestOptions options) {
             options = new CosmosConflictRequestOptions();
         }
         RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options);
-        return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), requestOptions)
-                   .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();
-
+        return withContext(context -> readInternal(requestOptions, context));
     }
 
     /**
@@ -85,8 +86,7 @@ public Mono<CosmosConflictResponse> delete(CosmosConflictRequestOptions options)
             options = new CosmosConflictRequestOptions();
         }
         RequestOptions requestOptions = ModelBridgeInternal.toRequestOptions(options);
-        return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), requestOptions)
-                   .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();
+        return withContext(context -> deleteInternal(requestOptions, context));
     }
 
     String getURIPathSegment() {
@@ -106,4 +106,33 @@ String getLink() {
         builder.append(getId());
         return builder.toString();
     }
+
+    private Mono<CosmosConflictResponse> readInternal(RequestOptions options, Context context) {
+        String spanName = "readConflict." + getId();
+        Mono<CosmosConflictResponse> responseMono = this.readInternal(options);
+        return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context,
+            spanName,
+            this.container.getDatabase().getId(),
+            this.container.getDatabase().getClient().getServiceEndpoint());
+
+    }
+
+    private Mono<CosmosConflictResponse> readInternal(RequestOptions options) {
+        return this.container.getDatabase().getDocClientWrapper().readConflict(getLink(), options)
+            .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();
+    }
+
+    private Mono<CosmosConflictResponse> deleteInternal(RequestOptions options, Context context) {
+        String spanName = "deleteConflict." + getId();
+        Mono<CosmosConflictResponse> responseMono = deleteInternal(options);
+        return this.container.getDatabase().getClient().getTracerProvider().traceEnabledCosmosResponsePublisher(responseMono, context,
+            spanName,
+            this.container.getDatabase().getId(),
+            this.container.getDatabase().getClient().getServiceEndpoint());
+    }
+
+    private Mono<CosmosConflictResponse> deleteInternal(RequestOptions options) {
+        return this.container.getDatabase().getDocClientWrapper().deleteConflict(getLink(), options)
+            .map(response -> ModelBridgeInternal.createCosmosConflictResponse(response)).single();
+    }
 }