Fix ThreadSanitizer data races between bundle processing and async background writer closing in FileBasedSink and WriteFiles

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<DestinationT, OutputT> {
     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,13 +926,19 @@ public abstract static class Writer<DestinationT, OutputT> {
      */
     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<DestinationT, OutputT> writeOperation, String mimeType) {
       checkNotNull(writeOperation);
       this.writeOperation = writeOperation;
       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);
         }

sdks/java/core/src/main/java/org/apache/beam/sdk/io/WriteFiles.java

@@ -1272,6 +1272,9 @@ public void processElement(
     }
 
     private void closeWriterInBackground(Writer<DestinationT, OutputT> 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(

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