feat(storage): wire operation metrics in gRPC veneer

java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GapicCopyWriter.java

@@ -23,6 +23,7 @@
 import com.google.cloud.storage.Retrying.RetrierWithAlg;
 import com.google.storage.v2.RewriteObjectRequest;
 import com.google.storage.v2.RewriteResponse;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class GapicCopyWriter extends CopyWriter {
 
@@ -32,6 +33,7 @@ final class GapicCopyWriter extends CopyWriter {
   private final RetrierWithAlg retrier;
   private final RewriteObjectRequest originalRequest;
   private final RewriteResponse initialResponse;
+  @Nullable private final StorageMetricsContext metricsContext;
 
   private RewriteResponse mostRecentResponse;
 
@@ -41,18 +43,41 @@ final class GapicCopyWriter extends CopyWriter {
       RetrierWithAlg retrier,
       RewriteObjectRequest originalRequest,
       RewriteResponse initialResponse) {
+    this(storage, callable, retrier, originalRequest, initialResponse, null);
+  }
+
+  GapicCopyWriter(
+      GrpcStorageImpl storage,
+      UnaryCallable<RewriteObjectRequest, RewriteResponse> callable,
+      RetrierWithAlg retrier,
+      RewriteObjectRequest originalRequest,
+      RewriteResponse initialResponse,
+      @Nullable StorageMetricsContext metricsContext) {
     this.storage = storage;
     this.callable = callable;
     this.retrier = retrier;
     this.initialResponse = initialResponse;
     this.mostRecentResponse = initialResponse;
     this.originalRequest = originalRequest;
+    this.metricsContext = metricsContext;
+  }
+
+  private void recordOperation(@Nullable Throwable error) {
+    if (metricsContext != null) {
+      metricsContext.recordOperation(storage.getOptions().getStorageClientMetrics(), error);
+    }
   }
 
   @Override
   public Blob getResult() {
-    while (!isDone()) {
-      copyChunk();
+    try {
+      while (!isDone()) {
+        copyChunk();
+      }
+      recordOperation(null);
+    } catch (Throwable t) {
+      recordOperation(t);
+      throw t;
     }
     BlobInfo info = Conversions.grpc().blobInfo().decode(mostRecentResponse.getResource());
     return info.asBlob(storage);
@@ -76,10 +101,29 @@ public long getTotalBytesCopied() {
   @Override
   public void copyChunk() {
     if (!isDone()) {
-      RewriteObjectRequest req =
-          originalRequest.toBuilder().setRewriteToken(mostRecentResponse.getRewriteToken()).build();
-      GrpcCallContext retryContext = Retrying.newCallContext();
-      mostRecentResponse = retrier.run(() -> callable.call(req, retryContext), Decoder.identity());
+      try {
+        RewriteObjectRequest req =
+            originalRequest.toBuilder()
+                .setRewriteToken(mostRecentResponse.getRewriteToken())
+                .build();
+        StorageMetricsContext chunkContext =
+            metricsContext != null
+                ? metricsContext
+                : StorageClientMetrics.createContext(
+                    storage.getOptions().getStorageClientMetrics(), "RewriteObject");
+        GrpcCallContext retryContext = Retrying.newCallContext(chunkContext);
+        RewriteResponse resp =
+            retrier.run(() -> callable.call(req, retryContext), Decoder.identity(), chunkContext);
+        if (resp != null) {
+          mostRecentResponse = resp;
+        }
+        if (isDone()) {
+          recordOperation(null);
+        }
+      } catch (Throwable t) {
+        recordOperation(t);
+        throw t;
+      }
     }
   }
 

java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcBlobReadChannel.java

@@ -26,6 +26,7 @@
 import com.google.storage.v2.ReadObjectRequest;
 import com.google.storage.v2.ReadObjectResponse;
 import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class GrpcBlobReadChannel extends BaseStorageReadChannel<Object> {
 
@@ -41,7 +42,18 @@ final class GrpcBlobReadChannel extends BaseStorageReadChannel<Object> {
       ResultRetryAlgorithm<?> resultRetryAlgorithm,
       ReadObjectRequest request,
       boolean autoGzipDecompression) {
-    super(Conversions.grpc().blobInfo());
+    this(read, retrier, resultRetryAlgorithm, request, autoGzipDecompression, null, null);
+  }
+
+  GrpcBlobReadChannel(
+      ZeroCopyServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> read,
+      Retrier retrier,
+      ResultRetryAlgorithm<?> resultRetryAlgorithm,
+      ReadObjectRequest request,
+      boolean autoGzipDecompression,
+      @Nullable StorageMetricsContext metricsContext,
+      @Nullable StorageClientMetrics storageClientMetrics) {
+    super(Conversions.grpc().blobInfo(), metricsContext, storageClientMetrics);
     this.read = read;
     this.retrier = retrier;
     this.resultRetryAlgorithm = resultRetryAlgorithm;
@@ -61,7 +73,8 @@ protected LazyReadChannel<?, Object> newLazyReadChannel() {
           ReadableByteChannelSessionBuilder b =
               ResumableMedia.gapic()
                   .read()
-                  .byteChannel(read, retrier, resultRetryAlgorithm)
+                  .byteChannel(
+                      read, retrier.withMetricsContext(getMetricsContext()), resultRetryAlgorithm)
                   .setHasher(Hasher.readHasher())
                   .setAutoGzipDecompression(autoGzipDecompression);
           BufferHandle bufferHandle = getBufferHandle();

java-storage/google-cloud-storage/src/main/java/com/google/cloud/storage/GrpcBlobWriteChannel.java

@@ -24,6 +24,7 @@
 import com.google.storage.v2.WriteObjectRequest;
 import com.google.storage.v2.WriteObjectResponse;
 import java.util.function.Supplier;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class GrpcBlobWriteChannel extends BaseStorageWriteChannel<WriteObjectResponse> {
 
@@ -37,7 +38,20 @@ final class GrpcBlobWriteChannel extends BaseStorageWriteChannel<WriteObjectResp
       RetrierWithAlg retrier,
       Supplier<ApiFuture<ResumableWrite>> start,
       Hasher hasher) {
-    super(Conversions.grpc().blobInfo().compose(WriteObjectResponse::getResource));
+    this(write, retrier, start, hasher, null, null);
+  }
+
+  GrpcBlobWriteChannel(
+      ClientStreamingCallable<WriteObjectRequest, WriteObjectResponse> write,
+      RetrierWithAlg retrier,
+      Supplier<ApiFuture<ResumableWrite>> start,
+      Hasher hasher,
+      @Nullable StorageMetricsContext metricsContext,
+      @Nullable StorageClientMetrics storageClientMetrics) {
+    super(
+        Conversions.grpc().blobInfo().compose(WriteObjectResponse::getResource),
+        metricsContext,
+        storageClientMetrics);
     this.write = write;
     this.retrier = retrier;
     this.start = start;