feat: implement BufferToDiskThenUpload BlobWriteSessionConfig

There are scenarios in which disk space is more plentiful than memory space. This new BlobWriteSessionConfig allows augmenting an instance of storage to prefer buffering to disk rather than keeping things in memory.

Once the file on disk is closed, the entire file will then be uploaded to GCS.

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

@@ -19,6 +19,11 @@
 import com.google.api.core.BetaApi;
 import com.google.cloud.storage.GrpcStorageOptions.GrpcStorageDefaults;
 import com.google.cloud.storage.Storage.BlobWriteOption;
+import com.google.common.collect.ImmutableList;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Collection;
 
 /**
  * Factory class to select and construct {@link BlobWriteSessionConfig}s.
@@ -46,4 +51,54 @@ private BlobWriteSessionConfigs() {}
   public static DefaultBlobWriteSessionConfig getDefault() {
     return new DefaultBlobWriteSessionConfig(ByteSizeConstants._16MiB);
   }
+
+  /**
+   * Create a new {@link BlobWriteSessionConfig} which will first buffer the content of the object
+   * to a temporary file under {@code java.io.tmpdir}.
+   *
+   * <p>Once the file on disk is closed, the entire file will then be uploaded to Google Cloud
+   * Storage.
+   *
+   * @see Storage#blobWriteSession(BlobInfo, BlobWriteOption...)
+   * @see GrpcStorageOptions.Builder#setBlobWriteSessionConfig(BlobWriteSessionConfig)
+   */
+  @BetaApi
+  public static BlobWriteSessionConfig bufferToTempDirThenUpload() throws IOException {
+    return bufferToDiskThenUpload(
+        Paths.get(System.getProperty("java.io.tmpdir"), "google-cloud-storage"));
+  }
+
+  /**
+   * Create a new {@link BlobWriteSessionConfig} which will first buffer the content of the object
+   * to a temporary file under the specified {@code path}.
+   *
+   * <p>Once the file on disk is closed, the entire file will then be uploaded to Google Cloud
+   * Storage.
+   *
+   * @see Storage#blobWriteSession(BlobInfo, BlobWriteOption...)
+   * @see GrpcStorageOptions.Builder#setBlobWriteSessionConfig(BlobWriteSessionConfig)
+   */
+  @BetaApi
+  public static BufferToDiskThenUpload bufferToDiskThenUpload(Path path) throws IOException {
+    return bufferToDiskThenUpload(ImmutableList.of(path));
+  }
+
+  /**
+   * Create a new {@link BlobWriteSessionConfig} which will first buffer the content of the object
+   * to a temporary file under one of the specified {@code paths}.
+   *
+   * <p>Once the file on disk is closed, the entire file will then be uploaded to Google Cloud
+   * Storage.
+   *
+   * <p>The specifics of how the work is spread across multiple paths is undefined and subject to
+   * change.
+   *
+   * @see Storage#blobWriteSession(BlobInfo, BlobWriteOption...)
+   * @see GrpcStorageOptions.Builder#setBlobWriteSessionConfig(BlobWriteSessionConfig)
+   */
+  @BetaApi
+  public static BufferToDiskThenUpload bufferToDiskThenUpload(Collection<Path> paths)
+      throws IOException {
+    return new BufferToDiskThenUpload(ImmutableList.copyOf(paths), false);
+  }
 }

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

