Cosmos: Diagnostic improvements

eng/code-quality-reports/src/main/resources/spotbugs/spotbugs-exclude.xml

@@ -2061,6 +2061,7 @@
   <Match>
     <Or>
       <Class name="~com.azure.cosmos.implementation.ImplementationBridgeHelpers(\$.*)?"/>
+      <Class name="~com.azure.cosmos.implementation.DiagnosticsProvider(\$.*)?"/>
       <Class name="~com.azure.cosmos.test.implementation.ImplementationBridgeHelpers(\$.*)?"/>
     </Or>
     <Bug pattern="DM_EXIT"/>

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

@@ -10,6 +10,7 @@
 import com.azure.cosmos.implementation.CosmosError;
 import com.azure.cosmos.implementation.DatabaseAccount;
 import com.azure.cosmos.implementation.DiagnosticsClientContext;
+import com.azure.cosmos.implementation.DiagnosticsProvider;
 import com.azure.cosmos.implementation.Document;
 import com.azure.cosmos.implementation.FeedResponseDiagnostics;
 import com.azure.cosmos.implementation.GlobalEndpointManager;
@@ -27,7 +28,6 @@
 import com.azure.cosmos.implementation.SerializationDiagnosticsContext;
 import com.azure.cosmos.implementation.ServiceUnavailableException;
 import com.azure.cosmos.implementation.StoredProcedureResponse;
-import com.azure.cosmos.implementation.TracerProvider;
 import com.azure.cosmos.implementation.Warning;
 import com.azure.cosmos.implementation.directconnectivity.StoreResponse;
 import com.azure.cosmos.implementation.directconnectivity.StoreResponseDiagnostics;
@@ -82,6 +82,9 @@ public static CosmosDiagnostics createCosmosDiagnostics(DiagnosticsClientContext
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static Set<String> getRegionsContacted(CosmosDiagnostics cosmosDiagnostics) {
+        if (cosmosDiagnostics.clientSideRequestStatistics() == null) {
+            return Collections.<String>emptySet();
+        }
         return cosmosDiagnostics.clientSideRequestStatistics().getContactedRegionNames();
     }
 
@@ -188,8 +191,23 @@ public static <T> boolean noChanges(FeedResponse<T> page) {
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
     public static <T> FeedResponse<T> createFeedResponse(List<T> results,
-            Map<String, String> headers) {
-        return ModelBridgeInternal.createFeedResponse(results, headers);
+            Map<String, String> headers, CosmosDiagnostics cosmosDiagnostics) {
+        FeedResponse<T> feedResponseWithDiagnostics = ModelBridgeInternal.createFeedResponse(results, headers);
+
+        if (cosmosDiagnostics == null) {
+            return feedResponseWithDiagnostics;
+        }
+
+        ClientSideRequestStatistics requestStatistics = cosmosDiagnostics.clientSideRequestStatistics();
+        if (requestStatistics != null) {
+            BridgeInternal.addClientSideDiagnosticsToFeed(feedResponseWithDiagnostics.getCosmosDiagnostics(),
+                Collections.singletonList(requestStatistics));
+        }
+        BridgeInternal.addClientSideDiagnosticsToFeed(feedResponseWithDiagnostics.getCosmosDiagnostics(),
+            cosmosDiagnostics.getFeedResponseDiagnostics()
+                             .getClientSideRequestStatisticsList());
+
+        return feedResponseWithDiagnostics;
     }
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
@@ -721,8 +739,8 @@ public static CosmosDatabase createCosmosDatabase(String id, CosmosClient client
     }
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
-    public static TracerProvider getTracerProvider(CosmosAsyncClient client) {
-        return client.getTracerProvider();
+    public static DiagnosticsProvider getTracerProvider(CosmosAsyncClient client) {
+        return client.getDiagnosticsProvider();
     }
 
     @Warning(value = INTERNAL_USE_ONLY_WARNING)
@@ -772,8 +790,10 @@ public static void  initializeAllAccessors() {
         CosmosAsyncDatabase.initialize();
         CosmosClientBuilder.initialize();
         CosmosDiagnostics.initialize();
+        CosmosDiagnosticsContext.initialize();
         CosmosException.initialize();
         DirectConnectionConfig.initialize();
         CosmosAsyncClient.initialize();
+        CosmosDiagnosticsThresholds.initialize();
     }
 }