feat: add transaction metrics collection for latency and per-attempt count

- Add ATTRIBUTES_KEY_STATUS and ATTRIBUTES_KEY_METHOD_NAME to TelemetryConstants
- Make MetricsRecorder public with @InternalExtensionOnly annotation
- Add MetricsRecorder field and getter to DatastoreOptions
- Wire MetricsRecorder into DatastoreImpl for transaction metrics
- Refactor TracedReadWriteTransactionCallable to delegate to ReadWriteTransactionCallable
- Record per-attempt transaction count with gRPC status code and method name
- Record overall transaction latency using Guava Stopwatch
- Add unit tests for OpenTelemetryMetricsRecorder, MetricsRecorder, and DatastoreImpl

Author: lqiu96

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

@@ -45,10 +45,13 @@
 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;
@@ -65,6 +68,7 @@
 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 +77,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 +94,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 +128,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 +174,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;
     }
 
@@ -235,9 +214,11 @@ public T call() throws DatastoreException {
         transaction = datastore.newTransaction(options);
         T value = callable.run(transaction);
         transaction.commit();
+        recordAttempt("OK");
         return value;
       } catch (Exception ex) {
         transaction.rollback();
+        recordAttempt(extractStatus(ex));
         throw DatastoreException.propagateUserException(ex);
       } finally {
         if (transaction.isActive()) {
@@ -249,54 +230,92 @@ 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_NAME,
+          TelemetryConstants.METHOD_TRANSACTION_RUN);
+      metricsRecorder.recordTransactionAttemptCount(1, attributes);
+    }
+
+    /**
+     * Extracts the gRPC status code from the given exception. Falls back to "UNKNOWN" if the
+     * status cannot be determined.
+     */
+    private static String extractStatus(Exception ex) {
+      if (ex instanceof DatastoreException) {
+        String reason = ((DatastoreException) ex).getReason();
+        if (reason != null && !reason.isEmpty()) {
+          return reason;
+        }
+      }
+      return "UNKNOWN";
+    }
   }
 
   @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);
+    Stopwatch stopwatch = Stopwatch.createStarted();
 
+    ReadWriteTransactionCallable<T> baseCallable =
+        new ReadWriteTransactionCallable<>(this, callable, transactionOptions, metricsRecorder);
     Callable<T> transactionCallable =
-        (getOptions().getOpenTelemetryOptions().isTracingEnabled()
-            ? new TracedReadWriteTransactionCallable<T>(this, callable, transactionOptions, span)
-            : new ReadWriteTransactionCallable<T>(this, callable, transactionOptions));
+        getOptions().getOpenTelemetryOptions().isTracingEnabled()
+            ? new TracedReadWriteTransactionCallable<>(baseCallable, span)
+            : baseCallable;
 
+    String status = "OK";
     try (Scope ignored = span.makeCurrent()) {
       return RetryHelper.runWithRetries(
           transactionCallable,
           retrySettings,
           TRANSACTION_EXCEPTION_HANDLER,
           getOptions().getClock());
     } catch (RetryHelperException e) {
+      status = extractGrpcStatusCode(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_NAME,
+          TelemetryConstants.METHOD_TRANSACTION_RUN);
+      metricsRecorder.recordTransactionLatency(latencyMs, attributes);
       span.end();
     }
   }
 
+  /**
+   * Extracts the gRPC status code from a {@link RetryHelperException}. The underlying cause is
+   * expected to be a {@link DatastoreException} which carries the gRPC reason. Falls back to
+   * "UNKNOWN" if the status cannot be determined.
+   */
+  private static String extractGrpcStatusCode(RetryHelperException e) {
+    Throwable cause = e.getCause();
+    if (cause instanceof DatastoreException) {
+      String reason = ((DatastoreException) cause).getReason();
+      if (reason != null && !reason.isEmpty()) {
+        return reason;
+      }
+    }
+    return "UNKNOWN";
+  }
+
   @Override
   public <T> QueryResults<T> run(Query<T> query) {
     return run(Optional.empty(), query, null);

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/TelemetryConstants.java

@@ -33,6 +33,12 @@ 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_NAME = "method_name";
+
   /* 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";