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 16acb90d11a0..b5246a87d2ba 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 GoogleCloudStorageReadOptions create(PipelineOptions options) { /** * Maximum number of files rewritten concurrently when copying or moving files. If unset, a - * default of 32 is used. + * default of 128 is used. * *

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 863bc9030e94..c0f40a9aec4b 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 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 @@ Storage storage() { /** * Default maximum number of files rewritten concurrently by {@link #copy} and {@link #move}, used * when {@link GcsOptions#getGcsMaxConcurrentRewrites} is unset. + * + *

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 @@ Storage storage() { */ 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 @@ Storage storage() { 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 @@ public SeekableByteChannel open(GcsPath path, BlobSourceOption... sourceOptions) } } - /** A bridge that allows a GCS WriteChannel to behave as a WritableByteChannel. */ + /** + * A bridge that allows a GCS WriteChannel to behave as a WritableByteChannel. + * + *

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 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 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 @@ public int write(ByteBuffer src) throws IOException { @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 @@ public WritableByteChannel create( 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 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 0654f9744440..9335ab6f56e0 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 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.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 void testV2WriteChannelCloseTranslatesStorageException() throws IOExcepti 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 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 @@ private static Storage storageWritingTo(WriteChannel writer) { @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 void testV2CreatePreconditions() throws IOException { 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 void testV2UploadChunkSize() throws IOException { 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 void testV2WriteMetricsIsSet() { 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);