This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 8cae4991fca [Java] Close GcsUtilV2's small-file performance gaps with
GcsUtilV1 (#40677)
8cae4991fca is described below
commit 8cae4991fca407818454ae5bc5ac1e5f70c03411
Author: Shunping Huang <[email protected]>
AuthorDate: Wed Oct 21 11:43:35 2026 -0400
[Java] Close GcsUtilV2's small-file performance gaps with GcsUtilV1 (#40677)
* Use the shared HTTP transport in GcsUtilV2
* Start GcsUtilV2 upload sessions in the background
* Raise the default GcsUtilV2 rewrite concurrency to 128
* Handle closed and failed upload sessions in GcsWritableByteChannel
---
.../sdk/extensions/gcp/options/GcsOptions.java | 4 +-
.../beam/sdk/extensions/gcp/util/GcsUtilV2.java | 125 +++++++++++--
.../sdk/extensions/gcp/util/GcsUtilV2Test.java | 198 +++++++++++++++++++--
3 files changed, 300 insertions(+), 27 deletions(-)
diff --git
a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java
b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java
index 16acb90d11a..b5246a87d2b 100644
---
a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java
+++
b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/options/GcsOptions.java
@@ -193,13 +193,13 @@ public interface GcsOptions extends
ApplicationNameOptions, GcpOptions, Pipeline
/**
* Maximum number of files rewritten concurrently when copying or moving
files. If unset, a
- * default of 32 is used.
+ * default of 128 is used.
*
* <p>Only applicable to GcsUtil V2, which is enabled with the {@code
use_gcsutil_v2} experiment.
*/
@Description(
"Maximum number of files rewritten concurrently when copying or moving
files in GCS."
- + " If unset, defaults to 32. Only applicable to GcsUtil V2 (enabled
with the"
+ + " If unset, defaults to 128. Only applicable to GcsUtil V2
(enabled with the"
+ " use_gcsutil_v2 experiment).")
@Nullable Integer getGcsMaxConcurrentRewrites();
diff --git
a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java
b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java
index 863bc9030e9..c0f40a9aec4 100644
---
a/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java
+++
b/sdks/java/extensions/google-cloud-platform-core/src/main/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2.java
@@ -56,9 +56,11 @@ import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
import java.io.Closeable;
import java.io.FileNotFoundException;
import java.io.IOException;
+import java.io.InterruptedIOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.ByteBuffer;
+import java.nio.channels.ClosedChannelException;
import java.nio.channels.SeekableByteChannel;
import java.nio.channels.WritableByteChannel;
import java.nio.file.AccessDeniedException;
@@ -159,8 +161,12 @@ class GcsUtilV2 {
/**
* Default maximum number of files rewritten concurrently by {@link #copy}
and {@link #move}, used
* when {@link GcsOptions#getGcsMaxConcurrentRewrites} is unset.
+ *
+ * <p>A copy or move never uses more threads than it has files, so this only
caps large ones, such
+ * as the temp-to-final renames that finish a write of thousands of shards.
There 128 brought the
+ * rename time close to {@code GcsUtilV1}'s batched rewrites without any
request being throttled.
*/
- private static final int DEFAULT_MAX_CONCURRENT_REWRITES = 32;
+ private static final int DEFAULT_MAX_CONCURRENT_REWRITES = 128;
/**
* How long {@link #copy} and {@link #move} wait for in-flight rewrites to
finish after a failure
@@ -168,6 +174,18 @@ class GcsUtilV2 {
*/
private static final long REWRITE_TERMINATION_TIMEOUT_MINUTES = 5;
+ /**
+ * Starts the resumable upload sessions of {@link #create} in the
background. Unbounded, so that a
+ * burst of creates is never queued behind one another, and daemon, so that
an abandoned upload
+ * cannot keep the JVM alive.
+ */
+ private static final ExecutorService UPLOAD_SESSION_EXECUTOR =
+ Executors.newCachedThreadPool(
+ new ThreadFactoryBuilder()
+ .setDaemon(true)
+ .setNameFormat("gcsutil-v2-upload-session-%d")
+ .build());
+
/** Maximum number of files rewritten concurrently by {@link #copy} and
{@link #move}. */
private final int maxConcurrentRewrites;
@@ -192,6 +210,13 @@ class GcsUtilV2 {
storageOptionsBuilder.setHost(rootUrlOf(endpoint));
}
+ // Use the process-wide shared HTTP transport rather than letting each
client create its own, so
+ // a new client reuses established connections instead of opening new ones
on its first request.
+ storageOptionsBuilder.setTransportOptions(
+ StorageOptions.getDefaultHttpTransportOptions().toBuilder()
+ .setHttpTransportFactory(Transport::getTransport)
+ .build());
+
storage = storageOptionsBuilder.build().getService();
uploadBufferSizeBytes = gcsOptions.getGcsUploadBufferSizeBytes();
this.gcsCountersOptions =
@@ -994,20 +1019,67 @@ class GcsUtilV2 {
}
}
- /** A bridge that allows a GCS WriteChannel to behave as a
WritableByteChannel. */
+ /**
+ * A bridge that allows a GCS WriteChannel to behave as a
WritableByteChannel.
+ *
+ * <p>The WriteChannel is opened in the background, because java-storage
starts the resumable
+ * upload session with a blocking request when it opens one. Waiting for it
on the first write, or
+ * on close, rather than in {@link #create} lets that request overlap with
the caller's own work,
+ * as gcsio does for {@code GcsUtilV1}. A failure to start the session
surfaces there instead.
+ */
private static class GcsWritableByteChannel implements WritableByteChannel {
- private final WriteChannel writer;
+ private final Future<WriteChannel> pendingWriter;
private final GcsPath gcsPath;
+ private boolean closed = false;
+ // Set when the upload session could not be started (it failed, or waiting
for it was
+ // interrupted). There is then nothing to finalize, and the failure has
already been reported.
+ private boolean sessionFailed = false;
- GcsWritableByteChannel(WriteChannel writer, GcsPath gcsPath) {
- this.writer = writer;
+ GcsWritableByteChannel(Future<WriteChannel> pendingWriter, GcsPath
gcsPath) {
+ this.pendingWriter = pendingWriter;
this.gcsPath = gcsPath;
}
+ /** Waits for the upload session to start and returns its channel. */
+ private WriteChannel writer() throws IOException {
+ try {
+ return pendingWriter.get();
+ } catch (InterruptedException e) {
+ sessionFailed = true;
+ Thread.currentThread().interrupt();
+ pendingWriter.cancel(true);
+ throw new InterruptedIOException("Interrupted while starting the
upload to " + gcsPath);
+ } catch (CancellationException e) {
+ // Only cancelled above, by an earlier interrupted wait, which
abandoned the upload.
+ sessionFailed = true;
+ throw new InterruptedIOException(
+ "The upload to " + gcsPath + " was abandoned after an earlier
interrupt");
+ } catch (ExecutionException e) {
+ sessionFailed = true;
+ Throwable cause = e.getCause();
+ if (cause instanceof StorageException) {
+ throw translateStorageException(gcsPath, (StorageException) cause);
+ }
+ if (cause instanceof IOException) {
+ throw (IOException) cause;
+ }
+ if (cause instanceof RuntimeException) {
+ throw (RuntimeException) cause;
+ }
+ if (cause instanceof Error) {
+ throw (Error) cause;
+ }
+ throw new IOException("Failed to start the upload to " + gcsPath,
cause);
+ }
+ }
+
@Override
public int write(ByteBuffer src) throws IOException {
+ if (closed) {
+ throw new ClosedChannelException();
+ }
try {
- return writer.write(src);
+ return writer().write(src);
} catch (StorageException e) {
throw translateStorageException(gcsPath, e);
}
@@ -1015,14 +1087,23 @@ class GcsUtilV2 {
@Override
public boolean isOpen() {
- return writer.isOpen();
+ return !closed;
}
@Override
public void close() throws IOException {
+ if (closed) {
+ return;
+ }
+ closed = true;
+ if (sessionFailed) {
+ // An earlier write already reported that the session could not be
started, and there is
+ // nothing to finalize.
+ return;
+ }
// The upload is finalized here, so this is where a failed precondition
surfaces.
try {
- writer.close();
+ writer().close();
} catch (StorageException e) {
throw translateStorageException(gcsPath, e);
}
@@ -1058,19 +1139,33 @@ class GcsUtilV2 {
writeOptionList.add(BlobWriteOption.generationMatch(blob.getGeneration()));
}
}
- // Open a WriteChannel from the storage service
- WriteChannel writer =
- client.writer(blobInfo, writeOptionList.toArray(new
BlobWriteOption[0]));
+ // Open a WriteChannel from the storage service. Opening one starts the
resumable upload
+ // session with a blocking request, so it is done in the background and
the returned channel
+ // waits for it on its first write or on close (see
GcsWritableByteChannel).
+ BlobWriteOption[] writeOptionArray = writeOptionList.toArray(new
BlobWriteOption[0]);
Integer uploadBufferSizeBytes =
options.getUploadBufferSizeBytes() != null
? options.getUploadBufferSizeBytes()
: this.uploadBufferSizeBytes;
- writer.setChunkSize(
- uploadBufferSizeBytes != null ? uploadBufferSizeBytes :
DEFAULT_UPLOAD_CHUNK_SIZE_BYTES);
+ int chunkSize =
+ uploadBufferSizeBytes != null ? uploadBufferSizeBytes :
DEFAULT_UPLOAD_CHUNK_SIZE_BYTES;
+ Future<WriteChannel> pendingWriter =
+ UPLOAD_SESSION_EXECUTOR.submit(
+ () -> {
+ try {
+ WriteChannel writer = client.writer(blobInfo,
writeOptionArray);
+ writer.setChunkSize(chunkSize);
+ serviceCallMetric.call("ok");
+ return writer;
+ } catch (StorageException e) {
+ serviceCallMetric.call(e.getCode());
+ throw e;
+ }
+ });
- serviceCallMetric.call("ok");
// Return the bridge wrapper
- return wrapInCounting(new GcsWritableByteChannel(writer, path),
path.getBucket(), container);
+ return wrapInCounting(
+ new GcsWritableByteChannel(pendingWriter, path), path.getBucket(),
container);
} catch (StorageException e) {
serviceCallMetric.call(e.getCode());
diff --git
a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java
b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java
index 0654f974444..9335ab6f56e 100644
---
a/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java
+++
b/sdks/java/extensions/google-cloud-platform-core/src/test/java/org/apache/beam/sdk/extensions/gcp/util/GcsUtilV2Test.java
@@ -58,9 +58,11 @@ import com.google.cloud.storage.StorageOptions;
import java.io.ByteArrayOutputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
+import java.io.InterruptedIOException;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
+import java.nio.channels.ClosedChannelException;
import java.nio.channels.SeekableByteChannel;
import java.nio.channels.WritableByteChannel;
import java.nio.file.AccessDeniedException;
@@ -71,6 +73,7 @@ import java.util.List;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
import
org.apache.beam.repackaged.core.org.apache.commons.compress.utils.SeekableInMemoryByteChannel;
import org.apache.beam.runners.core.metrics.CounterCell;
import org.apache.beam.runners.core.metrics.GcpResourceIdentifiers;
@@ -560,6 +563,175 @@ public class GcsUtilV2Test {
assertSame(preconditionFailed, thrown.getCause());
}
+ /**
+ * java-storage starts the resumable upload session with a blocking request
when it opens a
+ * WriteChannel. create() must leave that to a background thread, as gcsio
does for V1, so that
+ * the caller does not wait for it.
+ */
+ @Test
+ public void testV2CreateStartsUploadSessionInBackground() throws Exception {
+ WriteChannel writer = Mockito.mock(WriteChannel.class);
+ CountDownLatch release = new CountDownLatch(1);
+ AtomicReference<Thread> openingThread = new AtomicReference<>();
+ Storage storage = Mockito.mock(Storage.class);
+ when(storage.writer(any(BlobInfo.class), any()))
+ .thenAnswer(
+ invocation -> {
+ openingThread.set(Thread.currentThread());
+ if (!release.await(10, TimeUnit.SECONDS)) {
+ throw new AssertionError("the upload session was never
released");
+ }
+ return writer;
+ });
+ GcsUtil gcsUtil = gcsUtilWithV2Storage(storage);
+
+ WritableByteChannel channel =
+ gcsUtil.create(
+ GcsPath.fromComponents("testbucket", "testobject"),
+ CreateOptions.builder().setExpectFileToNotExist(true).build());
+ assertTrue(channel.isOpen());
+
+ release.countDown();
+ channel.write(ByteBuffer.wrap(new byte[] {1, 2, 3}));
+ channel.close();
+
+ assertNotSame(Thread.currentThread(), openingThread.get());
+ verify(writer).write(any(ByteBuffer.class));
+ verify(writer).close();
+ assertFalse(channel.isOpen());
+ }
+
+ /**
+ * A failure to start the upload session surfaces on the first write, as an
IOException. Close is
+ * then a no-op, and later writes report the closed channel.
+ */
+ @Test
+ public void testV2CreateReportsUploadSessionFailureOnWrite() throws
IOException {
+ Storage storage = Mockito.mock(Storage.class);
+ when(storage.writer(any(BlobInfo.class), any()))
+ .thenThrow(new StorageException(403, "Forbidden"));
+ GcsUtil gcsUtil = gcsUtilWithV2Storage(storage);
+
+ WritableByteChannel channel =
+ gcsUtil.create(
+ GcsPath.fromComponents("testbucket", "testobject"),
+ CreateOptions.builder().setExpectFileToNotExist(true).build());
+
+ assertThrows(AccessDeniedException.class, () ->
channel.write(ByteBuffer.wrap(new byte[] {1})));
+ // The failure was already reported and there is no session to finalize,
so close is a no-op.
+ channel.close();
+ assertFalse(channel.isOpen());
+ assertThrows(
+ ClosedChannelException.class, () -> channel.write(ByteBuffer.wrap(new
byte[] {1})));
+ }
+
+ /** A java-storage client whose upload sessions do not start until {@code
release} fires. */
+ private static Storage storageWithBlockedUploadSession(
+ CountDownLatch release, WriteChannel writer) {
+ Storage storage = Mockito.mock(Storage.class);
+ when(storage.writer(any(BlobInfo.class), any()))
+ .thenAnswer(
+ invocation -> {
+ release.await();
+ return writer;
+ });
+ return storage;
+ }
+
+ private static WritableByteChannel createTestObject(GcsUtil gcsUtil) throws
IOException {
+ return gcsUtil.create(
+ GcsPath.fromComponents("testbucket", "testobject"),
+ CreateOptions.builder().setExpectFileToNotExist(true).build());
+ }
+
+ /** A closed channel rejects writes itself, rather than relying on the
underlying channel. */
+ @Test
+ public void testV2WriteAfterCloseThrowsClosedChannelException() throws
IOException {
+ WriteChannel writer = Mockito.mock(WriteChannel.class);
+ WritableByteChannel channel =
createTestObject(gcsUtilWithV2Storage(storageWritingTo(writer)));
+
+ channel.close();
+
+ assertThrows(
+ ClosedChannelException.class, () -> channel.write(ByteBuffer.wrap(new
byte[] {1})));
+ verify(writer, never()).write(any(ByteBuffer.class));
+ }
+
+ /**
+ * When the session fails to start and nothing was written, close reports
the failure. A later
+ * write must then report the closed channel, not the session failure again.
+ */
+ @Test
+ public void testV2WriteAfterFailedCloseThrowsClosedChannelException() throws
IOException {
+ Storage storage = Mockito.mock(Storage.class);
+ when(storage.writer(any(BlobInfo.class), any()))
+ .thenThrow(new StorageException(403, "Forbidden"));
+ WritableByteChannel channel =
createTestObject(gcsUtilWithV2Storage(storage));
+
+ assertThrows(AccessDeniedException.class, channel::close);
+
+ assertFalse(channel.isOpen());
+ assertThrows(
+ ClosedChannelException.class, () -> channel.write(ByteBuffer.wrap(new
byte[] {1})));
+ }
+
+ /**
+ * An interrupt while the first write waits for the session abandons the
upload. Later writes
+ * report that as an IOException rather than leaking a
CancellationException, and close is a
+ * no-op.
+ */
+ @Test
+ public void testV2WriteInterruptedWhileStartingSession() throws IOException {
+ CountDownLatch release = new CountDownLatch(1);
+ WriteChannel writer = Mockito.mock(WriteChannel.class);
+ WritableByteChannel channel =
+
createTestObject(gcsUtilWithV2Storage(storageWithBlockedUploadSession(release,
writer)));
+ try {
+ Thread.currentThread().interrupt();
+ assertThrows(
+ InterruptedIOException.class, () ->
channel.write(ByteBuffer.wrap(new byte[] {1})));
+ // The interrupt is preserved for the caller.
+ assertTrue(Thread.interrupted());
+
+ assertThrows(
+ InterruptedIOException.class, () ->
channel.write(ByteBuffer.wrap(new byte[] {1})));
+ channel.close();
+ assertFalse(channel.isOpen());
+ verify(writer, never()).write(any(ByteBuffer.class));
+ verify(writer, never()).close();
+ } finally {
+ Thread.interrupted();
+ release.countDown();
+ }
+ }
+
+ /**
+ * An interrupt while close waits for the session of an empty file is
reported by close. A later
+ * write must then report the closed channel.
+ */
+ @Test
+ public void testV2CloseInterruptedWhileStartingSession() throws IOException {
+ CountDownLatch release = new CountDownLatch(1);
+ WriteChannel writer = Mockito.mock(WriteChannel.class);
+ WritableByteChannel channel =
+
createTestObject(gcsUtilWithV2Storage(storageWithBlockedUploadSession(release,
writer)));
+ try {
+ Thread.currentThread().interrupt();
+ assertThrows(InterruptedIOException.class, channel::close);
+ assertTrue(Thread.interrupted());
+
+ assertFalse(channel.isOpen());
+ assertThrows(
+ ClosedChannelException.class, () ->
channel.write(ByteBuffer.wrap(new byte[] {1})));
+ // A second close is a no-op.
+ channel.close();
+ verify(writer, never()).close();
+ } finally {
+ Thread.interrupted();
+ release.countDown();
+ }
+ }
+
/** Mirrors {@link GcsUtilV1Test#testBucketDoesNotExist} for V2. */
@Test
public void testV2BucketAccessibleIsFalseWhenBucketDoesNotExist() throws
IOException {
@@ -724,18 +896,21 @@ public class GcsUtilV2Test {
@Test
public void testV2CreatePreconditions() throws IOException {
GcsPath path = GcsPath.fromComponents("testbucket", "testobject");
+ // Each channel is closed so that its upload session, started in the
background, has been
+ // opened before it is verified.
// Expected not to exist: no lookup is made.
Storage expectedMissing =
storageWritingTo(Mockito.mock(WriteChannel.class));
gcsUtilWithV2Storage(expectedMissing)
- .create(path,
CreateOptions.builder().setExpectFileToNotExist(true).build());
+ .create(path,
CreateOptions.builder().setExpectFileToNotExist(true).build())
+ .close();
verify(expectedMissing).writer(any(BlobInfo.class),
Mockito.eq(BlobWriteOption.doesNotExist()));
verify(expectedMissing, never()).get(anyString(), anyString(), any());
// Looked up and missing. An unstubbed get() returns null, which is how
java-storage reports a
// missing object.
Storage missing = storageWritingTo(Mockito.mock(WriteChannel.class));
- gcsUtilWithV2Storage(missing).create(path,
CreateOptions.builder().build());
+ gcsUtilWithV2Storage(missing).create(path,
CreateOptions.builder().build()).close();
verify(missing).writer(any(BlobInfo.class),
Mockito.eq(BlobWriteOption.doesNotExist()));
// Looked up and found.
@@ -743,7 +918,7 @@ public class GcsUtilV2Test {
Blob blob = Mockito.mock(Blob.class);
when(blob.getGeneration()).thenReturn(42L);
when(existing.get(Mockito.eq("testbucket"), Mockito.eq("testobject"),
any())).thenReturn(blob);
- gcsUtilWithV2Storage(existing).create(path,
CreateOptions.builder().build());
+ gcsUtilWithV2Storage(existing).create(path,
CreateOptions.builder().build()).close();
verify(existing).writer(any(BlobInfo.class),
Mockito.eq(BlobWriteOption.generationMatch(42L)));
}
@@ -765,21 +940,23 @@ public class GcsUtilV2Test {
pipelineWithSize.setGcsUploadBufferSizeBytes(2048);
WriteChannel writer = Mockito.mock(WriteChannel.class);
- gcsUtilWithV2Storage(pipelineWithSize,
storageWritingTo(writer)).create(path, withSize);
+ gcsUtilWithV2Storage(pipelineWithSize,
storageWritingTo(writer)).create(path, withSize).close();
verify(writer).setChunkSize(1024);
writer = Mockito.mock(WriteChannel.class);
- gcsUtilWithV2Storage(pipelineWithSize,
storageWritingTo(writer)).create(path, withoutSize);
+ gcsUtilWithV2Storage(pipelineWithSize, storageWritingTo(writer))
+ .create(path, withoutSize)
+ .close();
verify(writer).setChunkSize(2048);
writer = Mockito.mock(WriteChannel.class);
- gcsUtilWithV2Storage(storageWritingTo(writer)).create(path, withoutSize);
+ gcsUtilWithV2Storage(storageWritingTo(writer)).create(path,
withoutSize).close();
verify(writer).setChunkSize(GcsUtilV2.DEFAULT_UPLOAD_CHUNK_SIZE_BYTES);
}
/** Mirrors {@link GcsUtilV1Test#testGCSWriteMetricsIsSet}. */
@Test
- public void testV2WriteMetricsIsSet() {
+ public void testV2WriteMetricsIsSet() throws IOException {
MetricsContainerImpl container = new MetricsContainerImpl(null);
MetricsEnvironment.setProcessWideContainer(container);
Storage storage = Mockito.mock(Storage.class);
@@ -788,9 +965,10 @@ public class GcsUtilV2Test {
GcsUtil gcsUtil = gcsUtilWithV2Storage(storage);
GcsPath path = GcsPath.fromComponents("testbucket", "testobject");
- assertThrows(
- AccessDeniedException.class,
- () -> gcsUtil.create(path,
CreateOptions.builder().setExpectFileToNotExist(true).build()));
+ // The upload session is started in the background, so its failure
surfaces on close.
+ WritableByteChannel channel =
+ gcsUtil.create(path,
CreateOptions.builder().setExpectFileToNotExist(true).build());
+ assertThrows(AccessDeniedException.class, channel::close);
verifyMetricWasSet("my_project", "testbucket", "GcsInsert",
"permission_denied", 1);
verifyMetricWasSet("my_project", "testbucket", "GcsInsert", "ok", 0);
