oxrocksdb-sys table properties and user collected properties example

table_properties_demo.rs

//! Demo and integration coverage for the oxrocksdb-sys table-properties
//! collector, collection iterator, and event-listener APIs.
//!
//! Run with:
//! `cargo test -p oxrocksdb-sys --test table_properties_demo -- --nocapture`

#![allow(
    unsafe_code,
    clippy::expect_used,
    clippy::print_stdout,
    clippy::print_stderr
)]

use oxrocksdb_sys::*;
use std::ffi::{CStr, CString, c_char, c_void};
use std::ptr;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use tempfile::TempDir;

const COLLECTOR_NAME: &CStr = c"OxDemoCollector";
const FACTORY_NAME: &CStr = c"OxDemoCollectorFactory";
const PROP_NUM_PUTS: &[u8] = b"ox.num_puts";
const PROP_NUM_BYTES: &[u8] = b"ox.num_bytes";
const PROP_MARKER: &[u8] = b"ox.collector";
const PROP_MARKER_VALUE: &[u8] = b"demo";

struct CollectorState {
    puts: AtomicU64,
    bytes: AtomicU64,
}

struct FactoryState {
    collectors_created: AtomicU64,
}

struct ListenerState {
    flush_begins: AtomicU64,
    flushes: AtomicU64,
    compaction_begins: AtomicU64,
    compactions: AtomicU64,
    messages: Mutex<Vec<String>>,
}

fn cstring(s: &str) -> CString {
    CString::new(s).expect("string must not contain interior nuls")
}

unsafe fn take_error(err: *mut c_char) -> Result<(), String> {
    if err.is_null() {
        Ok(())
    } else {
        // SAFETY: RocksDB allocated `err` with malloc and documented it as a
        // NUL-terminated C string that the caller must free with rocksdb_free.
        let msg = unsafe { CStr::from_ptr(err) }
            .to_string_lossy()
            .into_owned();
        unsafe { rocksdb_free(err.cast()) };
        Err(msg)
    }
}

unsafe fn slice_from_raw<'a>(ptr: *const c_char, len: usize) -> &'a [u8] {
    if ptr.is_null() || len == 0 {
        &[]
    } else {
        // SAFETY: caller guarantees `ptr` is valid for `len` bytes.
        unsafe { std::slice::from_raw_parts(ptr.cast::<u8>(), len) }
    }
}

unsafe extern "C" fn collector_name(_state: *mut c_void) -> *const c_char {
    COLLECTOR_NAME.as_ptr()
}

unsafe extern "C" fn collector_destruct(state: *mut c_void) {
    if !state.is_null() {
        // SAFETY: `state` is a `Box<CollectorState>` from `Box::into_raw`.
        drop(unsafe { Box::from_raw(state.cast::<CollectorState>()) });
    }
}

unsafe extern "C" fn collector_add(
    state: *mut c_void,
    _key: *const c_char,
    _key_len: usize,
    value: *const c_char,
    value_len: usize,
    _entry_type: u32,
    _seq: u64,
    _file_size: u64,
) {
    // SAFETY: RocksDB passes the collector state we boxed in `create_collector`.
    let state = unsafe { &*state.cast::<CollectorState>() };
    state.puts.fetch_add(1, Ordering::Relaxed);
    if !value.is_null() {
        state
            .bytes
            .fetch_add(u64::try_from(value_len).unwrap_or(u64::MAX), Ordering::Relaxed);
    }
}

