rust-rocksdb/tests/test_column_family.rs

// Copyright 2020 Tyler Neely
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

mod util;

use pretty_assertions::assert_eq;

use rocksdb::{ColumnFamilyDescriptor, MergeOperands, Options, DB, DEFAULT_COLUMN_FAMILY_NAME};
use rocksdb::{TransactionDB, TransactionDBOptions};
use util::DBPath;

use std::fs;
use std::io;
use std::path::Path;

#[cfg(feature = "multi-threaded-cf")]
use rocksdb::MultiThreaded;
#[cfg(not(feature = "multi-threaded-cf"))]
use rocksdb::SingleThreaded;

fn dir_size(path: impl AsRef<Path>) -> io::Result<u64> {
    fn dir_size(mut dir: fs::ReadDir) -> io::Result<u64> {
        dir.try_fold(0, |acc, file| {
            let file = file?;
            let size = match file.metadata()? {
                data if data.is_dir() => dir_size(fs::read_dir(file.path())?)?,
                data => data.len(),
            };
            Ok(acc + size)
        })
    }

    dir_size(fs::read_dir(path)?)
}

#[test]
fn test_column_family() {
    let n = DBPath::new("_rust_rocksdb_cftest");

    // should be able to create column families
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.set_merge_operator_associative("test operator", test_provided_merge);
        #[cfg(feature = "multi-threaded-cf")]
        let db = DB::open(&opts, &n).unwrap();
        #[cfg(not(feature = "multi-threaded-cf"))]
        let mut db = DB::open(&opts, &n).unwrap();
        let opts = Options::default();
        match db.create_cf("cf1", &opts) {
            Ok(()) => println!("cf1 created successfully"),
            Err(e) => {
                panic!("could not create column family: {}", e);
            }
        }
    }

    // should fail to open db without specifying same column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator_associative("test operator", test_provided_merge);
        match DB::open(&opts, &n) {
            Ok(_db) => panic!(
                "should not have opened DB successfully without \
                        specifying column
            families"
            ),
            Err(e) => assert!(e.to_string().starts_with("Invalid argument")),
        }
    }

    // should properly open db when specyfing all column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator_associative("test operator", test_provided_merge);
        match DB::open_cf(&opts, &n, ["cf1"]) {
            Ok(_db) => println!("successfully opened db with column family"),
            Err(e) => panic!("failed to open db with column family: {}", e),
        }
    }

    // should be able to list a cf
    {
        let opts = Options::default();
        let vec = DB::list_cf(&opts, &n);
        match vec {
            Ok(vec) => assert_eq!(vec, vec![DEFAULT_COLUMN_FAMILY_NAME, "cf1"]),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }

    // TODO should be able to use writebatch ops with a cf
    {}
    // TODO should be able to iterate over a cf
    {}
    // should b able to drop a cf
    {
        #[cfg(feature = "multi-threaded-cf")]
        let db = DB::open_cf(&Options::default(), &n, ["cf1"]).unwrap();
        #[cfg(not(feature = "multi-threaded-cf"))]
        let mut db = DB::open_cf(&Options::default(), &n, ["cf1"]).unwrap();

        match db.drop_cf("cf1") {
            Ok(_) => println!("cf1 successfully dropped."),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }
}

#[test]
fn test_column_family_with_transactiondb() {
    let n = DBPath::new("_rust_rocksdb_cftest");

    // should be able to create column families
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.set_merge_operator_associative("test operator", test_provided_merge);
        #[cfg(feature = "multi-threaded-cf")]
        let db = TransactionDB::open(&opts, &TransactionDBOptions::default(), &n).unwrap();
        #[cfg(not(feature = "multi-threaded-cf"))]
        let db = TransactionDB::open(&opts, &TransactionDBOptions::default(), &n).unwrap();
        let opts = Options::default();
        match db.create_cf("cf1", &opts) {
            Ok(()) => println!("cf1 created successfully"),
            Err(e) => {
                panic!("could not create column family: {}", e);
            }
        }
    }

    // should fail to open db without specifying same column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator_associative("test operator", test_provided_merge);
        #[cfg(feature = "multi-threaded-cf")]
        let db = TransactionDB::<MultiThreaded>::open(&opts, &TransactionDBOptions::default(), &n);
        #[cfg(not(feature = "multi-threaded-cf"))]
        let db = TransactionDB::<SingleThreaded>::open(&opts, &TransactionDBOptions::default(), &n);
        match db {
            Ok(_db) => panic!(
                "should not have opened TransactionDB successfully without \
                        specifying column
            families"
            ),
            Err(e) => assert!(e.to_string().starts_with("Invalid argument")),
        }
    }

    // should properly open db when specyfing all column families
    {
        let mut opts = Options::default();
        opts.set_merge_operator_associative("test operator", test_provided_merge);
        let cfs = &["cf1"];
        #[cfg(feature = "multi-threaded-cf")]
        let db = TransactionDB::<MultiThreaded>::open_cf(
            &opts,
            &TransactionDBOptions::default(),
            &n,
            cfs,
        );
        #[cfg(not(feature = "multi-threaded-cf"))]
        let db = TransactionDB::<SingleThreaded>::open_cf(
            &opts,
            &TransactionDBOptions::default(),
            &n,
            cfs,
        );
        match db {
            Ok(_db) => println!("successfully opened db with column family"),
            Err(e) => panic!("failed to open db with column family: {}", e),
        }
    }

    // should be able to list a cf
    {
        let opts = Options::default();
        let vec = DB::list_cf(&opts, &n);
        match vec {
            Ok(vec) => assert_eq!(vec, vec![DEFAULT_COLUMN_FAMILY_NAME, "cf1"]),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }

    // should b able to drop a cf
    {
        let opts = Options::default();
        let cfs = &["cf1"];
        #[cfg(feature = "multi-threaded-cf")]
        let db = TransactionDB::<MultiThreaded>::open_cf(
            &opts,
            &TransactionDBOptions::default(),
            &n,
            cfs,
        )
        .unwrap();
        #[cfg(not(feature = "multi-threaded-cf"))]
        let mut db = TransactionDB::<SingleThreaded>::open_cf(
            &opts,
            &TransactionDBOptions::default(),
            &n,
            cfs,
        )
        .unwrap();
        match db.drop_cf("cf1") {
            Ok(_) => println!("cf1 successfully dropped."),
            Err(e) => panic!("failed to drop column family: {}", e),
        }
    }
    // should not be able to open cf after dropping.
    {
        let opts = Options::default();
        let cfs = &["cf1"];
        #[cfg(feature = "multi-threaded-cf")]
        let db = TransactionDB::<MultiThreaded>::open_cf(
            &opts,
            &TransactionDBOptions::default(),
            &n,
            cfs,
        );
        #[cfg(not(feature = "multi-threaded-cf"))]
        let db = TransactionDB::<SingleThreaded>::open_cf(
            &opts,
            &TransactionDBOptions::default(),
            &n,
            cfs,
        );
        assert!(db.is_err())
    }
}

