feat: Record transaction metrics for Datastore

Author: lqiu96

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreException.java

@@ -160,4 +160,33 @@ static DatastoreException throwInvalidRequest(String massage, Object... params)
   static DatastoreException propagateUserException(Exception ex) {
     throw new DatastoreException(BaseServiceException.UNKNOWN_CODE, ex.getMessage(), null, ex);
   }
+
+  /**
+   * Extracts the status code name from the given throwable. Walks the exception cause chain looking
+   * for a {@link DatastoreException} that carries a reason string representing the status code
+   * (e.g. "ABORTED", "UNAVAILABLE"). The reason is set from {@link
+   * com.google.api.gax.rpc.StatusCode.Code} which is transport-neutral, supporting both gRPC and
+   * HttpJson. Falls back to "UNKNOWN" if the status cannot be determined.
+   *
+   * <p>Note: Some {@link DatastoreException} instances are constructed without a reason (e.g. via
+   * {@link DatastoreException#DatastoreException(int, String, Throwable)}). If all {@link
+   * DatastoreException} instances in the cause chain have a null or empty reason, this method
+   * returns "UNKNOWN" even if the underlying error carries a meaningful status.
+   *
+   * @param throwable the throwable to extract the status code from
+   * @return the status code name, or "UNKNOWN" if not determinable
+   */
+  public static String extractStatusCode(Throwable throwable) {
+    Throwable current = throwable;
+    while (current != null) {
+      if (current instanceof DatastoreException) {
+        String reason = ((DatastoreException) current).getReason();
+        if (!Strings.isNullOrEmpty(reason)) {
+          return reason;
+        }
+      }
+      current = current.getCause();
+    }
+    return StatusCode.Code.UNKNOWN.toString();
+  }
 }

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreImpl.java

@@ -38,17 +38,21 @@
 
 import com.google.api.core.BetaApi;
 import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.rpc.StatusCode;
 import com.google.cloud.BaseService;
 import com.google.cloud.ExceptionHandler;
 import com.google.cloud.RetryHelper;
 import com.google.cloud.RetryHelper.RetryHelperException;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.datastore.execution.AggregationQueryExecutor;
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
+import com.google.cloud.datastore.telemetry.MetricsRecorder;
+import com.google.cloud.datastore.telemetry.TelemetryConstants;
 import com.google.cloud.datastore.telemetry.TraceUtil;
 import com.google.cloud.datastore.telemetry.TraceUtil.Scope;
 import com.google.common.base.MoreObjects;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Stopwatch;
 import com.google.common.collect.AbstractIterator;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -61,10 +65,12 @@
 import com.google.datastore.v1.RunQueryResponse;
 import com.google.datastore.v1.TransactionOptions;
 import com.google.protobuf.ByteString;
+import io.grpc.Status;
 import io.opentelemetry.context.Context;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
@@ -73,6 +79,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import javax.annotation.Nullable;
@@ -89,6 +96,7 @@ final class DatastoreImpl extends BaseService<DatastoreOptions> implements Datas
 
   private final com.google.cloud.datastore.telemetry.TraceUtil otelTraceUtil =
       getOptions().getTraceUtil();
+  private final MetricsRecorder metricsRecorder = getOptions().getMetricsRecorder();
 
   private final ReadOptionProtoPreparer readOptionProtoPreparer;
   private final AggregationQueryExecutor aggregationQueryExecutor;
@@ -122,63 +130,31 @@ public Transaction newTransaction() {
     return new TransactionImpl(this);
   }
 