unsafe extern "C" fn collector_finish(
    state: *mut c_void,
    props: *mut oxrocksdb_user_collected_properties_t,
) {
    // SAFETY: Finish is called once with the collector state and a
    // RocksDB-owned properties map that we must only mutate, not free.
    let state = unsafe { &*state.cast::<CollectorState>() };
    let puts = state.puts.load(Ordering::Relaxed).to_string();
    let bytes = state.bytes.load(Ordering::Relaxed).to_string();
    unsafe {
        oxrocksdb_user_collected_properties_add(
            props,
            PROP_NUM_PUTS.as_ptr().cast(),
            PROP_NUM_PUTS.len(),
            puts.as_ptr().cast(),
            puts.len(),
        );
        oxrocksdb_user_collected_properties_add(
            props,
            PROP_NUM_BYTES.as_ptr().cast(),
            PROP_NUM_BYTES.len(),
            bytes.as_ptr().cast(),
            bytes.len(),
        );
        oxrocksdb_user_collected_properties_add(
            props,
            PROP_MARKER.as_ptr().cast(),
            PROP_MARKER.len(),
            PROP_MARKER_VALUE.as_ptr().cast(),
            PROP_MARKER_VALUE.len(),
        );
    }
}

unsafe extern "C" fn factory_name(_state: *mut c_void) -> *const c_char {
    FACTORY_NAME.as_ptr()
}

unsafe extern "C" fn factory_destruct(state: *mut c_void) {
    if !state.is_null() {
        // SAFETY: `state` is a `Box<FactoryState>` from `Box::into_raw`.
        drop(unsafe { Box::from_raw(state.cast::<FactoryState>()) });
    }
}

unsafe extern "C" fn create_collector(
    state: *mut c_void,
    _cf: u32,
) -> *mut oxrocksdb_table_properties_collector_t {
    // SAFETY: factory state lives as long as RocksDB holds the factory.
    let factory = unsafe { &*state.cast::<FactoryState>() };
    factory.collectors_created.fetch_add(1, Ordering::Relaxed);

    let collector_state = Box::into_raw(Box::new(CollectorState {
        puts: AtomicU64::new(0),
        bytes: AtomicU64::new(0),
    }));
    // RocksDB takes ownership of this pointer and deletes it. Do not call
    // oxrocksdb_table_properties_collector_destroy on the returned value.
    unsafe {
        oxrocksdb_table_properties_collector_create(
            collector_state.cast(),
            Some(collector_name),
            Some(collector_destruct),
            Some(collector_add),
            Some(collector_finish),
        )
    }
}

unsafe extern "C" fn listener_destruct(state: *mut c_void) {
    if !state.is_null() {
        // SAFETY: `state` is an `Arc<ListenerState>` reconstructed from a raw pointer.
        drop(unsafe { Arc::from_raw(state.cast::<ListenerState>()) });
    }
}

unsafe extern "C" fn on_flush_begin(
    state: *mut c_void,
    _db: *mut rocksdb_t,
    info: *const rocksdb_flushjobinfo_t,
) {
    let listener = unsafe { &*state.cast::<ListenerState>() };
    listener.flush_begins.fetch_add(1, Ordering::SeqCst);

    let mut cf_len = 0;
    let cf = unsafe { rocksdb_flushjobinfo_cf_name(info, &raw mut cf_len) };
    let cf_name = String::from_utf8_lossy(unsafe { slice_from_raw(cf, cf_len) }).into_owned();
    let message = format!("flush begin on CF '{cf_name}'");
    eprintln!("{message}");
    listener
        .messages
        .lock()
        .expect("listener mutex")
        .push(message);
}

unsafe extern "C" fn on_flush_completed(
    state: *mut c_void,
    _db: *mut rocksdb_t,
    info: *const rocksdb_flushjobinfo_t,
) {
    // SAFETY: listener state is an Arc kept alive until listener_destruct.
    let listener = unsafe { &*state.cast::<ListenerState>() };
    listener.flushes.fetch_add(1, Ordering::SeqCst);

    let mut cf_len = 0;
    // SAFETY: `info` is valid for the duration of this callback only.
    let cf = unsafe { rocksdb_flushjobinfo_cf_name(info, &raw mut cf_len) };
    let cf_name = String::from_utf8_lossy(unsafe { slice_from_raw(cf, cf_len) }).into_owned();

    let table_props = unsafe { oxrocksdb_flushjobinfo_table_properties(info) };
    let user_props = unsafe { oxrocksdb_table_properties_get_user_properties(table_props) };
    let marker = lookup_user_prop(user_props, PROP_MARKER);
    let message = format!(
        "flush completed on CF '{cf_name}', user marker={marker:?}"
    );
    eprintln!("{message}");
    listener
        .messages
        .lock()
        .expect("listener mutex")
        .push(message);
}