#[test]
fn test_can_open_db_with_results_of_list_cf() {
    // Test scenario derived from GitHub issue #175 and 177

    let n = DBPath::new("_rust_rocksdb_cftest_with_list_cf");

    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        #[cfg(feature = "multi-threaded-cf")]
        let db = DB::open(&opts, &n).unwrap();
        #[cfg(not(feature = "multi-threaded-cf"))]
        let mut db = DB::open(&opts, &n).unwrap();
        let opts = Options::default();

        assert!(db.create_cf("cf1", &opts).is_ok());
    }

    {
        let options = Options::default();
        let cfs = DB::list_cf(&options, &n).unwrap();
        let db = DB::open_cf(&options, &n, cfs).unwrap();

        assert!(db.cf_handle("cf1").is_some());
    }
}

#[test]
fn test_create_missing_column_family() {
    let n = DBPath::new("_rust_rocksdb_missing_cftest");

    // should be able to create new column families when opening a new database
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        match DB::open_cf(&opts, &n, ["cf1"]) {
            Ok(_db) => println!("successfully created new column family"),
            Err(e) => panic!("failed to create new column family: {}", e),
        }
    }
}

#[test]
fn test_open_column_family_with_opts() {
    let n = DBPath::new("_rust_rocksdb_open_cf_with_opts");

    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        // We can use different parameters for different column family.
        let mut cf1_opts = Options::default();
        cf1_opts.set_min_write_buffer_number(2);
        cf1_opts.set_min_write_buffer_number_to_merge(4);
        let mut cf2_opts = Options::default();
        cf2_opts.set_min_write_buffer_number(5);
        cf2_opts.set_min_write_buffer_number_to_merge(10);

        let cfs = vec![("cf1", cf1_opts), ("cf2", cf2_opts)];
        match DB::open_cf_with_opts(&opts, &n, cfs) {
            Ok(_db) => println!("successfully opened column family with the specified options"),
            Err(e) => panic!("failed to open cf with options: {}", e),
        }
    }
}

#[test]
#[ignore]
fn test_merge_operator() {
    let n = DBPath::new("_rust_rocksdb_cftest_merge");
    // TODO should be able to write, read, merge, batch, and iterate over a cf
    {
        let mut opts = Options::default();
        opts.set_merge_operator_associative("test operator", test_provided_merge);
        let db = match DB::open_cf(&opts, &n, ["cf1"]) {
            Ok(db) => {
                println!("successfully opened db with column family");
                db
            }
            Err(e) => panic!("failed to open db with column family: {}", e),
        };
        let cf1 = db.cf_handle("cf1").unwrap();
        assert!(db.put_cf(&cf1, b"k1", b"v1").is_ok());
        assert_eq!(db.get_cf(&cf1, b"k1").unwrap().unwrap(), b"v1");
        let p = db.put_cf(&cf1, b"k1", b"a");
        assert!(p.is_ok());
        db.merge_cf(&cf1, b"k1", b"b").unwrap();
        db.merge_cf(&cf1, b"k1", b"c").unwrap();
        db.merge_cf(&cf1, b"k1", b"d").unwrap();
        db.merge_cf(&cf1, b"k1", b"efg").unwrap();
        let m = db.merge_cf(&cf1, b"k1", b"h");
        println!("m is {m:?}");
        // TODO assert!(m.is_ok());
        match db.get(b"k1") {
            Ok(Some(value)) => match std::str::from_utf8(&value) {
                Ok(v) => println!("retrieved utf8 value: {v}"),
                Err(_) => println!("did not read valid utf-8 out of the db"),
            },
            Err(_) => println!("error reading value"),
            _ => panic!("value not present!"),
        }

        let _ = db.get_cf(&cf1, b"k1");
        // TODO assert!(r.unwrap().as_ref() == b"abcdefgh");
        assert!(db.delete(b"k1").is_ok());
        assert!(db.get(b"k1").unwrap().is_none());
    }
}

