Repository: GoogleCloudPlatform/pgadapter
Pull request #2504 (2024-11-01 13:05:01)

feat: add option for end-to-end tracing
Add a command line argument and option for enabling end-to-end tracing when running PGAdapter. This will propagate the current trace context to the Spanner server, and collect server-side traces together with the traces that are collected by PGAdapter.

The server-side traces can only be exported to Google Cloud Trace. Make sure that you have set up an OpenTelemetry instance that exports traces to Google Cloud Trace to use this feature.


Patch:
diff --git a/benchmarks/latency-comparison/golang/go.mod b/benchmarks/latency-comparison/golang/go.mod
index 8fb4974c..269782c1 100644
--- a/benchmarks/latency-comparison/golang/go.mod
+++ b/benchmarks/latency-comparison/golang/go.mod
@@ -5,7 +5,7 @@ go 1.22
 toolchain go1.23.2
 
 require (
-	cloud.google.com/go/spanner v1.70.0
+	cloud.google.com/go/spanner v1.71.0
 	github.com/docker/docker v27.3.1+incompatible
 	github.com/jackc/pgx/v4 v4.18.3
 	github.com/jackc/pgx/v5 v5.7.1
diff --git a/benchmarks/latency-comparison/golang/go.sum b/benchmarks/latency-comparison/golang/go.sum
index b7e8e7f5..929be750 100644
--- a/benchmarks/latency-comparison/golang/go.sum
+++ b/benchmarks/latency-comparison/golang/go.sum
@@ -580,6 +580,8 @@ cloud.google.com/go/spanner v1.69.0 h1:/llYvGj7j8ocRDagvHfJW1cM7url9Oa94IScX37eI
 cloud.google.com/go/spanner v1.69.0/go.mod h1:X5T0XftydYp0K1adeJQDJtdWpbrOeJ7wHecM4tK6FiE=
 cloud.google.com/go/spanner v1.70.0 h1:nj6p/GJTgMDiSQ1gQ034ItsKuJgHiMOjtOlONOg8PSo=
 cloud.google.com/go/spanner v1.70.0/go.mod h1:X5T0XftydYp0K1adeJQDJtdWpbrOeJ7wHecM4tK6FiE=
+cloud.google.com/go/spanner v1.71.0 h1:Gl4eX16JEvqlm3yvZ2lKENaQ6FvCkM3YZfO/Hq8OqZU=
+cloud.google.com/go/spanner v1.71.0/go.mod h1:mw98ua5ggQXVWwp83yjwggqEmW9t8rjs9Po1ohcUGW4=
 cloud.google.com/go/speech v1.6.0/go.mod h1:79tcr4FHCimOp56lwC01xnt/WPJZc4v3gzyT7FoBkCM=
 cloud.google.com/go/speech v1.7.0/go.mod h1:KptqL+BAQIhMsj1kOP2la5DSEEerPDuOP/2mmkhHhZQ=
 cloud.google.com/go/speech v1.8.0/go.mod h1:9bYIl1/tjsAnMgKGHKmBZzXKEkGgtU+MpdDPTE9f7y0=
diff --git a/benchmarks/latency-comparison/java/pom.xml b/benchmarks/latency-comparison/java/pom.xml
index 9915e638..eb5fcd6f 100644
--- a/benchmarks/latency-comparison/java/pom.xml
+++ b/benchmarks/latency-comparison/java/pom.xml
@@ -109,7 +109,7 @@
     <dependency>
       <groupId>org.checkerframework</groupId>
       <artifactId>checker-qual</artifactId>
-      <version>3.48.1</version>
+      <version>3.48.2</version>
     </dependency>
 
     <dependency>
diff --git a/pom.xml b/pom.xml
index 2c35bd55..a295209a 100644
--- a/pom.xml
+++ b/pom.xml
@@ -110,7 +110,7 @@
     <dependency>
       <groupId>org.checkerframework</groupId>
       <artifactId>checker-qual</artifactId>
-      <version>3.48.1</version>
+      <version>3.48.2</version>
     </dependency>
     <dependency>
       <groupId>io.opentelemetry</groupId>
diff --git a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
index 1249305d..34589953 100644
--- a/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
+++ b/src/main/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadata.java
@@ -14,6 +14,8 @@
 
 package com.google.cloud.spanner.pgadapter.metadata;
 
+import static com.google.cloud.spanner.connection.ConnectionOptions.ENABLE_END_TO_END_TRACING_PROPERTY_NAME;
+
 import com.google.auth.Credentials;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.NoCredentials;
@@ -92,6 +94,7 @@ public class OptionsMetadata {
     private boolean enableOpenTelemetry;
     private boolean enableOpenTelemetryMetrics;
     private Double openTelemetryTraceRatio;
+    private boolean enableEndToEndTracing;
     private boolean skipLocalhostCheck;
     private boolean useVirtualThreads;
     private boolean useVirtualGrpcTransportThreads;
@@ -266,6 +269,17 @@ public class OptionsMetadata {
       return this;
     }
 
+    /**
+     * Enables end-to-end tracing for RPCs on Spanner. This generates traces for both the time +
+     * that is spent in the client and 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.
+     */
+    public Builder setEnableEndToEndTracing(boolean enableEndToEndTracing) {
+      this.enableEndToEndTracing = enableEndToEndTracing;
+      return this;
+    }
+
     /**
      * PGAdapter by default only allows connections from localhost. Call this method to disable this
      * check. You should only allow connections from private networks, unless you are also using SSL
@@ -447,6 +461,9 @@ public class OptionsMetadata {
         addLongOption(
             args, OPTION_OPEN_TELEMETRY_TRACE_RATIO, String.valueOf(openTelemetryTraceRatio));
       }
+      if (enableEndToEndTracing) {
+        addOption(args, OPTION_ENABLE_END_TO_END_TRACING);
+      }
       if (skipLocalhostCheck) {
         addOption(args, OPTION_DISABLE_LOCALHOST_CHECK);
       }
@@ -467,7 +484,8 @@ public class OptionsMetadata {
           || databaseRole != null
           || autoConfigEmulator
           || useVirtualThreads
-          || useVirtualGrpcTransportThreads) {
+          || useVirtualGrpcTransportThreads
+          || enableEndToEndTracing) {
         StringBuilder jdbcOptionBuilder = new StringBuilder();
         if (usePlainText) {
           jdbcOptionBuilder.append("usePlainText=true;");
@@ -491,6 +509,9 @@ public class OptionsMetadata {
               .append(ConnectionOptions.USE_VIRTUAL_GRPC_TRANSPORT_THREADS_PROPERTY_NAME)
               .append("=true;");
         }
+        if (enableEndToEndTracing) {
+          jdbcOptionBuilder.append(ENABLE_END_TO_END_TRACING_PROPERTY_NAME).append("=true;");
+        }
         addOption(args, OPTION_JDBC_PROPERTIES, jdbcOptionBuilder.toString());
       }
       if (logGrpcMessages) {
@@ -588,6 +609,7 @@ public class OptionsMetadata {
   private static final String OPTION_ENABLE_OPEN_TELEMETRY = "enable_otel";
   private static final String OPTION_ENABLE_OPEN_TELEMETRY_METRICS = "enable_otel_metrics";
   private static final String OPTION_OPEN_TELEMETRY_TRACE_RATIO = "otel_trace_ratio";
+  private static final String OPTION_ENABLE_END_TO_END_TRACING = "enable_end_to_end_tracing";
   private static final String OPTION_SSL = "ssl";
   private static final String OPTION_DISABLE_AUTO_DETECT_CLIENT = "disable_auto_detect_client";
   private static final String OPTION_DISABLE_DEFAULT_LOCAL_STATEMENTS =
@@ -635,6 +657,7 @@ public class OptionsMetadata {
   private final boolean enableOpenTelemetry;
   private final boolean enableOpenTelemetryMetrics;
   private final Double openTelemetryTraceRatio;
+  private final boolean enableEndToEndTracing;
   private final SslMode sslMode;
   private final boolean disableAutoDetectClient;
   private final boolean disableDefaultLocalStatements;
@@ -725,6 +748,7 @@ public class OptionsMetadata {
     this.enableOpenTelemetryMetrics = commandLine.hasOption(OPTION_ENABLE_OPEN_TELEMETRY_METRICS);
     this.openTelemetryTraceRatio =
         parseOpenTelemetryTraceRatio(commandLine.getOptionValue(OPTION_OPEN_TELEMETRY_TRACE_RATIO));
+    this.enableEndToEndTracing = commandLine.hasOption(OPTION_ENABLE_END_TO_END_TRACING);
     this.sslMode = parseSslMode(commandLine.getOptionValue(OPTION_SSL));
     this.disableAutoDetectClient = commandLine.hasOption(OPTION_DISABLE_AUTO_DETECT_CLIENT);
     this.disableDefaultLocalStatements =
@@ -805,6 +829,7 @@ public class OptionsMetadata {
     this.enableOpenTelemetry = false;
     this.enableOpenTelemetryMetrics = false;
     this.openTelemetryTraceRatio = null;
+    this.enableEndToEndTracing = false;
     this.sslMode = SslMode.Disable;
     this.disableAutoDetectClient = false;
     this.disableDefaultLocalStatements = false;
@@ -1180,6 +1205,14 @@ public class OptionsMetadata {
         OPTION_OPEN_TELEMETRY_TRACE_RATIO,
         true,
         "OpenTelemetry trace sampling ration. Value must be in the range [0.0, 1.0].");
+    options.addOption(
+        null,
+        OPTION_ENABLE_END_TO_END_TRACING,
+        false,
+        "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.");
     options.addOption(
         OPTION_SSL,
         "sslmode",
@@ -1549,6 +1582,10 @@ public class OptionsMetadata {
     return this.openTelemetryTraceRatio;
   }
 
+  public boolean isEnableEndToEndTracing() {
+    return this.enableEndToEndTracing;
+  }
+
   public SslMode getSslMode() {
     return this.sslMode;
   }


Test patch:
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
index fae84944..fd706cc3 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/AbortedMockServerTest.java
@@ -879,6 +879,7 @@ public class AbortedMockServerTest extends AbstractMockServerTest {
               .setCredentials(NoCredentials.getInstance())
               .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
               .setClientLibToken("pg-adapter")
+              .setEnableEndToEndTracing(true)
               .build()
               .getService();
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
index 3220a52a..7c806a45 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/AbstractMockServerTest.java
@@ -16,6 +16,7 @@ package com.google.cloud.spanner.pgadapter;
 
 import static com.google.cloud.spanner.pgadapter.statements.PgCatalog.PG_TYPE_CTE_EMULATED;
 import static com.google.cloud.spanner.pgadapter.statements.PgCatalog.PgNamespace.PG_NAMESPACE_CTE;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertTrue;
 
@@ -1280,6 +1281,13 @@ public abstract class AbstractMockServerTest {
                       assertTrue(userAgent.contains("pg-adapter"));
                       assertTrue(
                           userAgent.contains(ServiceOptions.getGoogApiClientLibName() + "/"));
+
+                      String endToEndTracing =
+                          metadata.get(
+                              Metadata.Key.of(
+                                  "x-goog-spanner-end-to-end-tracing",
+                                  Metadata.ASCII_STRING_MARSHALLER));
+                      assertEquals("true", endToEndTracing);
                     }
                     return Contexts.interceptCall(
                         Context.current(), serverCall, metadata, serverCallHandler);
@@ -1297,7 +1305,8 @@ public abstract class AbstractMockServerTest {
         .enableDebugMode()
         .setUsePlainText()
         .setEndpoint(String.format("localhost:%d", spannerServer.getPort()))
-        .setCredentials(NoCredentials.getInstance());
+        .setCredentials(NoCredentials.getInstance())
+        .setEnableEndToEndTracing(true);
     optionsConfigurator.accept(builder);
     pgServer = new ProxyServer(builder.build(), openTelemetry);
     pgServer.startServer();
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
index f4f8eb86..6339c99e 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/JdbcMockServerTest.java
@@ -3249,6 +3249,7 @@ public class JdbcMockServerTest extends AbstractMockServerTest {
                 .setCredentials(NoCredentials.getInstance())
                 .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
                 .setClientLibToken("pg-adapter")
+                .setEnableEndToEndTracing(true)
                 .build()
                 .getService();
         DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of("p", "i", "d"));
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
index ff3e5ca8..66cf77c7 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/ShutdownModeMockServerTest.java
@@ -63,7 +63,8 @@ public class ShutdownModeMockServerTest extends AbstractMockServerTest {
         .setUsePlainText()
         .setEndpoint(String.format("localhost:%d", spannerServer.getPort()))
         .setCredentials(NoCredentials.getInstance())
-        .setAllowShutdownStatement(true);
+        .setAllowShutdownStatement(true)
+        .setEnableEndToEndTracing(true);
     proxyServer = new ProxyServer(builder.build(), OpenTelemetry.noop());
     proxyServer.startServer();
   }
diff --git a/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java b/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java
index 8d76a12c..47cba071 100644
--- a/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java
+++ b/src/test/java/com/google/cloud/spanner/pgadapter/metadata/OptionsMetadataTest.java
@@ -457,6 +457,31 @@ public class OptionsMetadataTest {
             .useVirtualGrpcTransportThreads()
             .build()
             .isUseGrpcTransportVirtualThreads());
+    assertNull(
+        OptionsMetadata.newBuilder()
+            .setCredentials(NoCredentials.getInstance())
+            .build()
+            .getPropertyMap()
+            .get("enableEndToEndTracing"));
+    assertEquals(
+        "true",
+        OptionsMetadata.newBuilder()
+            .setEnableEndToEndTracing(true)
+            .setCredentials(NoCredentials.getInstance())
+            .build()
+            .getPropertyMap()
+            .get("enableEndToEndTracing"));
+    assertFalse(OptionsMetadata.newBuilder().build().isEnableEndToEndTracing());
+    assertTrue(
+        OptionsMetadata.newBuilder()
+            .setEnableEndToEndTracing(true)
+            .build()
+            .isEnableEndToEndTracing());
+    assertFalse(
+        OptionsMetadata.newBuilder()
+            .setEnableEndToEndTracing(false)
+            .build()
+            .isEnableEndToEndTracing());
 
     assertEquals(
         DdlTransactionMode.Batch,