feat(storage): add gcp.storage.client.response.body.size metric for read operations

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

@@ -45,6 +45,7 @@ abstract class BaseStorageReadChannel<T> implements StorageReadChannel {
   protected final ReentrantLock lock;
   @Nullable private final StorageMetricsContext metricsContext;
   @Nullable private final StorageClientMetrics storageClientMetrics;
+  @Nullable private Throwable lastReadError;
 
   protected BaseStorageReadChannel(Decoder<T, BlobInfo> objectDecoder) {
     this(objectDecoder, null, null);
@@ -102,7 +103,8 @@ public final void close() {
     } finally {
       try {
         if (metricsContext != null && storageClientMetrics != null) {
-          metricsContext.recordOperation(storageClientMetrics, error);
+          metricsContext.recordOperation(
+              storageClientMetrics, error != null ? error : lastReadError);
         }
       } finally {
         lock.unlock();
@@ -151,8 +153,12 @@ public final int read(ByteBuffer dst) throws IOException {
           return -1;
         }
         int read = tmp.read(dst);
+        lastReadError = null;
         if (read != -1) {
           byteRangeSpec = byteRangeSpec.withShiftBeginOffset(read);
+          if (read > 0 && metricsContext != null) {
+            metricsContext.addResponseBytes(read);
+          }
         }
         return read;
       } catch (StorageException e) {
@@ -163,12 +169,16 @@ public final int read(ByteBuffer dst) throws IOException {
           // be removed in the next major version.
           return -1;
         } else {
+          lastReadError = e;
           throw new IOException(e);
         }
       } catch (IOException e) {
+        lastReadError = e;
         throw e;
       } catch (Exception e) {
-        throw new IOException(StorageException.coalesce(e));
+        Throwable coalesced = StorageException.coalesce(e);
+        lastReadError = coalesced;
+        throw new IOException(coalesced);
       }
     } finally {
       lock.unlock();

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

@@ -17,17 +17,35 @@
 package com.google.cloud.storage;
 
 import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
 import com.google.api.core.ApiFutures;
+import com.google.cloud.storage.ZeroCopySupport.DisposableByteString;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.ScatteringByteChannel;
+import java.nio.channels.SeekableByteChannel;
+import org.checkerframework.checker.nullness.qual.Nullable;
 
 final class BlobReadSessionAdapter implements BlobReadSession {
 
   @VisibleForTesting final ObjectReadSession session;
+  @Nullable private final StorageMetricsContext metricsContext;
+  @Nullable private final StorageClientMetrics storageClientMetrics;
+  @Nullable private volatile Throwable lastOperationError;
 
   BlobReadSessionAdapter(ObjectReadSession session) {
+    this(session, null, null);
+  }
+
+  BlobReadSessionAdapter(
+      ObjectReadSession session,
+      @Nullable StorageMetricsContext metricsContext,
+      @Nullable StorageClientMetrics storageClientMetrics) {
     this.session = session;
+    this.metricsContext = metricsContext;
+    this.storageClientMetrics = storageClientMetrics;
   }
 
   @Override
@@ -40,23 +58,208 @@ public BlobInfo getBlobInfo() {
   @SuppressWarnings({"rawtypes", "unchecked"})
   @Override
   public <Projection> Projection readAs(ReadProjectionConfig<Projection> config) {
-    Projection projection = session.readAs(config);
+    Projection projection;
+    try {
+      projection = session.readAs(config);
+    } catch (Throwable t) {
+      lastOperationError = t;
+      throw t;
+    }
     if (projection instanceof ApiFuture) {
-      ApiFuture apiFuture = (ApiFuture) projection;
-      return (Projection) StorageException.coalesceAsync(apiFuture);
+      ApiFuture apiFuture = StorageException.coalesceAsync((ApiFuture) projection);
+      if (metricsContext != null) {
+        ApiFutures.addCallback(
+            apiFuture,
+            new ApiFutureCallback<Object>() {
+              @Override
+              public void onSuccess(Object result) {
+                if (result instanceof byte[]) {
+                  int len = ((byte[]) result).length;
+                  if (len > 0) {
+                    metricsContext.addResponseBytes(len);
+                  }
+                } else if (result instanceof DisposableByteString) {
+                  int len = ((DisposableByteString) result).byteString().size();
+                  if (len > 0) {
+                    metricsContext.addResponseBytes(len);
+                  }
+                }
+              }
+
+              @Override
+              public void onFailure(Throwable t) {
+                if (!(t instanceof AsyncSessionClosedException)
+                    && !(t.getCause() instanceof AsyncSessionClosedException)) {
+                  lastOperationError = t;
+                }
+              }
+            },
+            MoreExecutors.directExecutor());
+      }
+      return (Projection) apiFuture;
+    }
+    if (metricsContext != null && projection instanceof SeekableByteChannel) {
+      return (Projection) wrapSeekableByteChannel((SeekableByteChannel) projection);
+    }
+    if (metricsContext != null && projection instanceof ScatteringByteChannel) {
+      return (Projection) wrapScatteringByteChannel((ScatteringByteChannel) projection);
     }
     return projection;
   }
 
   @Override
   public void close() throws IOException {
-    session.close();
+    Throwable priorError = lastOperationError;
+    Throwable error = null;
+    try {
+      session.close();
+    } catch (Throwable t) {
+      error = t;
+      throw t;
+    } finally {
+      if (metricsContext != null && storageClientMetrics != null) {
+        metricsContext.recordOperation(storageClientMetrics, error != null ? error : priorError);
+      }
+    }
+  }
+
+  private SeekableByteChannel wrapSeekableByteChannel(SeekableByteChannel delegate) {
+    return new SeekableByteChannel() {
+      @Override
+      public int read(ByteBuffer dst) throws IOException {
+        try {
+          int read = delegate.read(dst);
+          if (read > 0 && metricsContext != null) {
+            metricsContext.addResponseBytes(read);
+          }
+          return read;
+        } catch (IOException | RuntimeException | Error e) {
+          lastOperationError = e;
+          throw e;
+        }
+      }
+
+      @Override
+      public int write(ByteBuffer src) throws IOException {
+        return delegate.write(src);
+      }
+
+      @Override
+      public long position() throws IOException {
+        return delegate.position();
+      }
+
+      @Override
+      public SeekableByteChannel position(long newPosition) throws IOException {
+        delegate.position(newPosition);
+        return this;
+      }
+
+      @Override
+      public long size() throws IOException {
+        return delegate.size();
+      }
+
+      @Override
+      public SeekableByteChannel truncate(long size) throws IOException {
+        delegate.truncate(size);
+        return this;
+      }
+
+      @Override
+      public boolean isOpen() {
+        return delegate.isOpen();
+      }
+
+      @Override
+      public void close() throws IOException {
+        delegate.close();
+      }
+    };
+  }
+
+  private ScatteringByteChannel wrapScatteringByteChannel(ScatteringByteChannel delegate) {
+    return new ScatteringByteChannel() {
+      @Override
+      public int read(ByteBuffer dst) throws IOException {
+        try {
+          int read = delegate.read(dst);
+          if (read > 0 && metricsContext != null) {
+            metricsContext.addResponseBytes(read);
+          }
+          return read;
+        } catch (IOException | RuntimeException | Error e) {
+          lastOperationError = e;
+          throw e;
+        }
+      }
+
+      @Override
+      public long read(ByteBuffer[] dsts, int offset, int length) throws IOException {
+        try {
+          long read = delegate.read(dsts, offset, length);
+          if (read > 0 && metricsContext != null) {
+            metricsContext.addResponseBytes(read);
+          }
+          return read;
+        } catch (IOException | RuntimeException | Error e) {
+          lastOperationError = e;
+          throw e;
+        }
+      }
+
+      @Override
+      public long read(ByteBuffer[] dsts) throws IOException {
+        try {
+          long read = delegate.read(dsts);
+          if (read > 0 && metricsContext != null) {
+            metricsContext.addResponseBytes(read);
+          }
+          return read;
+        } catch (IOException | RuntimeException | Error e) {
+          lastOperationError = e;
+          throw e;
+        }
+      }
+
+      @Override
+      public boolean isOpen() {
+        return delegate.isOpen();
+      }
+
+      @Override
+      public void close() throws IOException {
+        delegate.close();
+      }
+    };
   }
 
   static ApiFuture<BlobReadSession> wrap(ApiFuture<ObjectReadSession> session) {
+    return wrap(session, null, null);
+  }
+
+  static ApiFuture<BlobReadSession> wrap(
+      ApiFuture<ObjectReadSession> session,
+      @Nullable StorageMetricsContext metricsContext,
+      @Nullable StorageClientMetrics storageClientMetrics) {
+    ApiFuture<ObjectReadSession> coalesced = StorageException.coalesceAsync(session);
+    if (metricsContext != null && storageClientMetrics != null) {
+      ApiFutures.addCallback(
+          coalesced,
+          new ApiFutureCallback<ObjectReadSession>() {
+            @Override
+            public void onSuccess(ObjectReadSession result) {}
+
+            @Override
+            public void onFailure(Throwable t) {
+              metricsContext.recordOperation(storageClientMetrics, t);
+            }
+          },
+          MoreExecutors.directExecutor());
+    }
     return ApiFutures.transform(
-        StorageException.coalesceAsync(session),
-        BlobReadSessionAdapter::new,
+        coalesced,
+        s -> new BlobReadSessionAdapter(s, metricsContext, storageClientMetrics),
         MoreExecutors.directExecutor());
   }
 }

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

@@ -819,6 +819,9 @@ public byte[] readAllBytes(BlobId blob, BlobSourceOption... options) {
       throw e;
     } finally {
       if (metricsContext != null) {
+        if (error == null) {
+          metricsContext.setResponseBytes(baos.size());
+        }
         metricsContext.recordOperation(getOptions().getStorageClientMetrics(), error);
       }
     }
@@ -864,7 +867,7 @@ public void downloadTo(BlobId blob, Path path, BlobSourceOption... options) {
 
     Throwable error = null;
     try (UnbufferedReadableByteChannel r = session.open();
-        WritableByteChannel w = Files.newByteChannel(path, WRITE_OPS)) {
+        WritableByteChannel w = countBytes(Files.newByteChannel(path, WRITE_OPS), metricsContext)) {
       ByteStreams.copy(r, w);
     } catch (ApiException | IOException e) {
       error = e;
@@ -887,7 +890,7 @@ public void downloadTo(BlobId blob, OutputStream outputStream, BlobSourceOption.
 
     Throwable error = null;
     try (UnbufferedReadableByteChannel r = session.open();
-        WritableByteChannel w = Channels.newChannel(outputStream)) {
+        WritableByteChannel w = countBytes(Channels.newChannel(outputStream), metricsContext)) {
       ByteStreams.copy(r, w);
     } catch (ApiException | IOException e) {
       error = e;
@@ -1710,10 +1713,22 @@ public ApiFuture<BlobReadSession> blobReadSession(BlobId id, BlobSourceOption...
     opts.bidiReadObjectRequest().apply(b);
     BidiReadObjectRequest req = b.build();
 
-    GrpcCallContext context = opts.grpcMetadataMapper().apply(GrpcCallContext.createDefault());
-    ApiFuture<ObjectReadSession> session = storageDataClient.readSession(req, context);
-
-    return BlobReadSessionAdapter.wrap(session);
+    StorageMetricsContext metricsContext = getOptions().createMultiStepMetricsContext("ReadObject");
+    GrpcCallContext baseContext =
+        metricsContext != null
+            ? metricsContext.attach(GrpcCallContext.createDefault())
+            : GrpcCallContext.createDefault();
+    GrpcCallContext context = opts.grpcMetadataMapper().apply(baseContext);
+    try {
+      ApiFuture<ObjectReadSession> session = storageDataClient.readSession(req, context);
+      return BlobReadSessionAdapter.wrap(
+          session, metricsContext, getOptions().getStorageClientMetrics());
+    } catch (Throwable t) {
+      if (metricsContext != null) {
+        metricsContext.recordOperation(getOptions().getStorageClientMetrics(), t);
+      }
+      throw t;
+    }
   }
 
   @Override
@@ -2421,4 +2436,31 @@ private ZeroCopyServerStreamingCallable<ReadObjectRequest, ReadObjectResponse> r
         storageClient.readObjectCallable().withDefaultCallContext(grpcCallContext),
         responseContentLifecycleManager);
   }
+
+  private static WritableByteChannel countBytes(
+      WritableByteChannel delegate, @Nullable StorageMetricsContext metricsContext) {
+    if (metricsContext == null) {
+      return delegate;
+    }
+    return new WritableByteChannel() {
+      @Override
+      public int write(ByteBuffer src) throws IOException {
+        int written = delegate.write(src);
+        if (written > 0) {
+          metricsContext.addResponseBytes(written);
+        }
+        return written;
+      }
+
+      @Override
+      public boolean isOpen() {
+        return delegate.isOpen();
+      }
+
+      @Override
+      public void close() throws IOException {
+        delegate.close();
+      }
+    };
+  }
 }