@@ -0,0 +1,194 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.BetaApi;
+import com.google.api.core.InternalApi;
+import com.google.api.core.SettableApiFuture;
+import com.google.cloud.storage.Conversions.Decoder;
+import com.google.cloud.storage.Storage.BlobWriteOption;
+import com.google.cloud.storage.UnifiedOpts.ObjectTargetOpt;
+import com.google.cloud.storage.UnifiedOpts.Opts;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableList;
+import com.google.common.util.concurrent.MoreExecutors;
+import com.google.storage.v2.WriteObjectResponse;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.time.Clock;
+import java.time.Duration;
+import java.util.Collection;
+import javax.annotation.concurrent.Immutable;
+
+/**
+ * There are scenarios in which disk space is more plentiful than memory space. This new {@link
+ * BlobWriteSessionConfig} allows augmenting an instance of storage to produce {@link
+ * BlobWriteSession}s which will buffer to disk rather than holding things in memory.
+ *
+ * <p>Once the file on disk is closed, the entire file will then be uploaded to GCS.
+ *
+ * @see Storage#blobWriteSession(BlobInfo, BlobWriteOption...)
+ * @see GrpcStorageOptions.Builder#setBlobWriteSessionConfig(BlobWriteSessionConfig)
+ * @see BlobWriteSessionConfigs#bufferToDiskThenUpload(Path)
+ * @see BlobWriteSessionConfigs#bufferToDiskThenUpload(Collection)
+ */
+@Immutable
+@BetaApi
+public final class BufferToDiskThenUpload extends BlobWriteSessionConfig {
+
+  private final ImmutableList<Path> paths;
+  private final boolean includeLoggingSink;
+
+  @InternalApi
+  BufferToDiskThenUpload(ImmutableList<Path> paths, boolean includeLoggingSink) throws IOException {
+    this.paths = paths;
+    this.includeLoggingSink = includeLoggingSink;
+  }
+
+  @VisibleForTesting
+  @InternalApi
+  BufferToDiskThenUpload withIncludeLoggingSink() throws IOException {
+    return new BufferToDiskThenUpload(paths, true);
+  }
+
+  @InternalApi
+  @Override
+  WriterFactory createFactory(Clock clock) throws IOException {
+    Duration window = Duration.ofMinutes(10);
+    RecoveryFileManager recoveryFileManager =
+        RecoveryFileManager.of(paths, getRecoverVolumeSinkFactory(clock, window));
+    ThroughputSink gcs = ThroughputSink.windowed(ThroughputMovingWindow.of(window), clock);
+    gcs = includeLoggingSink ? ThroughputSink.tee(ThroughputSink.logged("gcs", clock), gcs) : gcs;
+    return new Factory(recoveryFileManager, clock, gcs);
+  }
+
+  private RecoveryFileManager.RecoverVolumeSinkFactory getRecoverVolumeSinkFactory(
+      Clock clock, Duration window) {
+    return path -> {
+      ThroughputSink windowed = ThroughputSink.windowed(ThroughputMovingWindow.of(window), clock);
+      if (includeLoggingSink) {
+        return ThroughputSink.tee(
+            ThroughputSink.logged(path.toAbsolutePath().toString(), clock), windowed);
+      } else {
+        return windowed;
+      }
+    };
+  }
+
+  private static final class Factory implements WriterFactory {
+
+    private final RecoveryFileManager recoveryFileManager;
+    private final Clock clock;
+    private final ThroughputSink gcs;
+
+    private Factory(RecoveryFileManager recoveryFileManager, Clock clock, ThroughputSink gcs) {
+      this.recoveryFileManager = recoveryFileManager;
+      this.clock = clock;
+      this.gcs = gcs;
+    }
+
+    @InternalApi
+    @Override
+    public WritableByteChannelSession<?, BlobInfo> writeSession(
+        StorageInternal storage,
+        BlobInfo info,
+        Opts<ObjectTargetOpt> opts,
+        Decoder<WriteObjectResponse, BlobInfo> d) {
+      return new Factory.WriteToFileThenUpload(
+          storage, info, opts, recoveryFileManager.newRecoveryFile(info));
+    }
+
+    private final class WriteToFileThenUpload
+        implements WritableByteChannelSession<WritableByteChannel, BlobInfo> {
+
+      private final StorageInternal storage;
+      private final BlobInfo info;
+      private final Opts<ObjectTargetOpt> opts;
+      private final RecoveryFile rf;
+      private final SettableApiFuture<BlobInfo> result;
+
+      private WriteToFileThenUpload(
+          StorageInternal storage, BlobInfo info, Opts<ObjectTargetOpt> opts, RecoveryFile rf) {
+        this.info = info;
+        this.opts = opts;
+        this.rf = rf;
+        this.storage = storage;
+        this.result = SettableApiFuture.create();
+      }
+
+      @Override
+      public ApiFuture<WritableByteChannel> openAsync() {
+        try {
+          ApiFuture<WritableByteChannel> f = ApiFutures.immediateFuture(rf.writer());
+          return ApiFutures.transform(
+              f, Factory.WriteToFileThenUpload.Flusher::new, MoreExecutors.directExecutor());
+        } catch (IOException e) {
+          throw StorageException.coalesce(e);
+        }
+      }
+
+      @Override
+      public ApiFuture<BlobInfo> getResult() {
+        return result;
+      }
+
+      private final class Flusher implements WritableByteChannel {
+
+        private final WritableByteChannel delegate;
+
+        private Flusher(WritableByteChannel delegate) {
+          this.delegate = delegate;
+        }
+
+        @Override
+        public int write(ByteBuffer src) throws IOException {
+          return delegate.write(src);
+        }
+
+        @Override
+        public boolean isOpen() {
+          return delegate.isOpen();
+        }
+
+        @Override
+        public void close() throws IOException {
+          delegate.close();
+          try (RecoveryFile rf = Factory.WriteToFileThenUpload.this.rf) {
+            Path path = rf.getPath();
+            long size = Files.size(path);
+            ThroughputSink.computeThroughput(
+                clock,
+                gcs,
+                size,
+                () -> {
+                  BlobInfo blob = storage.internalCreateFrom(path, info, opts);
+                  result.set(blob);
+                });
+          } catch (StorageException | IOException e) {
+            result.setException(e);
+            throw e;
+          }
+        }
+      }
+    }
+  }
+}

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