unsafe extern "C" fn on_compaction_begin(
    state: *mut c_void,
    _db: *mut rocksdb_t,
    info: *const rocksdb_compactionjobinfo_t,
) {
    let listener = unsafe { &*state.cast::<ListenerState>() };
    listener.compaction_begins.fetch_add(1, Ordering::SeqCst);

    let mut cf_len = 0;
    let cf = unsafe { rocksdb_compactionjobinfo_cf_name(info, &raw mut cf_len) };
    let cf_name = String::from_utf8_lossy(unsafe { slice_from_raw(cf, cf_len) }).into_owned();
    let message = format!("compaction begin on CF '{cf_name}'");
    eprintln!("{message}");
    listener
        .messages
        .lock()
        .expect("listener mutex")
        .push(message);
}

unsafe extern "C" fn on_compaction_completed(
    state: *mut c_void,
    _db: *mut rocksdb_t,
    info: *const rocksdb_compactionjobinfo_t,
) {
    let listener = unsafe { &*state.cast::<ListenerState>() };
    listener.compactions.fetch_add(1, Ordering::SeqCst);

    let mut cf_len = 0;
    let cf = unsafe { rocksdb_compactionjobinfo_cf_name(info, &raw mut cf_len) };
    let cf_name = String::from_utf8_lossy(unsafe { slice_from_raw(cf, cf_len) }).into_owned();

    // Borrowed view: do not call oxrocksdb_table_properties_collection_destroy.
    let collection = unsafe { oxrocksdb_compactionjobinfo_table_properties(info) };
    let len = unsafe { oxrocksdb_table_properties_collection_len(collection) };
    let message = format!(
        "compaction completed on CF '{cf_name}', table_properties collection len={len}"
    );
    eprintln!("{message}");
    listener
        .messages
        .lock()
        .expect("listener mutex")
        .push(message);
}

fn lookup_user_prop(
    props: *const oxrocksdb_user_collected_properties_t,
    key: &[u8],
) -> Option<Vec<u8>> {
    let mut vlen = 0;
    // SAFETY: `props` is a borrowed RocksDB map; get does not allocate the map.
    let val = unsafe {
        oxrocksdb_user_collected_properties_get(
            props,
            key.as_ptr().cast(),
            key.len(),
            &raw mut vlen,
        )
    };
    if val.is_null() {
        None
    } else {
        Some(unsafe { slice_from_raw(val, vlen) }.to_vec())
    }
}

fn user_props_to_map(
    props: *const oxrocksdb_user_collected_properties_t,
) -> std::collections::BTreeMap<Vec<u8>, Vec<u8>> {
    let mut out = std::collections::BTreeMap::new();
    // SAFETY: iterator is heap-allocated by the shim and must be destroyed.
    let it = unsafe { oxrocksdb_user_collected_properties_iter_create(props) };
    assert!(!it.is_null());
    unsafe {
        while oxrocksdb_user_collected_properties_iter_valid(it) != 0 {
            let mut klen = 0;
            let mut vlen = 0;
            let key = slice_from_raw(
                oxrocksdb_user_collected_properties_iter_key(it, &raw mut klen),
                klen,
            );
            let val = slice_from_raw(
                oxrocksdb_user_collected_properties_iter_value(it, &raw mut vlen),
                vlen,
            );
            out.insert(key.to_vec(), val.to_vec());
            oxrocksdb_user_collected_properties_iter_next(it);
        }
        oxrocksdb_user_collected_properties_iter_destroy(it);
    }
    out
}

struct DemoDb {
    dir: TempDir,
    db: *mut rocksdb_t,
    db_options: *mut rocksdb_options_t,
    cf_options: *mut rocksdb_options_t,
    cf: *mut rocksdb_column_family_handle_t,
    write_options: *mut rocksdb_writeoptions_t,
    flush_options: *mut rocksdb_flushoptions_t,
    listener: Arc<ListenerState>,
}

