From 9e05c5e2519444faaf147a79e842853cee6f20ff Mon Sep 17 00:00:00 2001
From: Tomas Kolda <koldat@gmail.com>
Date: Mon, 14 Mar 2022 14:12:30 -0700
Subject: [PATCH] NPE in
 Java_org_rocksdb_ColumnFamilyOptions_setSstPartitionerFactory (#9622)

Summary:
There was a mistake that incorrectly cast SstPartitionerFactory (missed shared pointer). It worked for database (correct cast), but not for family. Trying to set it in family has caused Access violation.

I have also added test and improved it. Older version was passing even without sst partitioner which is weird, because on Level1 we had two SST files with same key "aaaa1". I was not sure if it is a new feature and changed it to overlaping keys "aaaa0" - "aaaa2" overlaps "aaaa1".

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9622

Reviewed By: ajkr

Differential Revision: D34871968

Pulled By: pdillinger

fbshipit-source-id: a08009766da49fc198692a610e8beb19caf737e6
---
 java/rocksjni/options.cc                      |  5 +--
 .../java/org/rocksdb/SstPartitionerTest.java  | 33 +++++++++++++++++--
 2 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/java/rocksjni/options.cc b/java/rocksjni/options.cc
index 8eae762b3..b892bf056 100644
--- a/java/rocksjni/options.cc
+++ b/java/rocksjni/options.cc
@@ -4261,9 +4261,10 @@ void Java_org_rocksdb_ColumnFamilyOptions_setSstPartitionerFactory(
     JNIEnv*, jobject, jlong jhandle, jlong factory_handle) {
   auto* options =
       reinterpret_cast<ROCKSDB_NAMESPACE::ColumnFamilyOptions*>(jhandle);
-  auto* factory = reinterpret_cast<ROCKSDB_NAMESPACE::SstPartitionerFactory*>(
+  auto factory = reinterpret_cast<
+      std::shared_ptr<ROCKSDB_NAMESPACE::SstPartitionerFactory>*>(
       factory_handle);
-  options->sst_partitioner_factory.reset(factory);
+  options->sst_partitioner_factory = *factory;
 }
 
 /*
diff --git a/java/src/test/java/org/rocksdb/SstPartitionerTest.java b/java/src/test/java/org/rocksdb/SstPartitionerTest.java
index 8593e0045..74816db93 100644
--- a/java/src/test/java/org/rocksdb/SstPartitionerTest.java
+++ b/java/src/test/java/org/rocksdb/SstPartitionerTest.java
@@ -5,6 +5,7 @@
 
 package org.rocksdb;
 
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import java.util.List;
@@ -21,7 +22,7 @@ public class SstPartitionerTest {
   @Rule public TemporaryFolder dbFolder = new TemporaryFolder();
 
   @Test
-  public void sstFixedPrefix() throws InterruptedException, RocksDBException {
+  public void sstFixedPrefix() throws RocksDBException {
     try (SstPartitionerFixedPrefixFactory factory = new SstPartitionerFixedPrefixFactory(4);
          final Options opt =
              new Options().setCreateIfMissing(true).setSstPartitionerFactory(factory);
@@ -31,7 +32,8 @@ public class SstPartitionerTest {
       db.put("bbbb1".getBytes(), "B".getBytes());
       db.flush(new FlushOptions());
 
-      db.put("aaaa1".getBytes(), "A2".getBytes());
+      db.put("aaaa0".getBytes(), "A2".getBytes());
+      db.put("aaaa2".getBytes(), "A2".getBytes());
       db.flush(new FlushOptions());
 
       db.compactRange();
@@ -40,4 +42,31 @@ public class SstPartitionerTest {
       assertThat(metadata.size()).isEqualTo(2);
     }
   }
+
+  @Test
+  public void sstFixedPrefixFamily() throws RocksDBException {
+    final byte[] cfName = "new_cf".getBytes(UTF_8);
+    final ColumnFamilyDescriptor cfDescriptor = new ColumnFamilyDescriptor(cfName,
+        new ColumnFamilyOptions().setSstPartitionerFactory(
+            new SstPartitionerFixedPrefixFactory(4)));
+
+    try (final Options opt = new Options().setCreateIfMissing(true);
+         final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
+      final ColumnFamilyHandle columnFamilyHandle = db.createColumnFamily(cfDescriptor);
+
+      // writing (long)100 under key
+      db.put(columnFamilyHandle, "aaaa1".getBytes(), "A".getBytes());
+      db.put(columnFamilyHandle, "bbbb1".getBytes(), "B".getBytes());
+      db.flush(new FlushOptions(), columnFamilyHandle);
+
+      db.put(columnFamilyHandle, "aaaa0".getBytes(), "A2".getBytes());
+      db.put(columnFamilyHandle, "aaaa2".getBytes(), "A2".getBytes());
+      db.flush(new FlushOptions(), columnFamilyHandle);
+
+      db.compactRange(columnFamilyHandle);
+
+      List<LiveFileMetaData> metadata = db.getLiveFilesMetaData();
+      assertThat(metadata.size()).isEqualTo(2);
+    }
+  }
 }