feat: implement BufferToDiskThenUpload BlobWriteSessionConfig (#2139)

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.

Author: BenWhitehead

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,236 @@
+/*
+ * 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.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.nio.ByteBuffer;
+import java.nio.channels.WritableByteChannel;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.time.Clock;
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.stream.Collector;
+import javax.annotation.concurrent.Immutable;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+
+/**
+ * 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 static final long serialVersionUID = 9059242302276891867L;
+
+  /**
+   * non-final because of {@link java.io.Serializable}, however this field is effectively final as
+   * it is immutable and there is not reference mutator method.
+   */
+  @MonotonicNonNull private transient ImmutableList<Path> paths;
+
+  private final boolean includeLoggingSink;
+
+  /** Used for {@link java.io.Serializable} */
+  @MonotonicNonNull private volatile ArrayList<String> absolutePaths;
+
+  @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 void writeObject(ObjectOutputStream out) throws IOException {
+    if (absolutePaths == null) {
+      synchronized (this) {
+        if (absolutePaths == null) {
+          absolutePaths =
+              paths.stream()
+                  .map(Path::toAbsolutePath)
+                  .map(Path::toString)
+                  .collect(
+                      Collector.of(
+                          ArrayList::new,
+                          ArrayList::add,
+                          (left, right) -> {
+                            left.addAll(right);
+                            return left;
+                          }));
+        }
+      }
+    }
+    out.defaultWriteObject();
+  }
+
+  private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+    in.defaultReadObject();
+    this.paths = absolutePaths.stream().map(Paths::get).collect(ImmutableList.toImmutableList());
+  }
+
+  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);