@@ -152,7 +152,8 @@
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 @BetaApi
-final class GrpcStorageImpl extends BaseService<StorageOptions> implements StorageInternal {
+final class GrpcStorageImpl extends BaseService<StorageOptions>
+    implements Storage, StorageInternal {
 
   private static final byte[] ZERO_BYTES = new byte[0];
   private static final Set<OpenOption> READ_OPS = ImmutableSet.of(StandardOpenOption.READ);

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

@@ -0,0 +1,100 @@
+/*
+ * Copyright 2023 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.collect.ImmutableSet;
+import java.io.IOException;
+import java.nio.channels.SeekableByteChannel;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Files;
+import java.nio.file.OpenOption;
+import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.Set;
+
+/**
+ * When uploading to GCS, there are times when memory buffers are not preferable. This class
+ * encapsulates the logic and lifecycle for a file written to local disk which can be used for
+ * upload recovery in the case an upload is interrupted.
+ */
+final class RecoveryFile implements AutoCloseable {
+  private static final Set<OpenOption> writeOps =
+      ImmutableSet.of(StandardOpenOption.CREATE, StandardOpenOption.WRITE);
+  private static final Set<OpenOption> readOps = ImmutableSet.of(StandardOpenOption.READ);
+
+  private final Path path;
+  private final ThroughputSink throughputSink;
+  private final Runnable onCloseCallback;
+
+  RecoveryFile(Path path, ThroughputSink throughputSink, Runnable onCloseCallback) {
+    this.path = path;
+    this.throughputSink = throughputSink;
+    this.onCloseCallback = onCloseCallback;
+  }
+
+  public Path getPath() {
+    return path;
+  }
+
+  public Path touch() throws IOException {
+    return Files.createFile(path);
+  }
+
+  public SeekableByteChannel reader() throws IOException {
+    return Files.newByteChannel(path, readOps);
+  }
+
+  public WritableByteChannel writer() throws IOException {
+    return throughputSink.decorate(Files.newByteChannel(path, writeOps));
+  }
+
+  @Override
+  public void close() throws IOException {
+    Files.delete(path);
+    onCloseCallback.run();
+  }
+
+  @Override
+  public String toString() {
+    return MoreObjects.toStringHelper(this)
+        .add("path", path)
+        .add("throughputSink", throughputSink)
+        .add("onCloseCallback", onCloseCallback)
+        .toString();
+  }
+
+  Unsafe unsafe() {
+    return new Unsafe();
+  }
+
+  final class Unsafe {
+    public Path touch() throws UnsafeIOException {
+      try {
+        return RecoveryFile.this.touch();
+      } catch (IOException e) {
+        throw new UnsafeIOException(e);
+      }
+    }
+  }
+
+  static final class UnsafeIOException extends RuntimeException {
+    private UnsafeIOException(IOException cause) {
+      super(cause);
+    }
+  }
+}