impl DemoDb {
    fn open() -> Self {
        let dir = TempDir::new().expect("tempdir");
        let path = cstring(dir.path().to_str().expect("utf8 path"));

        // SAFETY: all pointers below are created by the RocksDB C API and
        // checked for NULL before use.
        unsafe {
            let db_options = rocksdb_options_create();
            assert!(!db_options.is_null());
            rocksdb_options_set_create_if_missing(db_options, 1);
            rocksdb_options_set_create_missing_column_families(db_options, 1);
            rocksdb_options_set_disable_auto_compactions(db_options, 1);

            let listener_state = Arc::new(ListenerState {
                flush_begins: AtomicU64::new(0),
                flushes: AtomicU64::new(0),
                compaction_begins: AtomicU64::new(0),
                compactions: AtomicU64::new(0),
                messages: Mutex::new(Vec::new()),
            });
            let listener_ptr = Arc::into_raw(Arc::clone(&listener_state));
            let listener = oxrocksdb_eventlistener_create(
                listener_ptr.cast::<c_void>().cast_mut(),
                Some(listener_destruct),
                Some(on_flush_begin),
                Some(on_flush_completed),
                Some(on_compaction_begin),
                Some(on_compaction_completed),
            );
            assert!(!listener.is_null());
            oxrocksdb_options_add_eventlistener(db_options, listener);
            // Wrapper only: RocksDB's shared_ptr keeps the impl alive.
            oxrocksdb_eventlistener_destroy(listener);

            let mut err = ptr::null_mut();
            let db = rocksdb_open(db_options, path.as_ptr(), &raw mut err);
            take_error(err).expect("rocksdb_open");
            assert!(!db.is_null());

            let cf_options = rocksdb_options_create();
            assert!(!cf_options.is_null());
            rocksdb_options_set_disable_auto_compactions(cf_options, 1);

            let factory_state = Box::into_raw(Box::new(FactoryState {
                collectors_created: AtomicU64::new(0),
            }));
            let factory = oxrocksdb_table_properties_collector_factory_create(
                factory_state.cast(),
                Some(factory_name),
                Some(factory_destruct),
                Some(create_collector),
            );
            assert!(!factory.is_null());
            oxrocksdb_options_add_table_properties_collector_factory(cf_options, factory);
            // Wrapper only: options hold a shared_ptr copy of the factory.
            oxrocksdb_table_properties_collector_factory_destroy(factory);

            let cf_name = cstring("props");
            err = ptr::null_mut();
            let cf = rocksdb_create_column_family(
                db,
                cf_options,
                cf_name.as_ptr(),
                &raw mut err,
            );
            take_error(err).expect("create_column_family");
            assert!(!cf.is_null());

            let write_options = rocksdb_writeoptions_create();
            let flush_options = rocksdb_flushoptions_create();
            assert!(!write_options.is_null() && !flush_options.is_null());

            Self {
                dir,
                db,
                db_options,
                cf_options,
                cf,
                write_options,
                flush_options,
                listener: listener_state,
            }
        }
    }

    fn put(&self, key: &[u8], value: &[u8]) {
        let mut err = ptr::null_mut();
        unsafe {
            rocksdb_put_cf(
                self.db,
                self.write_options,
                self.cf,
                key.as_ptr().cast(),
                key.len(),
                value.as_ptr().cast(),
                value.len(),
                &raw mut err,
            );
        }
        unsafe { take_error(err) }.expect("put_cf");
    }

    fn flush(&self) {
        let mut err = ptr::null_mut();
        unsafe {
            rocksdb_flush_cf(self.db, self.flush_options, self.cf, &raw mut err);
        }
        unsafe { take_error(err) }.expect("flush_cf");
    }

    fn compact(&self) {
        unsafe {
            rocksdb_compact_range_cf(self.db, self.cf, ptr::null(), 0, ptr::null(), 0);
        }
    }

