Fix issues with RocksJava dropColumnFamily (#4770)

Summary:
Closes https://github.com/facebook/rocksdb/issues/4409
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4770

Differential Revision: D13416802

Pulled By: ajkr

fbshipit-source-id: 8a351e9b80dc9eeb6073467fbc67cd2f544917b0

java/rocksjni/rocksjni.cc

@@ -1699,14 +1699,46 @@ jlong Java_org_rocksdb_RocksDB_createColumnFamily(JNIEnv* env, jobject /*jdb*/,
 void Java_org_rocksdb_RocksDB_dropColumnFamily(JNIEnv* env, jobject /*jdb*/,
                                                jlong jdb_handle,
                                                jlong jcf_handle) {
-  auto* cf_handle = reinterpret_cast<rocksdb::ColumnFamilyHandle*>(jcf_handle);
   auto* db_handle = reinterpret_cast<rocksdb::DB*>(jdb_handle);
+  auto* cf_handle = reinterpret_cast<rocksdb::ColumnFamilyHandle*>(jcf_handle);
   rocksdb::Status s = db_handle->DropColumnFamily(cf_handle);
   if (!s.ok()) {
     rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
   }
 }
 
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    dropColumnFamilies
+ * Signature: (J[J)V
+ */
+void Java_org_rocksdb_RocksDB_dropColumnFamilies(
+    JNIEnv* env, jobject, jlong jdb_handle, jlongArray jcolumn_family_handles) {
+  auto* db_handle = reinterpret_cast<rocksdb::DB*>(jdb_handle);
+
+  std::vector<rocksdb::ColumnFamilyHandle*> cf_handles;
+  if (jcolumn_family_handles != nullptr) {
+    const jsize len_cols = env->GetArrayLength(jcolumn_family_handles);
+
+    jlong* jcfh = env->GetLongArrayElements(jcolumn_family_handles, nullptr);
+    if (jcfh == nullptr) {
+      // exception thrown: OutOfMemoryError
+      return;
+    }
+
+    for (jsize i = 0; i < len_cols; i++) {
+      auto* cf_handle = reinterpret_cast<rocksdb::ColumnFamilyHandle*>(jcfh[i]);
+      cf_handles.push_back(cf_handle);
+    }
+    env->ReleaseLongArrayElements(jcolumn_family_handles, jcfh, JNI_ABORT);
+  }
+
+  rocksdb::Status s = db_handle->DropColumnFamilies(cf_handles);
+  if (!s.ok()) {
+    rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
+  }
+}
+
 /*
  * Method:    getSnapshot
  * Signature: (J)J

java/src/main/java/org/rocksdb/RocksDB.java

@@ -2343,9 +2343,9 @@ public class RocksDB extends RocksObject {
   }
 
   /**
-   * Drops the column family identified by columnFamilyName. Internal
-   * handles to this column family will be disposed. If the column family
-   * is not known removal will fail.
+   * Drops the column family specified by {@code columnFamilyHandle}. This call
+   * only records a drop record in the manifest and prevents the column
+   * family from flushing and compacting.
    *
    * @param columnFamilyHandle {@link org.rocksdb.ColumnFamilyHandle}
    *     instance
@@ -2354,11 +2354,21 @@ public class RocksDB extends RocksObject {
    *    native library.
    */
   public void dropColumnFamily(final ColumnFamilyHandle columnFamilyHandle)
-      throws RocksDBException, IllegalArgumentException {
-    // throws RocksDBException if something goes wrong
+      throws RocksDBException {
     dropColumnFamily(nativeHandle_, columnFamilyHandle.nativeHandle_);
-    // After the drop the native handle is not valid anymore
-    columnFamilyHandle.disOwnNativeHandle();
+  }
+
+  // Bulk drop column families. This call only records drop records in the
+  // manifest and prevents the column families from flushing and compacting.
+  // In case of error, the request may succeed partially. User may call
+  // ListColumnFamilies to check the result.
+  public void dropColumnFamilies(
+      final List<ColumnFamilyHandle> columnFamilies) throws RocksDBException {
+    final long[] cfHandles = new long[columnFamilies.size()];
+    for (int i = 0; i < columnFamilies.size(); i++) {
+      cfHandles[i] = columnFamilies.get(i).nativeHandle_;
+    }
+    dropColumnFamilies(nativeHandle_, cfHandles);
   }
 
   /**
@@ -3035,8 +3045,10 @@ public class RocksDB extends RocksObject {
   private native long createColumnFamily(final long handle,
       final byte[] columnFamilyName, final long columnFamilyOptions)
       throws RocksDBException;
-  private native void dropColumnFamily(long handle, long cfHandle)
+  private native void dropColumnFamily(final long handle, final long cfHandle)
       throws RocksDBException;
+  private native void dropColumnFamilies(final long handle,
+      final long[] cfHandles) throws RocksDBException;
   private native void flush(long handle, long flushOptHandle)
       throws RocksDBException;
   private native void flush(long handle, long flushOptHandle, long cfHandle)

java/src/test/java/org/rocksdb/ColumnFamilyTest.java

@@ -229,6 +229,7 @@ public class ColumnFamilyTest {
                 new ColumnFamilyOptions()));
         db.put(tmpColumnFamilyHandle, "key".getBytes(), "value".getBytes());
         db.dropColumnFamily(tmpColumnFamilyHandle);
+        assertThat(tmpColumnFamilyHandle.isOwningHandle()).isTrue();
       } finally {
         if (tmpColumnFamilyHandle != null) {
           tmpColumnFamilyHandle.close();
@@ -240,6 +241,46 @@ public class ColumnFamilyTest {
     }
   }
 
+  @Test
+  public void createWriteDropColumnFamilies() throws RocksDBException {
+    final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(
+        new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY),
+        new ColumnFamilyDescriptor("new_cf".getBytes()));
+    final List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>();
+    try (final DBOptions options = new DBOptions()
+        .setCreateIfMissing(true)
+        .setCreateMissingColumnFamilies(true);
+         final RocksDB db = RocksDB.open(options,
+             dbFolder.getRoot().getAbsolutePath(), cfDescriptors,
+             columnFamilyHandleList)) {
+      ColumnFamilyHandle tmpColumnFamilyHandle = null;
+      ColumnFamilyHandle tmpColumnFamilyHandle2 = null;
+      try {
+        tmpColumnFamilyHandle = db.createColumnFamily(
+            new ColumnFamilyDescriptor("tmpCF".getBytes(),
+                new ColumnFamilyOptions()));
+        tmpColumnFamilyHandle2 = db.createColumnFamily(
+            new ColumnFamilyDescriptor("tmpCF2".getBytes(),
+                new ColumnFamilyOptions()));
+        db.put(tmpColumnFamilyHandle, "key".getBytes(), "value".getBytes());
+        db.put(tmpColumnFamilyHandle2, "key".getBytes(), "value".getBytes());
+        db.dropColumnFamilies(Arrays.asList(tmpColumnFamilyHandle, tmpColumnFamilyHandle2));
+        assertThat(tmpColumnFamilyHandle.isOwningHandle()).isTrue();
+        assertThat(tmpColumnFamilyHandle2.isOwningHandle()).isTrue();
+      } finally {
+        if (tmpColumnFamilyHandle != null) {
+          tmpColumnFamilyHandle.close();
+        }
+        if (tmpColumnFamilyHandle2 != null) {
+          tmpColumnFamilyHandle2.close();
+        }
+        for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandleList) {
+          columnFamilyHandle.close();
+        }
+      }
+    }
+  }
+
   @Test
   public void writeBatch() throws RocksDBException {
     try (final StringAppendOperator stringAppendOperator = new StringAppendOperator();