Add SST partitioner interface (#512)

Add interface for `SstPartitioner`. The interface is to support implementing "compaction guard" to split SST by custom boundaries.
Signed-off-by: Yi Wu <yiwu@pingcap.com>

Cargo.toml

@@ -31,5 +31,6 @@ librocksdb_sys = { path = "librocksdb_sys" }
 
 [dev-dependencies]
 crc = "1.8"
+lazy_static = "1.4.0"
 rand = "0.7"
 tempfile = "3.1"

Makefile

@@ -2,7 +2,7 @@ all: format build test
 
 format:
 	@cargo fmt --all
-	@librocksdb_sys/crocksdb/format-diff.sh > /dev/null || true
+	@scripts/format-diff.sh
 
 build:
 	@cargo build

librocksdb_sys/crocksdb/crocksdb/c.h

@@ -154,6 +154,19 @@ typedef struct crocksdb_writestallcondition_t crocksdb_writestallcondition_t;
 typedef struct crocksdb_map_property_t crocksdb_map_property_t;
 typedef struct crocksdb_writebatch_iterator_t crocksdb_writebatch_iterator_t;
 
+typedef enum crocksdb_sst_partitioner_result_t {
+  kNotRequired = 0,
+  kRequired = 1,
+} crocksdb_sst_partitioner_result_t;
+
+typedef struct crocksdb_sst_partitioner_t crocksdb_sst_partitioner_t;
+typedef struct crocksdb_sst_partitioner_request_t
+    crocksdb_sst_partitioner_request_t;
+typedef struct crocksdb_sst_partitioner_context_t
+    crocksdb_sst_partitioner_context_t;
+typedef struct crocksdb_sst_partitioner_factory_t
+    crocksdb_sst_partitioner_factory_t;
+
 typedef enum crocksdb_table_property_t {
   kDataSize = 1,
   kIndexSize = 2,
@@ -1005,6 +1018,10 @@ crocksdb_options_get_max_bytes_for_level_multiplier(crocksdb_options_t*);
 extern C_ROCKSDB_LIBRARY_API void
 crocksdb_options_set_max_bytes_for_level_multiplier_additional(
     crocksdb_options_t*, int* level_values, size_t num_levels);
+extern C_ROCKSDB_LIBRARY_API crocksdb_sst_partitioner_factory_t*
+crocksdb_options_get_sst_partitioner_factory(crocksdb_options_t*);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_options_set_sst_partitioner_factory(
+    crocksdb_options_t*, crocksdb_sst_partitioner_factory_t*);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_enable_statistics(
     crocksdb_options_t*, unsigned char);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_options_reset_statistics(
@@ -2225,6 +2242,115 @@ crocksdb_iostats_context_prepare_write_nanos(crocksdb_iostats_context_t*);
 extern C_ROCKSDB_LIBRARY_API uint64_t
 crocksdb_iostats_context_logger_nanos(crocksdb_iostats_context_t*);
 
+/* SstPartitioner */
+
+extern C_ROCKSDB_LIBRARY_API crocksdb_sst_partitioner_request_t*
+crocksdb_sst_partitioner_request_create();
+extern C_ROCKSDB_LIBRARY_API void crocksdb_sst_partitioner_request_destroy(
+    crocksdb_sst_partitioner_request_t* req);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_sst_partitioner_request_prev_user_key(
+    crocksdb_sst_partitioner_request_t* req, size_t* len);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_sst_partitioner_request_current_user_key(
+    crocksdb_sst_partitioner_request_t* req, size_t* len);
+extern C_ROCKSDB_LIBRARY_API uint64_t
+crocksdb_sst_partitioner_request_current_output_file_size(
+    crocksdb_sst_partitioner_request_t* req);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_req_set_prev_user_key(
+    crocksdb_sst_partitioner_request_t* req, const char* key, size_t len);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_req_set_current_user_key(
+    crocksdb_sst_partitioner_request_t* req, const char* key, size_t len);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_request_set_current_output_file_size(
+    crocksdb_sst_partitioner_request_t* req, uint64_t current_output_file_size);
+
+typedef crocksdb_sst_partitioner_result_t (
+    *crocksdb_sst_partitioner_should_partition_cb)(
+    void* underlying, crocksdb_sst_partitioner_request_t* req);
+typedef unsigned char (*crocksdb_sst_partitioner_can_do_trivial_move_cb)(
+    void* underlying, const char* smallest_user_key,
+    size_t smallest_user_key_len, const char* largest_user_key,
+    size_t largest_user_key_len);
+
+extern C_ROCKSDB_LIBRARY_API crocksdb_sst_partitioner_t*
+crocksdb_sst_partitioner_create(
+    void* underlying, void (*destructor)(void*),
+    crocksdb_sst_partitioner_should_partition_cb should_partition_cb,
+    crocksdb_sst_partitioner_can_do_trivial_move_cb can_do_trivial_move_cb);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_sst_partitioner_destroy(
+    crocksdb_sst_partitioner_t* partitioner);
+extern C_ROCKSDB_LIBRARY_API crocksdb_sst_partitioner_result_t
+crocksdb_sst_partitioner_should_partition(
+    crocksdb_sst_partitioner_t* partitioner,
+    crocksdb_sst_partitioner_request_t* req);
+extern C_ROCKSDB_LIBRARY_API unsigned char
+crocksdb_sst_partitioner_can_do_trivial_move(
+    crocksdb_sst_partitioner_t* partitioner, const char* smallest_user_key,
+    size_t smallest_user_key_len, const char* largest_user_key,
+    size_t largest_user_key_len);
+
+extern C_ROCKSDB_LIBRARY_API crocksdb_sst_partitioner_context_t*
+crocksdb_sst_partitioner_context_create();
+extern C_ROCKSDB_LIBRARY_API void crocksdb_sst_partitioner_context_destroy(
+    crocksdb_sst_partitioner_context_t* context);
+extern C_ROCKSDB_LIBRARY_API unsigned char
+crocksdb_sst_partitioner_context_is_full_compaction(
+    crocksdb_sst_partitioner_context_t* context);
+extern C_ROCKSDB_LIBRARY_API unsigned char
+crocksdb_sst_partitioner_context_is_manual_compaction(
+    crocksdb_sst_partitioner_context_t* context);
+extern C_ROCKSDB_LIBRARY_API int crocksdb_sst_partitioner_context_output_level(
+    crocksdb_sst_partitioner_context_t* context);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_sst_partitioner_context_smallest_key(
+    crocksdb_sst_partitioner_context_t* context, size_t* key_len);
+extern C_ROCKSDB_LIBRARY_API const char*
+crocksdb_sst_partitioner_context_largest_key(
+    crocksdb_sst_partitioner_context_t* context, size_t* key_len);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_context_set_is_full_compaction(
+    crocksdb_sst_partitioner_context_t* context,
+    unsigned char is_full_compaction);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_context_set_is_manual_compaction(
+    crocksdb_sst_partitioner_context_t* context,
+    unsigned char is_manual_compaction);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_context_set_output_level(
+    crocksdb_sst_partitioner_context_t* context, int output_level);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_context_set_smallest_key(
+    crocksdb_sst_partitioner_context_t* context, const char* smallest_key,
+    size_t key_len);
+extern C_ROCKSDB_LIBRARY_API void
+crocksdb_sst_partitioner_context_set_largest_key(
+    crocksdb_sst_partitioner_context_t* context, const char* largest_key,
+    size_t key_len);
+
+typedef const char* (*crocksdb_sst_partitioner_factory_name_cb)(
+    void* underlying);
+typedef crocksdb_sst_partitioner_t* (
+    *crocksdb_sst_partitioner_factory_create_partitioner_cb)(
+    void* underlying, crocksdb_sst_partitioner_context_t* context);
+
+extern C_ROCKSDB_LIBRARY_API crocksdb_sst_partitioner_factory_t*
+crocksdb_sst_partitioner_factory_create(
+    void* underlying, void (*destructor)(void*),
+    crocksdb_sst_partitioner_factory_name_cb name_cb,
+    crocksdb_sst_partitioner_factory_create_partitioner_cb
+        create_partitioner_cb);
+extern C_ROCKSDB_LIBRARY_API void crocksdb_sst_partitioner_factory_destroy(
+    crocksdb_sst_partitioner_factory_t* factory);
+extern C_ROCKSDB_LIBRARY_API const char* crocksdb_sst_partitioner_factory_name(
+    crocksdb_sst_partitioner_factory_t* factory);
+extern C_ROCKSDB_LIBRARY_API crocksdb_sst_partitioner_t*
+crocksdb_sst_partitioner_factory_create_partitioner(
+    crocksdb_sst_partitioner_factory_t* factory,
+    crocksdb_sst_partitioner_context_t* context);
+
 extern C_ROCKSDB_LIBRARY_API void crocksdb_run_ldb_tool(
     int argc, char** argv, const crocksdb_options_t* opts);
 extern C_ROCKSDB_LIBRARY_API void crocksdb_run_sst_dump_tool(

librocksdb_sys/src/lib.rs

@@ -164,6 +164,14 @@ pub struct DBFileEncryptionInfo(c_void);
 #[repr(C)]
 pub struct DBEncryptionKeyManagerInstance(c_void);
 #[repr(C)]
+pub struct DBSstPartitioner(c_void);
+#[repr(C)]
+pub struct DBSstPartitionerRequest(c_void);
+#[repr(C)]
+pub struct DBSstPartitionerContext(c_void);
+#[repr(C)]
+pub struct DBSstPartitionerFactory(c_void);
+#[repr(C)]
 pub struct DBWriteBatchIterator(c_void);
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -431,6 +439,13 @@ impl fmt::Display for DBEncryptionMethod {
     }
 }
 
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+#[repr(C)]
+pub enum DBSstPartitionerResult {
+    NotRequired = 0,
+    Required = 1,
+}
+
 /// # Safety
 ///
 /// ptr must point to a valid CStr value
@@ -825,6 +840,13 @@ extern "C" {
     pub fn crocksdb_options_get_path_target_size(options: *mut Options, idx: size_t) -> u64;
     pub fn crocksdb_options_set_vector_memtable_factory(options: *mut Options, reserved_bytes: u64);
     pub fn crocksdb_options_set_atomic_flush(option: *mut Options, enable: bool);
+    pub fn crocksdb_options_get_sst_partitioner_factory(
+        option: *mut Options,
+    ) -> *mut DBSstPartitionerFactory;
+    pub fn crocksdb_options_set_sst_partitioner_factory(
+        option: *mut Options,
+        factory: *mut DBSstPartitionerFactory,
+    );
     pub fn crocksdb_filterpolicy_create_bloom_full(bits_per_key: c_int) -> *mut DBFilterPolicy;
     pub fn crocksdb_filterpolicy_create_bloom(bits_per_key: c_int) -> *mut DBFilterPolicy;
     pub fn crocksdb_open(
@@ -2265,6 +2287,122 @@ extern "C" {
     pub fn crocksdb_iostats_context_prepare_write_nanos(ctx: *mut DBIOStatsContext) -> u64;
     pub fn crocksdb_iostats_context_logger_nanos(ctx: *mut DBIOStatsContext) -> u64;
 
+    pub fn crocksdb_sst_partitioner_request_create() -> *mut DBSstPartitionerRequest;
+    pub fn crocksdb_sst_partitioner_request_destroy(state: *mut DBSstPartitionerRequest);
+    pub fn crocksdb_sst_partitioner_request_prev_user_key(
+        state: *mut DBSstPartitionerRequest,
+        len: *mut size_t,
+    ) -> *const c_char;
+    pub fn crocksdb_sst_partitioner_request_current_user_key(
+        state: *mut DBSstPartitionerRequest,
+        len: *mut size_t,
+    ) -> *const c_char;
+    pub fn crocksdb_sst_partitioner_request_current_output_file_size(
+        state: *mut DBSstPartitionerRequest,
+    ) -> u64;
+    pub fn crocksdb_sst_partitioner_request_set_prev_user_key(
+        state: *mut DBSstPartitionerRequest,
+        key: *const c_char,
+        len: size_t,
+    );
+    pub fn crocksdb_sst_partitioner_request_set_current_user_key(
+        state: *mut DBSstPartitionerRequest,
+        key: *const c_char,
+        len: size_t,
+    );
+    pub fn crocksdb_sst_partitioner_request_set_current_output_file_size(
+        state: *mut DBSstPartitionerRequest,
+        current_output_file_size: u64,
+    );
+
+    pub fn crocksdb_sst_partitioner_create(
+        underlying: *mut c_void,
+        destructor: extern "C" fn(*mut c_void),
+        should_partition_cb: extern "C" fn(
+            *mut c_void,
+            *mut DBSstPartitionerRequest,
+        ) -> DBSstPartitionerResult,
+        can_do_trivial_move_cb: extern "C" fn(
+            *mut c_void,
+            *const c_char,
+            size_t,
+            *const c_char,
+            size_t,
+        ) -> c_uchar,
+    ) -> *mut DBSstPartitioner;
+    pub fn crocksdb_sst_partitioner_destroy(partitioner: *mut DBSstPartitioner);
+    pub fn crocksdb_sst_partitioner_should_partition(
+        partitioner: *mut DBSstPartitioner,
+        req: *mut DBSstPartitionerRequest,
+    ) -> DBSstPartitionerResult;
+    pub fn crocksdb_sst_partitioner_can_do_trivial_move(
+        partitioner: *mut DBSstPartitioner,
+        smallest_key: *const c_char,
+        smallest_key_len: size_t,
+        largest_key: *const c_char,
+        largest_key_len: size_t,
+    ) -> bool;
+
+    pub fn crocksdb_sst_partitioner_context_create() -> *mut DBSstPartitionerContext;
+    pub fn crocksdb_sst_partitioner_context_destroy(context: *mut DBSstPartitionerContext);
+    pub fn crocksdb_sst_partitioner_context_is_full_compaction(
+        context: *mut DBSstPartitionerContext,
+    ) -> c_uchar;
+    pub fn crocksdb_sst_partitioner_context_is_manual_compaction(
+        context: *mut DBSstPartitionerContext,
+    ) -> c_uchar;
+    pub fn crocksdb_sst_partitioner_context_output_level(
+        context: *mut DBSstPartitionerContext,
+    ) -> c_int;
+    pub fn crocksdb_sst_partitioner_context_smallest_key(
+        context: *mut DBSstPartitionerContext,
+        key_len: *mut size_t,
+    ) -> *const c_char;
+    pub fn crocksdb_sst_partitioner_context_largest_key(
+        context: *mut DBSstPartitionerContext,
+        key_len: *mut size_t,
+    ) -> *const c_char;
+    pub fn crocksdb_sst_partitioner_context_set_is_full_compaction(
+        context: *mut DBSstPartitionerContext,
+        is_full_compaction: c_uchar,
+    );
+    pub fn crocksdb_sst_partitioner_context_set_is_manual_compaction(
+        context: *mut DBSstPartitionerContext,
+        is_manual_compaction: c_uchar,
+    );
+    pub fn crocksdb_sst_partitioner_context_set_output_level(
+        context: *mut DBSstPartitionerContext,
+        output_level: c_int,
+    );
+    pub fn crocksdb_sst_partitioner_context_set_smallest_key(
+        context: *mut DBSstPartitionerContext,
+        smallest_key: *const c_char,
+        key_len: size_t,
+    );
+    pub fn crocksdb_sst_partitioner_context_set_largest_key(
+        context: *mut DBSstPartitionerContext,
+        largest_key: *const c_char,
+        key_len: size_t,
+    );
+
+    pub fn crocksdb_sst_partitioner_factory_create(
+        underlying: *mut c_void,
+        destructor: extern "C" fn(*mut c_void),
+        name_cb: extern "C" fn(*mut c_void) -> *const c_char,
+        create_partitioner_cb: extern "C" fn(
+            *mut c_void,
+            *mut DBSstPartitionerContext,
+        ) -> *mut DBSstPartitioner,
+    ) -> *mut DBSstPartitionerFactory;
+    pub fn crocksdb_sst_partitioner_factory_destroy(factory: *mut DBSstPartitionerFactory);
+    pub fn crocksdb_sst_partitioner_factory_name(
+        factory: *mut DBSstPartitionerFactory,
+    ) -> *const c_char;
+    pub fn crocksdb_sst_partitioner_factory_create_partitioner(
+        factory: *mut DBSstPartitionerFactory,
+        context: *mut DBSstPartitionerContext,
+    ) -> *mut DBSstPartitioner;
+
     pub fn crocksdb_run_ldb_tool(argc: c_int, argv: *const *const c_char, opts: *const Options);
     pub fn crocksdb_run_sst_dump_tool(
         argc: c_int,

rust-toolchain

-nightly-2019-12-20
+nightly-2020-07-01

src/lib.rs

@@ -23,6 +23,9 @@ extern crate libc;
 pub extern crate librocksdb_sys;
 #[cfg(test)]
 extern crate tempfile;
+#[cfg(test)]
+#[macro_use]
+extern crate lazy_static;
 
 #[cfg(feature = "cloud")]
 pub use cloud::CloudEnvOptions;
@@ -39,7 +42,8 @@ pub use event_listener::{
 pub use librocksdb_sys::{
     self as crocksdb_ffi, new_bloom_filter, CompactionPriority, CompactionReason,
     DBBackgroundErrorReason, DBBottommostLevelCompaction, DBCompactionStyle, DBCompressionType,
-    DBEntryType, DBInfoLogLevel, DBRateLimiterMode, DBRecoveryMode, DBStatisticsHistogramType,
+    DBEntryType, DBInfoLogLevel, DBRateLimiterMode, DBRecoveryMode,
+    DBSstPartitionerResult as SstPartitionerResult, DBStatisticsHistogramType,
     DBStatisticsTickerType, DBStatusPtr, DBTitanDBBlobRunMode, IndexType, WriteStallCondition,
 };
 pub use logger::Logger;
@@ -58,6 +62,9 @@ pub use rocksdb_options::{
     WriteOptions,
 };
 pub use slice_transform::SliceTransform;
+pub use sst_partitioner::{
+    SstPartitioner, SstPartitionerContext, SstPartitionerFactory, SstPartitionerRequest,
+};
 pub use table_filter::TableFilter;
 pub use table_properties::{
     TableProperties, TablePropertiesCollection, TablePropertiesCollectionView,
@@ -85,6 +92,7 @@ mod perf_context;
 pub mod rocksdb;
 pub mod rocksdb_options;
 mod slice_transform;
+pub mod sst_partitioner;
 mod table_filter;
 mod table_properties;
 mod table_properties_collector;

src/rocksdb_options.rs

@@ -34,6 +34,7 @@ use merge_operator::{self, full_merge_callback, partial_merge_callback, MergeOpe
 use rocksdb::Env;
 use rocksdb::{Cache, MemoryAllocator};
 use slice_transform::{new_slice_transform, SliceTransform};
+use sst_partitioner::{new_sst_partitioner_factory, SstPartitionerFactory};
 use std::ffi::{CStr, CString};
 use std::path::Path;
 use std::ptr;
@@ -1740,6 +1741,13 @@ impl ColumnFamilyOptions {
             Some(CStr::from_ptr(memtable_name).to_str().unwrap())
         }
     }
+
+    pub fn set_sst_partitioner_factory<F: SstPartitionerFactory>(&mut self, factory: F) {
+        let f = new_sst_partitioner_factory(factory);
+        unsafe {
+            crocksdb_ffi::crocksdb_options_set_sst_partitioner_factory(self.inner, f);
+        }
+    }
 }
 
 // ColumnFamilyDescriptor is a pair of column family's name and options.