feat(storage): wire veneer operation metrics in StorageImpl and GrpcStorageImpl

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

@@ -43,13 +43,24 @@ abstract class BaseStorageReadChannel<T> implements StorageReadChannel {
   private BufferHandle bufferHandle;
   private LazyReadChannel<?, T> lazyReadChannel;
   protected final ReentrantLock lock;
+  @Nullable private final StorageMetricsContext metricsContext;
+  @Nullable private final StorageClientMetrics storageClientMetrics;
 
   protected BaseStorageReadChannel(Decoder<T, BlobInfo> objectDecoder) {
+    this(objectDecoder, null, null);
+  }
+
+  protected BaseStorageReadChannel(
+      Decoder<T, BlobInfo> objectDecoder,
+      @Nullable StorageMetricsContext metricsContext,
+      @Nullable StorageClientMetrics storageClientMetrics) {
     this.objectDecoder = objectDecoder;
     this.result = SettableApiFuture.create();
     this.open = true;
     this.byteRangeSpec = ByteRangeSpec.nullRange();
     this.lock = new ReentrantLock();
+    this.metricsContext = metricsContext;
+    this.storageClientMetrics = storageClientMetrics;
   }
 
   @Override
@@ -76,14 +87,26 @@ public final boolean isOpen() {
   @Override
   public final void close() {
     lock.lock();
+    Throwable error = null;
     try {
-      open = false;
-      if (internalGetLazyChannel().isOpen()) {
-        ReadableByteChannel channel = internalGetLazyChannel().getChannel();
-        StorageException.wrapIOException(channel::close);
+      if (open) {
+        open = false;
+        if (internalGetLazyChannel().isOpen()) {
+          ReadableByteChannel channel = internalGetLazyChannel().getChannel();
+          StorageException.wrapIOException(channel::close);
+        }
       }
+    } catch (Throwable t) {
+      error = t;
+      throw t;
     } finally {
-      lock.unlock();
+      try {
+        if (metricsContext != null && storageClientMetrics != null) {
+          metricsContext.recordOperation(storageClientMetrics, error);
+        }
+      } finally {
+        lock.unlock();
+      }
     }
   }
 

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

@@ -44,6 +44,8 @@ abstract class BaseStorageWriteChannel<T> implements StorageWriteChannel {
   private int chunkSize;
   private LazyWriteChannel<T> lazyWriteChannel;
   private BufferHandle bufferHandle;
+  @Nullable private final StorageMetricsContext metricsContext;
+  @Nullable private final StorageClientMetrics storageClientMetrics;
 
   /**
    * This is tracked for compatibility with BlobWriteChannel, such that simply creating a writer
@@ -55,12 +57,21 @@ abstract class BaseStorageWriteChannel<T> implements StorageWriteChannel {
   protected boolean writeCalledAtLeastOnce;
 
   protected BaseStorageWriteChannel(Decoder<T, BlobInfo> objectDecoder) {
+    this(objectDecoder, null, null);
+  }
+
+  protected BaseStorageWriteChannel(
+      Decoder<T, BlobInfo> objectDecoder,
+      @Nullable StorageMetricsContext metricsContext,
+      @Nullable StorageClientMetrics storageClientMetrics) {
     this.objectDecoder = objectDecoder;
     this.result = SettableApiFuture.create();
     this.lock = new ReentrantLock();
     this.open = true;
     this.chunkSize = _16MiB;
     this.writeCalledAtLeastOnce = false;
+    this.metricsContext = metricsContext;
+    this.storageClientMetrics = storageClientMetrics;
   }
 
   @Override
@@ -90,16 +101,26 @@ public final boolean isOpen() {
   @Override
   public final void close() throws IOException {
     lock.lock();
+    Throwable error = null;
     try {
       if (open && !writeCalledAtLeastOnce) {
         this.write(ByteBuffer.allocate(0));
       }
       if (internalGetLazyChannel().isOpen()) {
         StorageException.wrapIOException(internalGetLazyChannel().getChannel()::close);
       }
+    } catch (Throwable t) {
+      error = t;
+      throw t;
     } finally {
       open = false;
-      lock.unlock();
+      try {
+        if (metricsContext != null && storageClientMetrics != null) {
+          metricsContext.recordOperation(storageClientMetrics, error);
+        }
+      } finally {
+        lock.unlock();
+      }
     }
   }
 

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

@@ -28,6 +28,7 @@
 import java.io.Serializable;
 import java.util.Map;
 import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class BlobReadChannelV2 extends BaseStorageReadChannel<StorageObject> {
 
@@ -40,7 +41,20 @@ final class BlobReadChannelV2 extends BaseStorageReadChannel<StorageObject> {
       StorageObject storageObject,
       Map<StorageRpc.Option, ?> opts,
       BlobReadChannelContext blobReadChannelContext) {
-    super(Conversions.json().blobInfo());
+    this(storageObject, opts, blobReadChannelContext, null);
+  }
+
+  BlobReadChannelV2(
+      StorageObject storageObject,
+      Map<StorageRpc.Option, ?> opts,
+      BlobReadChannelContext blobReadChannelContext,
+      @Nullable StorageMetricsContext metricsContext) {
+    super(
+        Conversions.json().blobInfo(),
+        metricsContext,
+        blobReadChannelContext.getStorageOptions() != null
+            ? blobReadChannelContext.getStorageOptions().getStorageClientMetrics()
+            : null);
     this.storageObject = storageObject;
     this.opts = opts;
     this.blobReadChannelContext = blobReadChannelContext;

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

@@ -25,14 +25,27 @@
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Objects;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class BlobWriteChannelV2 extends BaseStorageWriteChannel<StorageObject> {
 
   private final BlobReadChannelContext blobChannelContext;
   private final JsonResumableWrite start;
 
   BlobWriteChannelV2(BlobReadChannelContext blobChannelContext, JsonResumableWrite start) {
-    super(Conversions.json().blobInfo());
+    this(blobChannelContext, start, null);
+  }
+
+  BlobWriteChannelV2(
+      BlobReadChannelContext blobChannelContext,
+      JsonResumableWrite start,
+      @Nullable StorageMetricsContext metricsContext) {
+    super(
+        Conversions.json().blobInfo(),
+        metricsContext,
+        blobChannelContext.getStorageOptions() != null
+            ? blobChannelContext.getStorageOptions().getStorageClientMetrics()
+            : null);
     this.start = start;
     this.blobChannelContext = blobChannelContext;
   }

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,19 +43,33 @@ 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;
   }
 
   @Override
   public Blob getResult() {
     while (!isDone()) {
       copyChunk();
     }
+    if (metricsContext != null) {
+      metricsContext.recordOperation(storage.getOptions().getStorageClientMetrics(), null);
+    }
     BlobInfo info = Conversions.grpc().blobInfo().decode(mostRecentResponse.getResource());
     return info.asBlob(storage);
   }
@@ -78,8 +94,19 @@ 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());
+      StorageMetricsContext chunkContext =
+          metricsContext != null
+              ? metricsContext.asSubAttemptContext()
+              : storage.getOptions().createMetricsContext("RewriteObject");
+      GrpcCallContext retryContext = Retrying.newCallContext(chunkContext);
+      RewriteResponse resp =
+          retrier.run(() -> callable.call(req, retryContext), Decoder.identity(), chunkContext);
+      if (resp != null) {
+        mostRecentResponse = resp;
+      }
+      if (isDone() && metricsContext != null) {
+        metricsContext.recordOperation(storage.getOptions().getStorageClientMetrics(), null);
+      }
     }
   }
 

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;

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;