fn test_provided_merge(
    _: &[u8],
    existing_val: Option<&[u8]>,
    operands: &MergeOperands,
) -> Option<Vec<u8>> {
    let nops = operands.len();
    let mut result: Vec<u8> = Vec::with_capacity(nops);
    if let Some(v) = existing_val {
        for e in v {
            result.push(*e);
        }
    }
    for op in operands {
        for e in op {
            result.push(*e);
        }
    }
    Some(result)
}

#[test]
fn test_column_family_with_options() {
    let n = DBPath::new("_rust_rocksdb_cf_with_optionstest");
    {
        let mut cfopts = Options::default();
        cfopts.set_max_write_buffer_number(16);
        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);

        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let cfs = vec![cf_descriptor];
        match DB::open_cf_descriptors(&opts, &n, cfs) {
            Ok(_db) => println!("created db with column family descriptors successfully"),
            Err(e) => {
                panic!(
                    "could not create new database with column family descriptors: {}",
                    e
                );
            }
        }
    }

    {
        let mut cfopts = Options::default();
        cfopts.set_max_write_buffer_number(16);
        let cf_descriptor = ColumnFamilyDescriptor::new("cf1", cfopts);

        let opts = Options::default();
        let cfs = vec![cf_descriptor];

        match DB::open_cf_descriptors(&opts, &n, cfs) {
            Ok(_db) => println!("successfully re-opened database with column family descriptors"),
            Err(e) => {
                panic!(
                    "unable to re-open database with column family descriptors: {}",
                    e
                );
            }
        }
    }
}

#[test]
fn test_create_duplicate_column_family() {
    let n = DBPath::new("_rust_rocksdb_create_duplicate_column_family");
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        #[cfg(feature = "multi-threaded-cf")]
        let db = DB::open_cf(&opts, &n, ["cf1"]).unwrap();
        #[cfg(not(feature = "multi-threaded-cf"))]
        let mut db = DB::open_cf(&opts, &n, ["cf1"]).unwrap();

        assert!(db.create_cf("cf1", &opts).is_err());
    }
}

#[test]
fn test_no_leaked_column_family() {
    let n = DBPath::new("_rust_rocksdb_no_leaked_column_family");
    {
        let mut opts = Options::default();
        opts.create_if_missing(true);
        opts.create_missing_column_families(true);

        let mut write_options = rocksdb::WriteOptions::default();
        write_options.set_sync(false);
        write_options.disable_wal(true);

        #[cfg(feature = "multi-threaded-cf")]
        let db = DB::open(&opts, &n).unwrap();
        #[cfg(not(feature = "multi-threaded-cf"))]
        let mut db = DB::open(&opts, &n).unwrap();

        #[cfg(feature = "multi-threaded-cf")]
        let mut outlived_cf = None;

        let large_blob = vec![0x20; 1024 * 1024];

        // repeat creating and dropping cfs many time to indirectly detect
        // possible leak via large dir.
        for cf_index in 0..20 {
            let cf_name = format!("cf{cf_index}");
            db.create_cf(&cf_name, &Options::default()).unwrap();
            let cf = db.cf_handle(&cf_name).unwrap();

            let mut batch = rocksdb::WriteBatch::default();
            for key_index in 0..100 {
                batch.put_cf(&cf, format!("k{key_index}"), &large_blob);
            }
            db.write_opt(batch, &write_options).unwrap();

            // force create an SST file
            db.flush_cf(&cf).unwrap();
            db.drop_cf(&cf_name).unwrap();

            #[cfg(feature = "multi-threaded-cf")]
            {
                outlived_cf = Some(cf);
            }
        }

        // if we're not leaking, the dir bytes should be well under 10M bytes in total
        let dir_bytes = dir_size(&n).unwrap();
        let leak_msg = format!("{dir_bytes} is too large (maybe leaking...)");
        assert!(dir_bytes < 10_000_000, "{}", leak_msg);

        // only if MultiThreaded, cf can outlive db.drop_cf() and shouldn't cause SEGV...
        #[cfg(feature = "multi-threaded-cf")]
        {
            let outlived_cf = outlived_cf.unwrap();
            assert_eq!(
                &db.get_cf(&outlived_cf, "k0").unwrap().unwrap(),
                &large_blob
            );
            drop(outlived_cf);
        }

        // make it explicit not to drop the db until we get dir size above...
        drop(db);
    }
}