+  /**
+   * A wrapper around {@link ReadWriteTransactionCallable} that adds OpenTelemetry tracing context
+   * propagation. All transaction logic (begin, run, commit, rollback, metrics recording) is
+   * delegated to the underlying {@link ReadWriteTransactionCallable}.
+   */
   static class TracedReadWriteTransactionCallable<T> implements Callable<T> {
-    private final Datastore datastore;
-    private final TransactionCallable<T> callable;
-    private volatile TransactionOptions options;
-    private volatile Transaction transaction;
-
+    private final ReadWriteTransactionCallable<T> delegate;
     private final TraceUtil.Span parentSpan;
 
     TracedReadWriteTransactionCallable(
-        Datastore datastore,
-        TransactionCallable<T> callable,
-        TransactionOptions options,
+        ReadWriteTransactionCallable<T> delegate,
         @Nullable com.google.cloud.datastore.telemetry.TraceUtil.Span parentSpan) {
-      this.datastore = datastore;
-      this.callable = callable;
-      this.options = options;
-      this.transaction = null;
+      this.delegate = delegate;
       this.parentSpan = parentSpan;
     }
 
-    Datastore getDatastore() {
-      return datastore;
-    }
-
-    TransactionOptions getOptions() {
-      return options;
-    }
-
-    Transaction getTransaction() {
-      return transaction;
-    }
-
-    void setPrevTransactionId(ByteString transactionId) {
-      TransactionOptions.ReadWrite readWrite =
-          TransactionOptions.ReadWrite.newBuilder().setPreviousTransaction(transactionId).build();
-      options = options.toBuilder().setReadWrite(readWrite).build();
+    ReadWriteTransactionCallable<T> getDelegate() {
+      return delegate;
     }
 
     @Override
     public T call() throws DatastoreException {
       try (io.opentelemetry.context.Scope ignored =
           Context.current().with(parentSpan.getSpan()).makeCurrent()) {
-        transaction = datastore.newTransaction(options);
-        T value = callable.run(transaction);
-        transaction.commit();
-        return value;
-      } catch (Exception ex) {
-        transaction.rollback();
-        throw DatastoreException.propagateUserException(ex);
-      } finally {
-        if (transaction.isActive()) {
-          transaction.rollback();
-        }
-        if (options != null
-            && options.getModeCase().equals(TransactionOptions.ModeCase.READ_WRITE)) {
-          setPrevTransactionId(transaction.getTransactionId());
-        }
+        return delegate.call();
       }
     }
   }
@@ -200,14 +176,19 @@ public boolean isClosed() {
   static class ReadWriteTransactionCallable<T> implements Callable<T> {
     private final Datastore datastore;
     private final TransactionCallable<T> callable;
+    private final MetricsRecorder metricsRecorder;
     private volatile TransactionOptions options;
     private volatile Transaction transaction;
 
     ReadWriteTransactionCallable(
-        Datastore datastore, TransactionCallable<T> callable, TransactionOptions options) {
+        Datastore datastore,
+        TransactionCallable<T> callable,
+        TransactionOptions options,
+        MetricsRecorder metricsRecorder) {
       this.datastore = datastore;
       this.callable = callable;
       this.options = options;
+      this.metricsRecorder = metricsRecorder;
       this.transaction = null;
     }
 
@@ -231,15 +212,28 @@ void setPrevTransactionId(ByteString transactionId) {
 
     @Override
     public T call() throws DatastoreException {
+      String attemptStatus = StatusCode.Code.UNKNOWN.toString();
       try {
         transaction = datastore.newTransaction(options);
         T value = callable.run(transaction);
         transaction.commit();
+        attemptStatus = Status.Code.OK.toString();
         return value;
       } catch (Exception ex) {
-        transaction.rollback();
+        attemptStatus = DatastoreException.extractStatusCode(ex);
+        // An exception here can originate from either callable.run() (before commit was attempted)
+        // or from transaction.commit() itself. In both cases the transaction is still active.
+        // isActive() returns false if the transaction was already committed or rolled back, so
+        // it is safe to use as the sole guard here without tracking a separate committed flag.
+        if (transaction.isActive()) {
+          transaction.rollback();
+        }
         throw DatastoreException.propagateUserException(ex);
       } finally {
+        recordAttempt(attemptStatus);
+        // transaction.isActive() returns false after both a successful commit or a completed
+        // rollback, so it already guards against rolling back a committed transaction or
+        // rolling back a transaction that has already been rolled back.
         if (transaction.isActive()) {
           transaction.rollback();
         }
@@ -249,50 +243,63 @@ public T call() throws DatastoreException {
         }
       }
     }
+
+    /**
+     * Records a single transaction commit attempt with the given status code. This is called once
+     * per invocation of {@link #call()}, capturing the outcome of each individual commit attempt.
+     */
+    private void recordAttempt(String status) {
+      Map<String, String> attributes = new HashMap<>();
+      attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, status);
+      attributes.put(
+          TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_COMMIT);
+      attributes.put(
+          TelemetryConstants.ATTRIBUTES_KEY_PROJECT_ID, datastore.getOptions().getProjectId());
+      attributes.put(
+          TelemetryConstants.ATTRIBUTES_KEY_DATABASE_ID, datastore.getOptions().getDatabaseId());
+      metricsRecorder.recordTransactionAttemptCount(1, attributes);
+    }
   }
 
   @Override
   public <T> T runInTransaction(final TransactionCallable<T> callable) {
-    TraceUtil.Span span = otelTraceUtil.startSpan(SPAN_NAME_TRANSACTION_RUN);
-    Callable<T> transactionCallable =
-        (getOptions().getOpenTelemetryOptions().isTracingEnabled()
-            ? new TracedReadWriteTransactionCallable<T>(
-                this, callable, /* transactionOptions= */ null, span)
-            : new ReadWriteTransactionCallable<T>(this, callable, /* transactionOptions= */ null));
-    try (Scope ignored = span.makeCurrent()) {
-      return RetryHelper.runWithRetries(
-          transactionCallable,
-          retrySettings,
-          TRANSACTION_EXCEPTION_HANDLER,
-          getOptions().getClock());
-    } catch (RetryHelperException e) {
-      span.end(e);
-      throw DatastoreException.translateAndThrow(e);
-    } finally {
-      span.end();
-    }
+    return runInTransaction(callable, /* transactionOptions= */ null);
   }
 
   @Override
   public <T> T runInTransaction(
       final TransactionCallable<T> callable, TransactionOptions transactionOptions) {
     TraceUtil.Span span = otelTraceUtil.startSpan(SPAN_NAME_TRANSACTION_RUN);
 
-    Callable<T> transactionCallable =
-        (getOptions().getOpenTelemetryOptions().isTracingEnabled()
-            ? new TracedReadWriteTransactionCallable<T>(this, callable, transactionOptions, span)
-            : new ReadWriteTransactionCallable<T>(this, callable, transactionOptions));
+    ReadWriteTransactionCallable<T> baseCallable =
+        new ReadWriteTransactionCallable<>(this, callable, transactionOptions, metricsRecorder);
+
+    Callable<T> transactionCallable = baseCallable;
+    if (getOptions().getOpenTelemetryOptions().isTracingEnabled()) {
+      transactionCallable = new TracedReadWriteTransactionCallable<>(baseCallable, span);
+    }
 
+    String status = StatusCode.Code.OK.toString();
+    Stopwatch stopwatch = Stopwatch.createStarted();
     try (Scope ignored = span.makeCurrent()) {
       return RetryHelper.runWithRetries(
           transactionCallable,
           retrySettings,
           TRANSACTION_EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      status = DatastoreException.extractStatusCode(e);
       span.end(e);
       throw DatastoreException.translateAndThrow(e);
     } finally {
+      long latencyMs = stopwatch.elapsed(TimeUnit.MILLISECONDS);
+      Map<String, String> attributes = new HashMap<>();
+      attributes.put(TelemetryConstants.ATTRIBUTES_KEY_STATUS, status);
+      attributes.put(
+          TelemetryConstants.ATTRIBUTES_KEY_METHOD, TelemetryConstants.METHOD_TRANSACTION_RUN);
+      attributes.put(TelemetryConstants.ATTRIBUTES_KEY_PROJECT_ID, getOptions().getProjectId());
+      attributes.put(TelemetryConstants.ATTRIBUTES_KEY_DATABASE_ID, getOptions().getDatabaseId());
+      metricsRecorder.recordTransactionLatency(latencyMs, attributes);
       span.end();
     }
   }

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/DatastoreOptions.java

@@ -31,6 +31,7 @@
 import com.google.cloud.datastore.spi.v1.DatastoreRpc;
 import com.google.cloud.datastore.spi.v1.GrpcDatastoreRpc;
 import com.google.cloud.datastore.spi.v1.HttpDatastoreRpc;
+import com.google.cloud.datastore.telemetry.MetricsRecorder;
 import com.google.cloud.datastore.v1.DatastoreSettings;
 import com.google.cloud.grpc.GrpcTransportOptions;
 import com.google.cloud.http.HttpTransportOptions;
@@ -64,6 +65,7 @@ public class DatastoreOptions extends ServiceOptions<Datastore, DatastoreOptions
 
   private final transient @Nonnull DatastoreOpenTelemetryOptions openTelemetryOptions;
   private final transient @Nonnull com.google.cloud.datastore.telemetry.TraceUtil traceUtil;
+  private final transient @Nonnull MetricsRecorder metricsRecorder;
 
   public static class DefaultDatastoreFactory implements DatastoreFactory {
 
@@ -104,6 +106,11 @@ public DatastoreOpenTelemetryOptions getOpenTelemetryOptions() {
     return openTelemetryOptions;
   }
 
+  @Nonnull
+  MetricsRecorder getMetricsRecorder() {
+    return metricsRecorder;
+  }
+
   public static class Builder extends ServiceOptions.Builder<Datastore, DatastoreOptions, Builder> {
 
     private String namespace;
@@ -216,6 +223,7 @@ private DatastoreOptions(Builder builder) {
             ? builder.openTelemetryOptions
             : DatastoreOpenTelemetryOptions.newBuilder().build();
     this.traceUtil = com.google.cloud.datastore.telemetry.TraceUtil.getInstance(this);
+    this.metricsRecorder = MetricsRecorder.getInstance(openTelemetryOptions);
 
     namespace = MoreObjects.firstNonNull(builder.namespace, defaultNamespace());
     databaseId = MoreObjects.firstNonNull(builder.databaseId, DEFAULT_DATABASE_ID);

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/MetricsRecorder.java

@@ -16,14 +16,21 @@
 
 package com.google.cloud.datastore.telemetry;
 
+import com.google.api.core.InternalExtensionOnly;
 import com.google.cloud.datastore.DatastoreOpenTelemetryOptions;
 import io.opentelemetry.api.GlobalOpenTelemetry;
 import io.opentelemetry.api.OpenTelemetry;
 import java.util.Map;
 import javax.annotation.Nonnull;
 
-/** Interface to record specific metric operations. */
-interface MetricsRecorder {
+/**
+ * Interface to record specific metric operations.
+ *
+ * <p><b>Warning:</b> This is an internal API and is not intended for external use. Do not implement
+ * or extend this interface.
+ */
+@InternalExtensionOnly
+public interface MetricsRecorder {
   /** Records the total latency of a transaction in milliseconds. */
   void recordTransactionLatency(double latencyMs, Map<String, String> attributes);
 

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/OpenTelemetryMetricsRecorder.java

@@ -42,7 +42,7 @@ class OpenTelemetryMetricsRecorder implements MetricsRecorder {
     this.transactionLatency =
         meter
             .histogramBuilder(TelemetryConstants.SERVICE_NAME + "/transaction_latency")
-            .setDescription("Total latency for successful transaction operations")
+            .setDescription("Total latency of transaction operations")
             .setUnit("ms")
             .build();
 

java-datastore/google-cloud-datastore/src/main/java/com/google/cloud/datastore/telemetry/TelemetryConstants.java

@@ -33,6 +33,18 @@ public class TelemetryConstants {
   public static final String ATTRIBUTES_KEY_DEFERRED = "Deferred";
   public static final String ATTRIBUTES_KEY_MORE_RESULTS = "more_results";
 
+  /** Attribute key for the gRPC status code (e.g. "OK", "ABORTED", "UNAVAILABLE"). */
+  public static final String ATTRIBUTES_KEY_STATUS = "status";
+
+  /** Attribute key for the RPC method name (e.g. "Transaction.Run"). */
+  public static final String ATTRIBUTES_KEY_METHOD = "method";
+
+  /** Attribute key for the GCP project ID. */
+  public static final String ATTRIBUTES_KEY_PROJECT_ID = "project_id";
+
+  /** Attribute key for the Datastore database ID. */
+  public static final String ATTRIBUTES_KEY_DATABASE_ID = "database_id";
+
   /* TODO(lawrenceqiu): For now, these are a duplicate of method names in TraceUtil. Those will use these eventually */
   // Format is not SnakeCase to keep backward compatibility with the existing values TraceUtil spans
   public static final String METHOD_ALLOCATE_IDS = "AllocateIds";