    fn properties_collection(&self) -> *mut oxrocksdb_table_properties_collection_t {
        let mut err = ptr::null_mut();
        let collection =
            unsafe { oxrocksdb_get_properties_of_all_tables_cf(self.db, self.cf, &raw mut err) };
        unsafe { take_error(err) }.expect("get_properties_of_all_tables_cf");
        assert!(!collection.is_null());
        collection
    }
}

impl Drop for DemoDb {
    fn drop(&mut self) {
        unsafe {
            rocksdb_column_family_handle_destroy(self.cf);
            rocksdb_close(self.db);
            rocksdb_options_destroy(self.cf_options);
            rocksdb_options_destroy(self.db_options);
            rocksdb_writeoptions_destroy(self.write_options);
            rocksdb_flushoptions_destroy(self.flush_options);
            let path = cstring(self.dir.path().to_str().expect("utf8 path"));
            let opts = rocksdb_options_create();
            let mut err = ptr::null_mut();
            rocksdb_destroy_db(opts, path.as_ptr(), &raw mut err);
            rocksdb_options_destroy(opts);
            drop(take_error(err));
        }
    }
}

#[test]
fn collector_create_and_destroy_without_handing_to_rocksdb() {
    let state = Box::into_raw(Box::new(CollectorState {
        puts: AtomicU64::new(0),
        bytes: AtomicU64::new(0),
    }));
    unsafe {
        let collector = oxrocksdb_table_properties_collector_create(
            state.cast(),
            Some(collector_name),
            Some(collector_destruct),
            Some(collector_add),
            Some(collector_finish),
        );
        assert!(!collector.is_null());
        // Never given to RocksDB, so we own the delete.
        oxrocksdb_table_properties_collector_destroy(collector);
    }
}

#[test]
fn factory_destroy_without_adding_to_options_runs_destructor() {
    let state = Box::into_raw(Box::new(FactoryState {
        collectors_created: AtomicU64::new(0),
    }));
    unsafe {
        let factory = oxrocksdb_table_properties_collector_factory_create(
            state.cast(),
            Some(factory_name),
            Some(factory_destruct),
            Some(create_collector),
        );
        assert!(!factory.is_null());
        oxrocksdb_table_properties_collector_factory_destroy(factory);
    }
}

#[test]
fn null_listener_callbacks_do_not_crash() {
    let dir = TempDir::new().expect("tempdir");
    let path = cstring(dir.path().to_str().expect("utf8 path"));
    unsafe {
        let options = rocksdb_options_create();
        rocksdb_options_set_create_if_missing(options, 1);
        let listener =
            oxrocksdb_eventlistener_create(
                ptr::null_mut(),
                None,
                None,
                None,
                None,
                None,
            );
        assert!(!listener.is_null());
        oxrocksdb_options_add_eventlistener(options, listener);
        oxrocksdb_eventlistener_destroy(listener);

        let mut err = ptr::null_mut();
        let db = rocksdb_open(options, path.as_ptr(), &raw mut err);
        take_error(err).expect("open");
        let write_options = rocksdb_writeoptions_create();
        let flush_options = rocksdb_flushoptions_create();
        rocksdb_put(
            db,
            write_options,
            b"k".as_ptr().cast(),
            1,
            b"v".as_ptr().cast(),
            1,
            &raw mut err,
        );
        take_error(err).expect("put");
        rocksdb_flush(db, flush_options, &raw mut err);
        take_error(err).expect("flush");

        rocksdb_writeoptions_destroy(write_options);
        rocksdb_flushoptions_destroy(flush_options);
        rocksdb_close(db);
        rocksdb_options_destroy(options);
    }
}

#[test]
fn empty_cf_table_properties_collection_is_empty() {
    let db = DemoDb::open();
    let collection = db.properties_collection();
    unsafe {
        assert_eq!(oxrocksdb_table_properties_collection_len(collection), 0);
        let it = oxrocksdb_table_properties_collection_iter_create(collection);
        assert_eq!(oxrocksdb_table_properties_collection_iter_valid(it), 0);
        oxrocksdb_table_properties_collection_iter_destroy(it);
        oxrocksdb_table_properties_collection_destroy(collection);
    }
}

