feat: Enabling endToEndTracing support in Connection API (#3412)

ankiaga · olavloite · web-flow · commit 16cc6eed58cf · 2024-10-22T09:28:12.000+05:30
* feat: Enabling endToEndTracing support in Connection API

* fix formatting

* Comments incorporated

* formatting fixed

* formatting fixed

* Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java

Co-authored-by: Knut Olav Løite &lt;koloite@gmail.com&gt;

* Update google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java

Co-authored-by: Knut Olav Løite &lt;koloite@gmail.com&gt;

---------

Co-authored-by: Knut Olav Løite &lt;koloite@gmail.com&gt;
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionOptions.java
@@ -26,6 +26,7 @@
 import static com.google.cloud.spanner.connection.ConnectionProperties.DATA_BOOST_ENABLED;
 import static com.google.cloud.spanner.connection.ConnectionProperties.DIALECT;
 import static com.google.cloud.spanner.connection.ConnectionProperties.ENABLE_API_TRACING;
+import static com.google.cloud.spanner.connection.ConnectionProperties.ENABLE_END_TO_END_TRACING;
 import static com.google.cloud.spanner.connection.ConnectionProperties.ENABLE_EXTENDED_TRACING;
 import static com.google.cloud.spanner.connection.ConnectionProperties.ENCODED_CREDENTIALS;
 import static com.google.cloud.spanner.connection.ConnectionProperties.ENDPOINT;
@@ -249,6 +250,7 @@ public String[] getValidValues() {
   static final int DEFAULT_MAX_PARTITIONED_PARALLELISM = 1;
   static final Boolean DEFAULT_ENABLE_EXTENDED_TRACING = null;
   static final Boolean DEFAULT_ENABLE_API_TRACING = null;
+  static final boolean DEFAULT_ENABLE_END_TO_END_TRACING = false;
   static final boolean DEFAULT_AUTO_BATCH_DML = false;
   static final long DEFAULT_AUTO_BATCH_DML_UPDATE_COUNT = 1L;
   static final boolean DEFAULT_AUTO_BATCH_DML_UPDATE_COUNT_VERIFICATION = true;
@@ -335,6 +337,7 @@ public String[] getValidValues() {
 
   public static final String ENABLE_EXTENDED_TRACING_PROPERTY_NAME = "enableExtendedTracing";
   public static final String ENABLE_API_TRACING_PROPERTY_NAME = "enableApiTracing";
+  public static final String ENABLE_END_TO_END_TRACING_PROPERTY_NAME = "enableEndToEndTracing";
 
   public static final String AUTO_BATCH_DML_PROPERTY_NAME = "auto_batch_dml";
   public static final String AUTO_BATCH_DML_UPDATE_COUNT_PROPERTY_NAME =
@@ -537,7 +540,14 @@ static boolean isEnableTransactionalConnectionStateForPostgreSQL() {
                           + "to get a detailed view of each RPC that is being executed by your application, "
                           + "or if you want to debug potential latency problems caused by RPCs that are "
                           + "being retried.",
-                      DEFAULT_ENABLE_API_TRACING))));
+                      DEFAULT_ENABLE_API_TRACING),
+                  ConnectionProperty.createBooleanProperty(
+                      ENABLE_END_TO_END_TRACING_PROPERTY_NAME,
+                      "Enable end-to-end tracing (true/false) to generate traces for both the time "
+                          + "that is spent in the client, as well as time that is spent in the Spanner server. "
+                          + "Server side traces can only go to Google Cloud Trace, so to see end to end traces, "
+                          + "the application should configure an exporter that exports the traces to Google Cloud Trace.",
+                      DEFAULT_ENABLE_END_TO_END_TRACING))));
 
   private static final Set<ConnectionProperty> INTERNAL_PROPERTIES =
       Collections.unmodifiableSet(
@@ -1205,6 +1215,11 @@ public boolean isRouteToLeader() {
     return getInitialConnectionPropertyValue(ROUTE_TO_LEADER);
   }
 
+  /** Whether end-to-end tracing is enabled. */
+  public boolean isEndToEndTracingEnabled() {
+    return getInitialConnectionPropertyValue(ENABLE_END_TO_END_TRACING);
+  }
+
   /**
    * The initial retryAbortsInternally value for connections created by this {@link
    * ConnectionOptions}
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/ConnectionProperties.java
@@ -39,6 +39,7 @@
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DDL_IN_TRANSACTION_MODE;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_API_TRACING;
+import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_END_TO_END_TRACING;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENABLE_EXTENDED_TRACING;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_ENDPOINT;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DEFAULT_KEEP_TRANSACTION_ALIVE;
@@ -65,6 +66,7 @@
 import static com.google.cloud.spanner.connection.ConnectionOptions.DELAY_TRANSACTION_START_UNTIL_FIRST_WRITE_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.DIALECT_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_API_TRACING_PROPERTY_NAME;
+import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_END_TO_END_TRACING_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_EXTENDED_TRACING_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENCODED_CREDENTIALS_PROPERTY_NAME;
 import static com.google.cloud.spanner.connection.ConnectionOptions.ENDPOINT_PROPERTY_NAME;
@@ -292,6 +294,16 @@ class ConnectionProperties {
           DEFAULT_ENABLE_API_TRACING,
           BooleanConverter.INSTANCE,
           Context.STARTUP);
+  static final ConnectionProperty<Boolean> ENABLE_END_TO_END_TRACING =
+      create(
+          ENABLE_END_TO_END_TRACING_PROPERTY_NAME,
+          "Enable end-to-end tracing (true/false) to generate traces for both the time "
+              + "that is spent in the client, as well as time that is spent in the Spanner server. "
+              + "Server side traces can only go to Google Cloud Trace, so to see end to end traces, "
+              + "the application should configure an exporter that exports the traces to Google Cloud Trace.",
+          DEFAULT_ENABLE_END_TO_END_TRACING,
+          BooleanConverter.INSTANCE,
+          Context.STARTUP);
   static final ConnectionProperty<Integer> MIN_SESSIONS =
       create(
           MIN_SESSIONS_PROPERTY_NAME,
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/connection/SpannerPool.java
@@ -160,6 +160,7 @@ static class SpannerPoolKey {
     private final OpenTelemetry openTelemetry;
     private final Boolean enableExtendedTracing;
     private final Boolean enableApiTracing;
+    private final boolean enableEndToEndTracing;
 
     @VisibleForTesting
     static SpannerPoolKey of(ConnectionOptions options) {
@@ -190,6 +191,7 @@ private SpannerPoolKey(ConnectionOptions options) throws IOException {
       this.openTelemetry = options.getOpenTelemetry();
       this.enableExtendedTracing = options.isEnableExtendedTracing();
       this.enableApiTracing = options.isEnableApiTracing();
+      this.enableEndToEndTracing = options.isEndToEndTracingEnabled();
     }
 
     @Override
@@ -211,7 +213,8 @@ public boolean equals(Object o) {
               this.useVirtualGrpcTransportThreads, other.useVirtualGrpcTransportThreads)
           && Objects.equals(this.openTelemetry, other.openTelemetry)
           && Objects.equals(this.enableExtendedTracing, other.enableExtendedTracing)
-          && Objects.equals(this.enableApiTracing, other.enableApiTracing);
+          && Objects.equals(this.enableApiTracing, other.enableApiTracing)
+          && Objects.equals(this.enableEndToEndTracing, other.enableEndToEndTracing);
     }
 
     @Override
@@ -229,7 +232,8 @@ public int hashCode() {
           this.useVirtualGrpcTransportThreads,
           this.openTelemetry,
           this.enableExtendedTracing,
-          this.enableApiTracing);
+          this.enableApiTracing,
+          this.enableEndToEndTracing);
     }
   }
 
@@ -380,6 +384,9 @@ Spanner createSpanner(SpannerPoolKey key, ConnectionOptions options) {
     if (!options.isRouteToLeader()) {
       builder.disableLeaderAwareRouting();
     }
+    if (options.isEndToEndTracingEnabled()) {
+      builder.setEnableEndToEndTracing(true);
+    }
     if (key.usePlainText) {
       // Credentials may not be sent over a plain text channel.
       builder.setCredentials(NoCredentials.getInstance());
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/ConnectionOptionsTest.java
@@ -322,6 +322,27 @@ public void testBuildWithRouteToLeader() {
     assertTrue(options.isRouteToLeader());
   }
 
+  @Test
+  public void testBuildWithEndToEndTracingEnabled() {
+    final String BASE_URI =
+        "cloudspanner:/projects/test-project-123/instances/test-instance-123/databases/test-database-123";
+    ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();
+    builder.setUri(BASE_URI + "?enableEndToEndTracing=true");
+    builder.setCredentialsUrl(FILE_TEST_PATH);
+    ConnectionOptions options = builder.build();
+    assertEquals(options.getHost(), DEFAULT_HOST);
+    assertEquals(options.getProjectId(), TEST_PROJECT);
+    assertEquals(options.getInstanceId(), TEST_INSTANCE);
+    assertEquals(options.getDatabaseName(), TEST_DATABASE);
+    assertTrue(options.isEndToEndTracingEnabled());
+
+    // Test for default behavior for enableEndToEndTracing property.
+    builder = ConnectionOptions.newBuilder().setUri(BASE_URI);
+    builder.setCredentialsUrl(FILE_TEST_PATH);
+    options = builder.build();
+    assertFalse(options.isEndToEndTracingEnabled());
+  }
+
   @Test
   public void testBuildWithAutoConfigEmulatorAndHost() {
     ConnectionOptions.Builder builder = ConnectionOptions.newBuilder();
diff --git a/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java b/google-cloud-spanner/src/test/java/com/google/cloud/spanner/connection/SpannerPoolTest.java
@@ -563,6 +563,58 @@ public void testEnableApiTracing() {
                 .build()));
   }
 
+  @Test
+  public void testEnableEndToEndTracing() {
+    SpannerPoolKey keyWithoutApiTracingConfig =
+        SpannerPoolKey.of(
+            ConnectionOptions.newBuilder()
+                .setUri("cloudspanner:/projects/p/instances/i/databases/d")
+                .setCredentials(NoCredentials.getInstance())
+                .build());
+    SpannerPoolKey keyWithApiTracingEnabled =
+        SpannerPoolKey.of(
+            ConnectionOptions.newBuilder()
+                .setUri(
+                    "cloudspanner:/projects/p/instances/i/databases/d?enableEndToEndTracing=true")
+                .setCredentials(NoCredentials.getInstance())
+                .build());
+    SpannerPoolKey keyWithApiTracingDisabled =
+        SpannerPoolKey.of(
+            ConnectionOptions.newBuilder()
+                .setUri(
+                    "cloudspanner:/projects/p/instances/i/databases/d?enableEndToEndTracing=false")
+                .setCredentials(NoCredentials.getInstance())
+                .build());
+
+    assertNotEquals(keyWithoutApiTracingConfig, keyWithApiTracingEnabled);
+    assertEquals(keyWithoutApiTracingConfig, keyWithApiTracingDisabled);
+    assertNotEquals(keyWithApiTracingEnabled, keyWithApiTracingDisabled);
+
+    assertEquals(
+        keyWithApiTracingEnabled,
+        SpannerPoolKey.of(
+            ConnectionOptions.newBuilder()
+                .setUri(
+                    "cloudspanner:/projects/p/instances/i/databases/d?enableEndToEndTracing=true")
+                .setCredentials(NoCredentials.getInstance())
+                .build()));
+    assertEquals(
+        keyWithApiTracingDisabled,
+        SpannerPoolKey.of(
+            ConnectionOptions.newBuilder()
+                .setUri(
+                    "cloudspanner:/projects/p/instances/i/databases/d?enableEndToEndTracing=false")
+                .setCredentials(NoCredentials.getInstance())
+                .build()));
+    assertEquals(
+        keyWithoutApiTracingConfig,
+        SpannerPoolKey.of(
+            ConnectionOptions.newBuilder()
+                .setUri("cloudspanner:/projects/p/instances/i/databases/d")
+                .setCredentials(NoCredentials.getInstance())
+                .build()));
+  }
+
   @Test
   public void testOpenTelemetry() {
     SpannerPool pool = createSubjectAndMocks();
