Tidy up and share (between processes) temporary copies of JNI shared library loaded from JAR

java/CMakeLists.txt

@@ -299,6 +299,7 @@ set(JAVA_MAIN_CLASSES
   src/main/java/org/rocksdb/util/Environment.java
   src/main/java/org/rocksdb/util/IntComparator.java
   src/main/java/org/rocksdb/util/ReverseBytewiseComparator.java
+  src/main/java/org/rocksdb/util/SharedTempFile.java
   src/main/java/org/rocksdb/util/SizeUnit.java
   src/main/java/org/rocksdb/util/StdErrLogger.java
   src/test/java/org/rocksdb/NativeComparatorWrapperTest.java
@@ -384,6 +385,9 @@ set(JAVA_TEST_CLASSES
   src/test/java/org/rocksdb/RocksIteratorTest.java
   src/test/java/org/rocksdb/RocksMemEnvTest.java
   src/test/java/org/rocksdb/SecondaryDBTest.java
+  src/test/java/org/rocksdb/SharedTempFileLoaderTest.java
+  src/test/java/org/rocksdb/SharedTempFileMockMain.java
+  src/test/java/org/rocksdb/SharedTempFileRocksDBMain.java
   src/test/java/org/rocksdb/SliceTest.java
   src/test/java/org/rocksdb/SnapshotTest.java
   src/test/java/org/rocksdb/SstFileManagerTest.java
@@ -415,6 +419,7 @@ set(JAVA_TEST_CLASSES
   src/test/java/org/rocksdb/test/RocksJunitRunner.java
   src/test/java/org/rocksdb/test/TestableEventListener.java
   src/test/java/org/rocksdb/test/TestableEventListener.java
+  src/main/java/org/rocksdb/util/ArgUtil.java
   src/test/java/org/rocksdb/util/ByteBufferAllocator.java
   src/test/java/org/rocksdb/util/BytewiseComparatorIntTest.java
   src/test/java/org/rocksdb/util/BytewiseComparatorTest.java
@@ -429,6 +434,10 @@ set(JAVA_TEST_CLASSES
   src/test/java/org/rocksdb/util/TestUtil.java
 )
 
+set(JAVA_TEST_RESOURCES
+  src/test/resources/shared-temp-file.txt
+)
+
 set(JAVA_TEST_RUNNING_CLASSES
   org.rocksdb.BackupEngineOptionsTest
   org.rocksdb.BackupEngineTest
@@ -503,6 +512,7 @@ set(JAVA_TEST_RUNNING_CLASSES
   org.rocksdb.RocksIteratorTest
   org.rocksdb.RocksMemEnvTest
   org.rocksdb.SecondaryDBTest
+  org.rocksdb.SharedTempFileLoaderTest
   org.rocksdb.SliceTest
   org.rocksdb.SnapshotTest
   org.rocksdb.SstFileManagerTest
@@ -598,6 +608,7 @@ add_jar(
     SOURCES
     ${JAVA_MAIN_CLASSES}
     ${JAVA_TEST_CLASSES}
+    ${JAVA_TEST_RESOURCES}
     INCLUDE_JARS ${JAVA_TESTCLASSPATH}
     GENERATE_NATIVE_HEADERS rocksdbjni_test_headers DESTINATION ${JNI_OUTPUT_DIR}
 )

java/Makefile

@@ -184,6 +184,7 @@ JAVA_TESTS = \
 	org.rocksdb.RocksMemEnvTest\
 	org.rocksdb.util.SizeUnitTest\
 	org.rocksdb.SecondaryDBTest\
+	org.rocksdb.SharedTempFileLoaderTest\
 	org.rocksdb.SliceTest\
 	org.rocksdb.SnapshotTest\
 	org.rocksdb.SstFileManagerTest\
@@ -212,6 +213,7 @@ JAVA_TESTS = \
 
 MAIN_SRC = src/main/java
 TEST_SRC = src/test/java
+TEST_RESOURCES = src/test/resources
 OUTPUT = target
 MAIN_CLASSES = $(OUTPUT)/classes
 TEST_CLASSES = $(OUTPUT)/test-classes
@@ -452,6 +454,7 @@ java_test: java resolve_test_deps
 	$(AM_V_GEN)mkdir -p $(TEST_CLASSES)
 	$(AM_V_at) $(JAVAC_CMD) $(JAVAC_ARGS) -cp $(MAIN_CLASSES):$(JAVA_TESTCLASSPATH) -h $(NATIVE_INCLUDE) -d $(TEST_CLASSES)\
 		$(TEST_SOURCES)
+	cp -v $(TEST_RESOURCES)/* $(TEST_CLASSES)
 
 test: java java_test
 	$(MAKE) run_test

java/src/main/java/org/rocksdb/NativeLibraryLoader.java

@@ -1,11 +1,12 @@
 // Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
 package org.rocksdb;
 
-import java.io.*;
-import java.nio.file.Files;
-import java.nio.file.StandardCopyOption;
-
+import java.io.IOException;
+import java.io.InputStream;
+import java.net.URL;
+import java.nio.file.Path;
 import org.rocksdb.util.Environment;
+import org.rocksdb.util.SharedTempFile;
 
 /**
  * This class is used to load the RocksDB shared library from within the jar.
@@ -104,63 +105,52 @@ public synchronized void loadLibrary(final String tmpDir) throws IOException {
   void loadLibraryFromJar(final String tmpDir)
       throws IOException {
     if (!initialized) {
-      System.load(loadLibraryFromJarToTemp(tmpDir).getAbsolutePath());
+      System.load(loadLibraryFromJarToTemp(tmpDir).toString());
       initialized = true;
     }
   }
 
-  private File createTemp(final String tmpDir, final String libraryFileName) throws IOException {
-    // create a temporary file to copy the library to
-    final File temp;
-    if (tmpDir == null || tmpDir.isEmpty()) {
-      temp = File.createTempFile(tempFilePrefix, tempFileSuffix);
-    } else {
-      final File parentDir = new File(tmpDir);
-      if (!parentDir.exists()) {
-        throw new RuntimeException(
-            "Directory: " + parentDir.getAbsolutePath() + " does not exist!");
-      }
-      temp = new File(parentDir, libraryFileName);
-      if (temp.exists() && !temp.delete()) {
-        throw new RuntimeException(
-            "File: " + temp.getAbsolutePath() + " already exists and cannot be removed.");
-      }
-      if (!temp.createNewFile()) {
-        throw new RuntimeException("File: " + temp.getAbsolutePath() + " could not be created.");
-      }
+  @SuppressWarnings({"PMD.UseProperClassLoader", "PMD.CloseResource"})
+  private InputStream libraryFromJar() {
+    InputStream is = getClass().getClassLoader().getResourceAsStream(jniLibraryFileName);
+    if (is == null) {
+      is = getClass().getClassLoader().getResourceAsStream(fallbackJniLibraryFileName);
     }
-    if (temp.exists()) {
-      temp.deleteOnExit();
-      return temp;
-    } else {
-      throw new RuntimeException("File " + temp.getAbsolutePath() + " does not exist.");
+    if (is == null) {
+      throw new RuntimeException(jniLibraryFileName + " was not found inside JAR.");
     }
+    return is;
   }
 
-  @SuppressWarnings({"PMD.UseProperClassLoader", "PMD.UseTryWithResources"})
-  File loadLibraryFromJarToTemp(final String tmpDir) throws IOException {
-    try (InputStream is = getClass().getClassLoader().getResourceAsStream(jniLibraryFileName)) {
-      if (is != null) {
-        final File temp = createTemp(tmpDir, jniLibraryFileName);
-        Files.copy(is, temp.toPath(), StandardCopyOption.REPLACE_EXISTING);
-        return temp;
-      }
+  @SuppressWarnings("PMD.UseProperClassLoader")
+  private String libraryResourcePath() {
+    URL resource = getClass().getClassLoader().getResource(jniLibraryFileName);
+    if (resource == null) {
+      resource = getClass().getClassLoader().getResource(fallbackJniLibraryFileName);
     }
-
-    if (fallbackJniLibraryFileName == null) {
-      throw new RuntimeException(fallbackJniLibraryFileName + " was not found inside JAR.");
+    if (resource == null) {
+      throw new RuntimeException(jniLibraryFileName + " was not found inside JAR.");
     }
+    return resource.getFile();
+  }
 
-    try (InputStream is =
-             getClass().getClassLoader().getResourceAsStream(fallbackJniLibraryFileName)) {
-      if (is != null) {
-        final File temp = createTemp(tmpDir, fallbackJniLibraryFileName);
-        Files.copy(is, temp.toPath(), StandardCopyOption.REPLACE_EXISTING);
-        return temp;
+  @SuppressWarnings({"PMD.UseTryWithResources", "PMD.CloseResource"})
+  Path loadLibraryFromJarToTemp(final String tmpDir) throws IOException {
+    String prefix = tempFilePrefix;
+    String suffix = tempFileSuffix;
+    if (jniLibraryFileName != null) {
+      String[] split = jniLibraryFileName.split("\\.");
+      if (split.length == 2) {
+        prefix = split[0];
+        suffix = split[1];
       }
     }
-
-    throw new RuntimeException(jniLibraryFileName + " was not found inside JAR.");
+    SharedTempFile.Instance instance =
+        new SharedTempFile.Instance(tmpDir, prefix, libraryResourcePath(), suffix);
+    SharedTempFile sharedTemp = instance.create();
+    SharedTempFile.Lock lock = sharedTemp.lock(this::libraryFromJar);
+    Runtime.getRuntime().addShutdownHook(new Thread(lock::close));
+    return sharedTemp.getContent();
   }
 
   /**