From 49085f8b3efa4641776c507b639aee433c5af847 Mon Sep 17 00:00:00 2001 From: Radek Stankiewicz Date: Thu, 23 Jul 2026 15:38:53 +0200 Subject: [PATCH] Fix ThreadSanitizer data races between bundle processing and async background writer closing in FileBasedSink and WriteFiles --- .../org/apache/beam/sdk/io/FileBasedSink.java | 20 ++++++++++++++----- .../org/apache/beam/sdk/io/WriteFiles.java | 3 +++ .../TFRecordSchemaTransformProviderTest.java | 7 ++++--- 3 files changed, 22 insertions(+), 8 deletions(-) diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java index bba9b1f82f5b..d134df6fe272 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/FileBasedSink.java @@ -40,6 +40,7 @@ import java.util.Objects; import java.util.Set; import java.util.UUID; +import java.util.concurrent.atomic.AtomicBoolean; import org.apache.beam.sdk.coders.CannotProvideCoderException; import org.apache.beam.sdk.coders.Coder; import org.apache.beam.sdk.coders.CoderException; @@ -913,7 +914,7 @@ public abstract static class Writer { private @Nullable ResourceId outputFile; /** The channel to write to. */ - private @Nullable WritableByteChannel channel; + private volatile @Nullable WritableByteChannel channel; /** * The MIME type used in the creation of the output channel (if the file system supports it). @@ -925,6 +926,8 @@ public abstract static class Writer { */ private final @Nullable String mimeType; + private final AtomicBoolean readyToClose = new AtomicBoolean(false); + /** Construct a new {@link Writer} that will produce files of the given MIME type. */ public Writer(WriteOperation writeOperation, String mimeType) { checkNotNull(writeOperation); @@ -932,6 +935,10 @@ public Writer(WriteOperation writeOperation, String mimeT this.mimeType = mimeType; } + void releaseForBackgroundClose() { + readyToClose.set(true); + } + /** * Called with the channel that a subclass will write its header, footer, and values to. * Subclasses should either keep a reference to the channel provided or create and keep a @@ -1037,6 +1044,7 @@ private static void closeChannelAndThrow( } public final void cleanup() throws Exception { + readyToClose.get(); if (outputFile != null) { LOG.info("Deleting temporary file {}", outputFile); // outputFile may be null if open() was not called or failed. @@ -1047,28 +1055,30 @@ public final void cleanup() throws Exception { /** Closes the channel and returns the bundle result. */ public final void close() throws Exception { + readyToClose.get(); + WritableByteChannel channelToClose = channel; checkState(outputFile != null, "FileResult.close cannot be called with a null outputFile"); LOG.debug("Closing {}", outputFile); try { writeFooter(); } catch (Exception e) { - closeChannelAndThrow(channel, outputFile, e); + closeChannelAndThrow(channelToClose, outputFile, e); } try { finishWrite(); } catch (Exception e) { - closeChannelAndThrow(channel, outputFile, e); + closeChannelAndThrow(channelToClose, outputFile, e); } // It is valid for a subclass to either close the channel or not. // They would typically close the channel e.g. if they are wrapping it in another channel // and the wrapper needs to be closed. - if (channel.isOpen()) { + if (channelToClose.isOpen()) { LOG.debug("Closing channel to {}.", outputFile); try { - channel.close(); + channelToClose.close(); } catch (Exception e) { throw new IOException(String.format("Failed closing channel to %s", outputFile), e); } diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java index b0b5051f3210..a55b5658a1f5 100644 --- a/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java +++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java @@ -1272,6 +1272,9 @@ public void processElement( } private void closeWriterInBackground(Writer writer) { + // Release memory barrier so background thread acquires all prior writes on Writer and + // channels. + writer.releaseForBackgroundClose(); // Close in parallel so flushing of buffered writes to files for many windows happens in // parallel. closeFutures.add( diff --git a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java index 9c067a533e0b..18de29698ab6 100644 --- a/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java +++ b/sdks/java/core/src/test/java/org/apache/beam/sdk/io/TFRecordSchemaTransformProviderTest.java @@ -434,9 +434,10 @@ private void runTestWrite(String[] elems, String... base64) throws IOException { assertTrue("File should exist", tmpFile.exists()); assertTrue("File should have content", tmpFile.length() > 0); - FileInputStream fis = new FileInputStream(tmpFile); - String written = BaseEncoding.base64().encode(ByteStreams.toByteArray(fis)); - assertThat(written, is(in(base64))); + try (FileInputStream fis = new FileInputStream(tmpFile)) { + String written = BaseEncoding.base64().encode(ByteStreams.toByteArray(fis)); + assertThat(written, is(in(base64))); + } } @Test