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

Author: stankiewicz

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

@@ -17,6 +17,7 @@
  */
 package org.apache.beam.sdk.io;
 
+import java.util.concurrent.atomic.AtomicBoolean;
 import static org.apache.beam.sdk.io.WriteFiles.UNKNOWN_SHARDNUM;
 import static org.apache.beam.sdk.values.TypeDescriptors.extractFromTypeParameters;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects.firstNonNull;
@@ -925,13 +926,20 @@ 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 +1045,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,6 +1056,7 @@ public final void cleanup() throws Exception {
 
     /** Closes the channel and returns the bundle result. */
     public final void close() throws Exception {
+      readyToClose.get();
       checkState(outputFile != null, "FileResult.close cannot be called with a null outputFile");
       LOG.debug("Closing {}", outputFile);
 

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