#[test]
fn user_collected_properties_end_to_end() {
    let db = DemoDb::open();

    let p_coll = db.properties_collection();
    unsafe {
        assert_eq!(oxrocksdb_table_properties_collection_len(p_coll), 0);
        oxrocksdb_table_properties_collection_destroy(p_coll);
    }

    const N: u64 = 50;
    for i in 0..N {
        let key = format!("k{i:03}");
        let value = format!("value-{i}");
        db.put(key.as_bytes(), value.as_bytes());
    }
    db.flush();

    assert!(
        db.listener.flush_begins.load(Ordering::SeqCst) >= 1,
        "expected at least one flush begin event"
    );
    assert!(
        db.listener.flushes.load(Ordering::SeqCst) >= 1,
        "expected at least one flush completed event"
    );

    let collection = db.properties_collection();
    let mut found_marker = false;
    let mut total_puts = 0_u64;
    unsafe {
        assert!(oxrocksdb_table_properties_collection_len(collection) >= 1);
        let it = oxrocksdb_table_properties_collection_iter_create(collection);
        assert!(!it.is_null());
        while oxrocksdb_table_properties_collection_iter_valid(it) != 0 {
            let mut klen = 0;
            let sst_name = slice_from_raw(
                oxrocksdb_table_properties_collection_iter_key(it, &raw mut klen),
                klen,
            );
            eprintln!(
                "SST in collection: {}",
                String::from_utf8_lossy(sst_name)
            );

            let table_props = oxrocksdb_table_properties_collection_iter_value(it);
            assert!(!table_props.is_null(), "iterator value must not be null");

            let user_props = oxrocksdb_table_properties_get_user_properties(table_props);
            assert!(!user_props.is_null());
            assert!(oxrocksdb_user_collected_properties_len(user_props) >= 3);

            let map = user_props_to_map(user_props);
            assert_eq!(map.get(PROP_MARKER).map(Vec::as_slice), Some(PROP_MARKER_VALUE));
            found_marker = true;

            let puts = lookup_user_prop(user_props, PROP_NUM_PUTS)
                .expect("ox.num_puts");
            total_puts += String::from_utf8(puts)
                .expect("utf8")
                .parse::<u64>()
                .expect("num_puts");

            assert!(lookup_user_prop(user_props, PROP_NUM_BYTES).is_some());
            assert!(
                lookup_user_prop(user_props, b"does-not-exist").is_none(),
                "missing user property key must return null"
            );

            oxrocksdb_table_properties_collection_iter_next(it);
        }
        oxrocksdb_table_properties_collection_iter_destroy(it);
        oxrocksdb_table_properties_collection_destroy(collection);
    }
    assert!(found_marker);
    assert_eq!(total_puts, N);

    for i in N..(N * 2) {
        let key = format!("k{i:03}");
        let value = format!("value-{i}");
        db.put(key.as_bytes(), value.as_bytes());
    }
    db.flush();
    db.compact();

    assert!(
        db.listener.compaction_begins.load(Ordering::SeqCst) >= 1,
        "expected at least one compaction begin event, messages={:?}",
        db.listener.messages.lock().expect("mutex")
    );
    assert!(
        db.listener.compactions.load(Ordering::SeqCst) >= 1,
        "expected at least one compaction completed event, messages={:?}",
        db.listener.messages.lock().expect("mutex")
    );

    let messages = db.listener.messages.lock().expect("mutex");
    assert!(
        messages.iter().any(|m| m.contains("flush begin")),
        "missing flush begin message: {messages:?}"
    );
    assert!(
        messages.iter().any(|m| m.contains("flush completed")),
        "missing flush completed message: {messages:?}"
    );
    assert!(
        messages.iter().any(|m| m.contains("compaction begin")),
        "missing compaction begin message: {messages:?}"
    );
    assert!(
        messages.iter().any(|m| m.contains("compaction completed")),
        "missing compaction completed message: {messages:?}"
    );
}