apacheGH-36069: [Java] Ensure S3 is finalized on shutdown

danepitkin · Jul 28, 2023 · 6058031 · 6058031
1 parent 8c4941b
commit 6058031
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 0 deletions.
diff --git a/java/dataset/src/main/cpp/jni_wrapper.cc b/java/dataset/src/main/cpp/jni_wrapper.cc
@@ -27,6 +27,7 @@
 #include "arrow/dataset/file_base.h"
 #include "arrow/filesystem/localfs.h"
 #include "arrow/filesystem/path_util.h"
+#include "arrow/filesystem/s3fs.h"
 #include "arrow/engine/substrait/util.h"
 #include "arrow/ipc/api.h"
 #include "arrow/util/iterator.h"
@@ -678,6 +679,18 @@ Java_org_apache_arrow_dataset_file_JniWrapper_writeFromScannerToFile(
   JNI_METHOD_END()
 }
 
+/*
+ * Class:     org_apache_arrow_dataset_file_JniWrapper
+ * Method:    ensureS3Finalized
+ * Signature: (J)V
+ */
+JNIEXPORT void JNICALL Java_org_apache_arrow_dataset_file_JniWrapper_ensureS3Finalized(
+    JNIEnv* env, jobject) {
+  JNI_METHOD_START
+  JniAssertOkOrThrow(arrow::fs::EnsureS3Finalized());
+  JNI_METHOD_END()
+}
+
 /*
  * Class:     org_apache_arrow_dataset_substrait_JniWrapper
  * Method:    executeSerializedPlan

diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/file/FileSystemDatasetFactory.java b/java/dataset/src/main/java/org/apache/arrow/dataset/file/FileSystemDatasetFactory.java
@@ -17,6 +17,8 @@
 
 package org.apache.arrow.dataset.file;
 
+import java.util.concurrent.atomic.AtomicBoolean;
+
 import org.apache.arrow.dataset.jni.NativeDatasetFactory;
 import org.apache.arrow.dataset.jni.NativeMemoryPool;
 import org.apache.arrow.memory.BufferAllocator;
@@ -26,14 +28,18 @@
  */
 public class FileSystemDatasetFactory extends NativeDatasetFactory {
 
+  private static final AtomicBoolean addedS3ShutdownHook = new AtomicBoolean(false);
+
   public FileSystemDatasetFactory(BufferAllocator allocator, NativeMemoryPool memoryPool, FileFormat format,
       String uri) {
     super(allocator, memoryPool, createNative(format, uri));
+    ensureS3FinalizedOnShutdown();
   }
 
   public FileSystemDatasetFactory(BufferAllocator allocator, NativeMemoryPool memoryPool, FileFormat format,
                                   String[] uris) {
     super(allocator, memoryPool, createNative(format, uris));
+    ensureS3FinalizedOnShutdown();
   }
 
   private static long createNative(FileFormat format, String uri) {
@@ -44,4 +50,10 @@ private static long createNative(FileFormat format, String[] uris) {
     return JniWrapper.get().makeFileSystemDatasetFactory(uris, format.id());
   }
 
+  private static void ensureS3FinalizedOnShutdown() {
+    if (addedS3ShutdownHook.compareAndSet(false, true)) {
+      Runtime.getRuntime().addShutdownHook(new Thread(() -> { JniWrapper.get().ensureS3Finalized(); }));
+    }
+  }
+
 }
diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/file/JniWrapper.java b/java/dataset/src/main/java/org/apache/arrow/dataset/file/JniWrapper.java
@@ -75,4 +75,10 @@ public native void writeFromScannerToFile(long streamAddress,
                                             int maxPartitions,
                                             String baseNameTemplate);
 
+  /**
+   * Ensure the S3 APIs are shutdown, but only if not already done. If the S3 APIs are unintialized,
+   * then this is a noop.
+   */
+  public native void ensureS3Finalized();
+
 }