diff --git a/Cargo.toml b/Cargo.toml index 62128dd93..04ac3aa82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,3 +32,6 @@ tempdir = "0.3.4" [dependencies.librocksdb_sys] path = "librocksdb_sys" + +[dev-dependencies] +byteorder = "1.0.0" diff --git a/librocksdb_sys/crocksdb/c.cc b/librocksdb_sys/crocksdb/c.cc index 2d56c36e5..02bf99acc 100644 --- a/librocksdb_sys/crocksdb/c.cc +++ b/librocksdb_sys/crocksdb/c.cc @@ -27,6 +27,7 @@ #include "rocksdb/statistics.h" #include "rocksdb/slice_transform.h" #include "rocksdb/table.h" +#include "rocksdb/table_properties.h" #include "rocksdb/rate_limiter.h" #include "rocksdb/utilities/backupable_db.h" @@ -83,6 +84,13 @@ using rocksdb::HistogramData; using rocksdb::PinnableSlice; using rocksdb::FilterBitsBuilder; using rocksdb::FilterBitsReader; +using rocksdb::EntryType; +using rocksdb::SequenceNumber; +using rocksdb::UserCollectedProperties; +using rocksdb::TableProperties; +using rocksdb::TablePropertiesCollection; +using rocksdb::TablePropertiesCollector; +using rocksdb::TablePropertiesCollectorFactory; using std::shared_ptr; @@ -654,6 +662,10 @@ void crocksdb_drop_column_family( SaveError(errptr, db->rep->DropColumnFamily(handle->rep)); } +uint32_t crocksdb_column_family_handle_id(crocksdb_column_family_handle_t* handle) { + return handle->rep->GetID(); +} + void crocksdb_column_family_handle_destroy(crocksdb_column_family_handle_t* handle) { delete handle->rep; delete handle; @@ -2919,4 +2931,341 @@ const char* crocksdb_pinnableslice_value(const crocksdb_pinnableslice_t* v, *vlen = v->rep.size(); return v->rep.data(); } + +/* Table Properties */ + +struct crocksdb_user_collected_properties_t { + UserCollectedProperties* rep_ = nullptr; +}; + +void crocksdb_user_collected_properties_add( + crocksdb_user_collected_properties_t* props, + const char* k, size_t klen, + const char* v, size_t vlen) { + props->rep_->emplace(std::make_pair(std::string(k, klen), std::string(v, vlen))); +} + +struct crocksdb_user_collected_properties_iterator_t { + UserCollectedProperties::iterator cur_; + UserCollectedProperties::iterator end_; +}; + +crocksdb_user_collected_properties_iterator_t* +crocksdb_user_collected_properties_iter_create( + crocksdb_user_collected_properties_t* props) { + auto it = new crocksdb_user_collected_properties_iterator_t; + it->cur_ = props->rep_->begin(); + it->end_ = props->rep_->end(); + return it; +} + +void crocksdb_user_collected_properties_iter_destroy( + crocksdb_user_collected_properties_iterator_t* it) { + delete it; +} + +unsigned char crocksdb_user_collected_properties_iter_valid( + crocksdb_user_collected_properties_iterator_t* it) { + return it->cur_ != it->end_; +} + +void crocksdb_user_collected_properties_iter_next( + crocksdb_user_collected_properties_iterator_t* it) { + ++(it->cur_); +} + +const char* crocksdb_user_collected_properties_iter_key( + crocksdb_user_collected_properties_iterator_t* it, size_t* klen) { + *klen = it->cur_->first.size(); + return it->cur_->first.data(); +} + +const char* crocksdb_user_collected_properties_iter_value( + crocksdb_user_collected_properties_iterator_t* it, size_t* vlen) { + *vlen = it->cur_->second.size(); + return it->cur_->second.data(); +} + +struct crocksdb_table_properties_t { + std::shared_ptr rep_; + crocksdb_user_collected_properties_t users_; + + void init(std::shared_ptr rep) { + rep_ = rep; + users_.rep_ = const_cast(&rep->user_collected_properties); + } +}; + +crocksdb_table_properties_t* crocksdb_table_properties_create() { + return new crocksdb_table_properties_t; +} + +void crocksdb_table_properties_destroy(crocksdb_table_properties_t* props) { + delete props; +} + +uint64_t crocksdb_table_properties_get_u64(crocksdb_table_properties_t* props, + crocksdb_table_property_t prop) { + auto rep = props->rep_; + switch (prop) { + case kDataSize: return rep->data_size; + case kIndexSize: return rep->index_size; + case kFilterSize: return rep->filter_size; + case kRawKeySize: return rep->raw_key_size; + case kRawValueSize: return rep->raw_value_size; + case kNumDataBlocks: return rep->num_data_blocks; + case kNumEntries: return rep->num_entries; + case kFormatVersion: return rep->format_version; + case kFixedKeyLen: return rep->data_size; + case kColumnFamilyID: return rep->column_family_id; + } + return 0; +} + +const char* crocksdb_table_properties_get_str(crocksdb_table_properties_t* props, + crocksdb_table_property_t prop, size_t* slen) { + auto rep = props->rep_; + switch (prop) { + case kColumnFamilyName: + *slen = rep->column_family_name.size(); + return rep->column_family_name.data(); + case kFilterPolicyName: + *slen = rep->filter_policy_name.size(); + return rep->filter_policy_name.data(); + case kComparatorName: + *slen = rep->comparator_name.size(); + return rep->comparator_name.data(); + case kMergeOperatorName: + *slen = rep->merge_operator_name.size(); + return rep->merge_operator_name.data(); + case kPrefixExtractorName: + *slen = rep->prefix_extractor_name.size(); + return rep->prefix_extractor_name.data(); + case kPropertyCollectorsNames: + *slen = rep->property_collectors_names.size(); + return rep->property_collectors_names.data(); + case kCompressionName: + *slen = rep->compression_name.size(); + return rep->compression_name.data(); + } + return nullptr; +} + +crocksdb_user_collected_properties_t* +crocksdb_table_properties_get_user_properties(crocksdb_table_properties_t* props) { + return &props->users_; +} + +/* Table Properties Collection */ + +struct crocksdb_table_properties_collection_t { + TablePropertiesCollection rep_; +}; + +crocksdb_table_properties_collection_t* +crocksdb_table_properties_collection_create() { + return new crocksdb_table_properties_collection_t; +} + +void crocksdb_table_properties_collection_destroy( + crocksdb_table_properties_collection_t* collection) { + delete collection; +} + +struct crocksdb_table_properties_collection_iterator_t { + TablePropertiesCollection::iterator cur_; + TablePropertiesCollection::iterator end_; +}; + +crocksdb_table_properties_collection_iterator_t* +crocksdb_table_properties_collection_iter_create( + crocksdb_table_properties_collection_t* collection) { + auto it = new crocksdb_table_properties_collection_iterator_t; + it->cur_ = collection->rep_.begin(); + it->end_ = collection->rep_.end(); + return it; +} + +void crocksdb_table_properties_collection_iter_destroy( + crocksdb_table_properties_collection_iterator_t* it) { + delete it; +} + +unsigned char crocksdb_table_properties_collection_iter_valid( + crocksdb_table_properties_collection_iterator_t* it) { + return it->cur_ != it->end_; +} + +void crocksdb_table_properties_collection_iter_next( + crocksdb_table_properties_collection_iterator_t* it) { + ++(it->cur_); +} + +const char* crocksdb_table_properties_collection_iter_key( + crocksdb_table_properties_collection_iterator_t* it, size_t* klen) { + *klen = it->cur_->first.size(); + return it->cur_->first.data(); +} + +void crocksdb_table_properties_collection_iter_value( + crocksdb_table_properties_collection_iterator_t* it, crocksdb_table_properties_t* props) { + props->init(it->cur_->second); +} + +/* Table Properties Collector */ + +struct crocksdb_table_properties_collector_t : public TablePropertiesCollector { + void* state_; + const char* (*name_)(void*); + void (*destruct_)(void*); + void (*add_)(void*, + const char* key, size_t key_len, + const char* value, size_t value_len, + int entry_type, uint64_t seq, uint64_t file_size); + void (*finish_)(void*, crocksdb_user_collected_properties_t* props); + + virtual ~crocksdb_table_properties_collector_t() { + destruct_(state_); + } + + virtual Status AddUserKey(const Slice& key, + const Slice& value, + EntryType entry_type, + SequenceNumber seq, + uint64_t file_size) override { + add_(state_, + key.data(), key.size(), + value.data(), value.size(), + entry_type, seq, file_size); + return Status::OK(); + } + + virtual Status Finish(UserCollectedProperties* rep) override { + crocksdb_user_collected_properties_t props; + props.rep_ = rep; + finish_(state_, &props); + return Status::OK(); + } + + virtual UserCollectedProperties GetReadableProperties() const override { + // Seems rocksdb will not return the readable properties and we don't need them too. + return UserCollectedProperties(); + } + + const char* Name() const override { + return name_(state_); + } +}; + +crocksdb_table_properties_collector_t* +crocksdb_table_properties_collector_create( + void* state, + const char* (*name)(void*), + void (*destruct)(void*), + void (*add)(void*, + const char* key, size_t key_len, + const char* value, size_t value_len, + int entry_type, uint64_t seq, uint64_t file_size), + void (*finish)(void*, crocksdb_user_collected_properties_t* props)) { + auto c = new crocksdb_table_properties_collector_t; + c->state_ = state; + c->name_ = name; + c->destruct_ = destruct; + c->add_ = add; + c->finish_ = finish; + return c; +} + +void crocksdb_table_properties_collector_destroy(crocksdb_table_properties_collector_t* c) { + delete c; +} + +/* Table Properties Collector Factory */ + +struct crocksdb_table_properties_collector_factory_t : public TablePropertiesCollectorFactory { + void* state_; + const char* (*name_)(void*); + void (*destruct_)(void*); + crocksdb_table_properties_collector_t* + (*create_table_properties_collector_)(void*, uint32_t cf); + + virtual ~crocksdb_table_properties_collector_factory_t() { + destruct_(state_); + } + + virtual TablePropertiesCollector* CreateTablePropertiesCollector ( + TablePropertiesCollectorFactory::Context ctx) override { + return create_table_properties_collector_(state_, ctx.column_family_id); + } + + const char* Name() const override { + return name_(state_); + } +}; + +crocksdb_table_properties_collector_factory_t* +crocksdb_table_properties_collector_factory_create( + void* state, + const char* (*name)(void*), + void (*destruct)(void*), + crocksdb_table_properties_collector_t* + (*create_table_properties_collector)(void*, uint32_t cf)) { + auto f = new crocksdb_table_properties_collector_factory_t; + f->state_ = state; + f->name_ = name; + f->destruct_ = destruct; + f->create_table_properties_collector_ = create_table_properties_collector; + return f; +} + +void crocksdb_table_properties_collector_factory_destroy( + crocksdb_table_properties_collector_factory_t* f) { + delete f; +} + +void crocksdb_options_add_table_properties_collector_factory( + crocksdb_options_t* opt, crocksdb_table_properties_collector_factory_t* f) { + opt->rep.table_properties_collector_factories.push_back( + std::shared_ptr(f)); +} + +/* Get Table Properties */ + +void crocksdb_get_properties_of_all_tables(crocksdb_t* db, + crocksdb_table_properties_collection_t* props, char** errptr) { + auto s = db->rep->GetPropertiesOfAllTables(&props->rep_); + if (!s.ok()) { + SaveError(errptr, s); + } +} + +void crocksdb_get_properties_of_all_tables_cf( + crocksdb_t* db, crocksdb_column_family_handle_t* cf, + crocksdb_table_properties_collection_t* props, char** errptr) { + auto s = db->rep->GetPropertiesOfAllTables(cf->rep, &props->rep_); + if (!s.ok()) { + SaveError(errptr, s); + } +} + +void crocksdb_get_properties_of_tables_in_range( + crocksdb_t* db, crocksdb_column_family_handle_t* cf, + int num_ranges, + const char* const* start_keys, const size_t* start_keys_lens, + const char* const* limit_keys, const size_t* limit_keys_lens, + crocksdb_table_properties_collection_t* props, char** errptr) { + std::vector ranges; + for (int i = 0; i < num_ranges; i++) { + ranges.emplace_back(Range(Slice(start_keys[i], start_keys_lens[i]), + Slice(limit_keys[i], limit_keys_lens[i]))); + } + auto s = db->rep->GetPropertiesOfTablesInRange(cf->rep, + ranges.data(), + ranges.size(), + &props->rep_); + if (!s.ok()) { + SaveError(errptr, s); + } +} + } // end extern "C" diff --git a/librocksdb_sys/crocksdb/rocksdb/c.h b/librocksdb_sys/crocksdb/rocksdb/c.h index 06f1fb43a..b9e8564ca 100644 --- a/librocksdb_sys/crocksdb/rocksdb/c.h +++ b/librocksdb_sys/crocksdb/rocksdb/c.h @@ -111,6 +111,39 @@ typedef struct crocksdb_ingestexternalfileoptions_t crocksdb_ingestexternalfileo typedef struct crocksdb_sstfilewriter_t crocksdb_sstfilewriter_t; typedef struct crocksdb_ratelimiter_t crocksdb_ratelimiter_t; typedef struct crocksdb_pinnableslice_t crocksdb_pinnableslice_t; +typedef struct crocksdb_user_collected_properties_t + crocksdb_user_collected_properties_t; +typedef struct crocksdb_user_collected_properties_iterator_t + crocksdb_user_collected_properties_iterator_t; +typedef struct crocksdb_table_properties_t crocksdb_table_properties_t; +typedef struct crocksdb_table_properties_collection_t + crocksdb_table_properties_collection_t; +typedef struct crocksdb_table_properties_collection_iterator_t + crocksdb_table_properties_collection_iterator_t; +typedef struct crocksdb_table_properties_collector_t + crocksdb_table_properties_collector_t; +typedef struct crocksdb_table_properties_collector_factory_t + crocksdb_table_properties_collector_factory_t; + +typedef enum crocksdb_table_property_t { + kDataSize = 1, + kIndexSize = 2, + kFilterSize = 3, + kRawKeySize = 4, + kRawValueSize = 5, + kNumDataBlocks = 6, + kNumEntries = 7, + kFormatVersion = 8, + kFixedKeyLen = 9, + kColumnFamilyID = 10, + kColumnFamilyName = 11, + kFilterPolicyName = 12, + kComparatorName = 13, + kMergeOperatorName = 14, + kPrefixExtractorName = 15, + kPropertyCollectorsNames = 16, + kCompressionName = 17, +} crocksdb_table_property_t; /* DB operations */ @@ -198,6 +231,9 @@ crocksdb_create_column_family(crocksdb_t* db, extern C_ROCKSDB_LIBRARY_API void crocksdb_drop_column_family( crocksdb_t* db, crocksdb_column_family_handle_t* handle, char** errptr); +extern C_ROCKSDB_LIBRARY_API uint32_t crocksdb_column_family_handle_id( + crocksdb_column_family_handle_t*); + extern C_ROCKSDB_LIBRARY_API void crocksdb_column_family_handle_destroy( crocksdb_column_family_handle_t*); @@ -1176,6 +1212,139 @@ extern C_ROCKSDB_LIBRARY_API void crocksdb_pinnableslice_destroy( extern C_ROCKSDB_LIBRARY_API const char* crocksdb_pinnableslice_value( const crocksdb_pinnableslice_t* t, size_t* vlen); +/* Table Properties */ + +extern C_ROCKSDB_LIBRARY_API crocksdb_table_properties_t* +crocksdb_table_properties_create(); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_table_properties_destroy(crocksdb_table_properties_t*); + +extern C_ROCKSDB_LIBRARY_API uint64_t +crocksdb_table_properties_get_u64(crocksdb_table_properties_t*, + crocksdb_table_property_t prop); + +extern C_ROCKSDB_LIBRARY_API const char* +crocksdb_table_properties_get_str(crocksdb_table_properties_t*, + crocksdb_table_property_t prop, size_t* slen); + +extern C_ROCKSDB_LIBRARY_API crocksdb_user_collected_properties_t* +crocksdb_table_properties_get_user_properties(crocksdb_table_properties_t*); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_user_collected_properties_add( + crocksdb_user_collected_properties_t*, + const char* key, size_t key_len, const char* value, size_t value_len); + +extern C_ROCKSDB_LIBRARY_API crocksdb_user_collected_properties_iterator_t* +crocksdb_user_collected_properties_iter_create( + crocksdb_user_collected_properties_t*); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_user_collected_properties_iter_destroy( + crocksdb_user_collected_properties_iterator_t*); + +extern C_ROCKSDB_LIBRARY_API unsigned char + crocksdb_user_collected_properties_iter_valid( + crocksdb_user_collected_properties_iterator_t*); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_user_collected_properties_iter_next( + crocksdb_user_collected_properties_iterator_t*); + +extern C_ROCKSDB_LIBRARY_API const char* +crocksdb_user_collected_properties_iter_key( + crocksdb_user_collected_properties_iterator_t*, size_t* klen); + +extern C_ROCKSDB_LIBRARY_API const char* +crocksdb_user_collected_properties_iter_value( + crocksdb_user_collected_properties_iterator_t*, size_t* vlen); + +/* Table Properties Collection */ + +extern C_ROCKSDB_LIBRARY_API crocksdb_table_properties_collection_t* +crocksdb_table_properties_collection_create(); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_table_properties_collection_destroy(crocksdb_table_properties_collection_t*); + +extern C_ROCKSDB_LIBRARY_API crocksdb_table_properties_collection_iterator_t* +crocksdb_table_properties_collection_iter_create( + crocksdb_table_properties_collection_t*); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_table_properties_collection_iter_destroy( + crocksdb_table_properties_collection_iterator_t*); + +extern C_ROCKSDB_LIBRARY_API unsigned char +crocksdb_table_properties_collection_iter_valid( + crocksdb_table_properties_collection_iterator_t*); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_table_properties_collection_iter_next( + crocksdb_table_properties_collection_iterator_t*); + +extern C_ROCKSDB_LIBRARY_API const char* +crocksdb_table_properties_collection_iter_key( + crocksdb_table_properties_collection_iterator_t*, size_t* klen); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_table_properties_collection_iter_value( + crocksdb_table_properties_collection_iterator_t*, crocksdb_table_properties_t* props); + +/* Table Properties Collector */ + +extern C_ROCKSDB_LIBRARY_API crocksdb_table_properties_collector_t* +crocksdb_table_properties_collector_create( + void* state, + const char* (*name)(void*), + void (*destruct)(void*), + void (*add)(void*, + const char* key, size_t key_len, + const char* value, size_t value_len, + int entry_type, uint64_t seq, uint64_t file_size), + void (*finish)(void*, crocksdb_user_collected_properties_t* props)); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_table_properties_collector_destroy(crocksdb_table_properties_collector_t*); + +/* Table Properties Collector Factory */ + +extern C_ROCKSDB_LIBRARY_API crocksdb_table_properties_collector_factory_t* +crocksdb_table_properties_collector_factory_create( + void* state, + const char* (*name)(void*), + void (*destruct)(void*), + crocksdb_table_properties_collector_t* + (*create_table_properties_collector)(void*, uint32_t cf)); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_table_properties_collector_factory_destroy( + crocksdb_table_properties_collector_factory_t*); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_options_add_table_properties_collector_factory( + crocksdb_options_t* opt, crocksdb_table_properties_collector_factory_t* f); + +/* Get Table Properties */ + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_get_propeties_of_all_tables(crocksdb_t* db, + crocksdb_table_properties_collection_t* props, char** errptr); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_get_propeties_of_all_tables_cf( + crocksdb_t* db, crocksdb_column_family_handle_t* cf, + crocksdb_table_properties_collection_t* props, char** errptr); + +extern C_ROCKSDB_LIBRARY_API void +crocksdb_get_propeties_of_tables_in_range( + crocksdb_t* db, crocksdb_column_family_handle_t* cf, + int num_ranges, + const char* const* start_keys, const size_t* start_keys_lens, + const char* const* limit_keys, const size_t* limit_keys_lens, + crocksdb_table_properties_collection_t* props, char** errptr); + #ifdef __cplusplus } /* end extern "C" */ #endif diff --git a/librocksdb_sys/src/lib.rs b/librocksdb_sys/src/lib.rs index 28c6e678e..1d3f39cae 100644 --- a/librocksdb_sys/src/lib.rs +++ b/librocksdb_sys/src/lib.rs @@ -17,7 +17,7 @@ extern crate libc; #[cfg(test)] extern crate tempdir; -use libc::{c_char, c_uchar, c_int, c_void, size_t, uint64_t, c_double}; +use libc::{c_char, c_uchar, c_int, c_void, size_t, uint8_t, uint32_t, uint64_t, c_double}; use std::ffi::CStr; pub enum DBOptions {} @@ -45,6 +45,13 @@ pub enum DBRateLimiter {} pub enum DBLogger {} pub enum DBCompactOptions {} pub enum DBPinnableSlice {} +pub enum DBUserCollectedProperties {} +pub enum DBUserCollectedPropertiesIterator {} +pub enum DBTableProperties {} +pub enum DBTablePropertiesCollection {} +pub enum DBTablePropertiesCollectionIterator {} +pub enum DBTablePropertiesCollector {} +pub enum DBTablePropertiesCollectorFactory {} pub fn new_bloom_filter(bits: c_int) -> *mut DBFilterPolicy { unsafe { crocksdb_filterpolicy_create_bloom(bits) } @@ -54,6 +61,16 @@ pub fn new_cache(capacity: size_t) -> *mut DBCache { unsafe { crocksdb_cache_create_lru(capacity) } } +#[repr(C)] +#[derive(PartialEq, Eq, Clone, Copy)] +pub enum DBEntryType { + Put = 0, + Delete = 1, + SingleDelete = 2, + Merge = 3, + Other = 4, +} + #[derive(Copy, Clone)] #[repr(C)] pub enum DBCompressionType { @@ -170,6 +187,28 @@ pub enum DBInfoLogLevel { DBNumInfoLog = 6, } +#[derive(Copy, Clone, Eq, PartialEq)] +#[repr(C)] +pub enum DBTableProperty { + DataSize = 1, + IndexSize = 2, + FilterSize = 3, + RawKeySize = 4, + RawValueSize = 5, + NumDataBlocks = 6, + NumEntries = 7, + FormatVersion = 8, + FixedKeyLen = 9, + ColumnFamilyId = 10, + ColumnFamilyName = 11, + FilterPolicyName = 12, + ComparatorName = 13, + MergeOperatorName = 14, + PrefixExtractorName = 15, + PropertyCollectorsNames = 16, + CompressionName = 17, +} + pub fn error_message(ptr: *mut c_char) -> String { let c_str = unsafe { CStr::from_ptr(ptr) }; let s = format!("{}", c_str.to_string_lossy()); @@ -599,6 +638,7 @@ extern "C" { pub fn crocksdb_drop_column_family(db: *mut DBInstance, column_family_handle: *mut DBCFHandle, err: *mut *mut c_char); + pub fn crocksdb_column_family_handle_id(column_family_handle: *mut DBCFHandle) -> u32; pub fn crocksdb_column_family_handle_destroy(column_family_handle: *mut DBCFHandle); pub fn crocksdb_list_column_families(db: *const DBOptions, path: *const c_char, @@ -808,6 +848,124 @@ extern "C" { valLen: *mut size_t) -> *const u8; pub fn crocksdb_pinnableslice_destroy(v: *mut DBPinnableSlice); + + pub fn crocksdb_user_collected_properties_add(props: *mut DBUserCollectedProperties, + key: *const uint8_t, + key_len: size_t, + value: *const uint8_t, + value_len: size_t); + + pub fn crocksdb_user_collected_properties_iter_create + (props: *mut DBUserCollectedProperties) + -> *mut DBUserCollectedPropertiesIterator; + + pub fn crocksdb_user_collected_properties_iter_destroy + (it: *mut DBUserCollectedPropertiesIterator); + + pub fn crocksdb_user_collected_properties_iter_valid + (it: *mut DBUserCollectedPropertiesIterator) -> bool; + + pub fn crocksdb_user_collected_properties_iter_next + (it: *mut DBUserCollectedPropertiesIterator); + + pub fn crocksdb_user_collected_properties_iter_key(it: *mut DBUserCollectedPropertiesIterator, + klen: *mut size_t) + -> *const uint8_t; + + pub fn crocksdb_user_collected_properties_iter_value + (it: *mut DBUserCollectedPropertiesIterator, vlen: *mut size_t) -> *const uint8_t; + + pub fn crocksdb_table_properties_create() -> *mut DBTableProperties; + + pub fn crocksdb_table_properties_destroy(props: *mut DBTableProperties); + + pub fn crocksdb_table_properties_get_u64(props: *mut DBTableProperties, + prop: DBTableProperty) + -> uint64_t; + + pub fn crocksdb_table_properties_get_str(props: *mut DBTableProperties, + prop: DBTableProperty, + slen: *mut size_t) + -> *const uint8_t; + + pub fn crocksdb_table_properties_get_user_properties(props: *mut DBTableProperties) + -> *mut DBUserCollectedProperties; + + pub fn crocksdb_table_properties_collection_create() -> *mut DBTablePropertiesCollection; + + pub fn crocksdb_table_properties_collection_destroy(props: *mut DBTablePropertiesCollection); + + pub fn crocksdb_table_properties_collection_iter_create + (props: *mut DBTablePropertiesCollection) + -> *mut DBTablePropertiesCollectionIterator; + + pub fn crocksdb_table_properties_collection_iter_destroy + (it: *mut DBTablePropertiesCollectionIterator); + + pub fn crocksdb_table_properties_collection_iter_valid + (it: *mut DBTablePropertiesCollectionIterator) -> bool; + + pub fn crocksdb_table_properties_collection_iter_next + (it: *mut DBTablePropertiesCollectionIterator); + + pub fn crocksdb_table_properties_collection_iter_key( + it: *mut DBTablePropertiesCollectionIterator, klen: *mut size_t) -> *const uint8_t; + + pub fn crocksdb_table_properties_collection_iter_value + (it: *mut DBTablePropertiesCollectionIterator, value: *mut DBTableProperties); + + pub fn crocksdb_table_properties_collector_create(state: *mut c_void, + name: extern "C" fn(*mut c_void) + -> *const c_char, + destruct: extern "C" fn(*mut c_void), + add_userkey: extern "C" fn(*mut c_void, + *const uint8_t, + size_t, + *const uint8_t, + size_t, + c_int, + uint64_t, + uint64_t), + finish: extern "C" fn( + *mut c_void, + *mut DBUserCollectedProperties)) + -> *mut DBTablePropertiesCollector; + + pub fn crocksdb_table_properties_collector_destroy(c: *mut DBTablePropertiesCollector); + + pub fn crocksdb_table_properties_collector_factory_create + (state: *mut c_void, + name: extern "C" fn(*mut c_void) -> *const c_char, + destruct: extern "C" fn(*mut c_void), + create_table_properties_collector: extern "C" fn(*mut c_void, uint32_t) + -> *mut DBTablePropertiesCollector) + -> *mut DBTablePropertiesCollectorFactory; + + pub fn crocksdb_table_properties_collector_factory_destroy( + f: *mut DBTablePropertiesCollectorFactory); + + pub fn crocksdb_options_add_table_properties_collector_factory( + options: *mut DBOptions, f: *mut DBTablePropertiesCollectorFactory); + + pub fn crocksdb_get_properties_of_all_tables(db: *mut DBInstance, + props: *mut DBTablePropertiesCollection, + errptr: *mut *mut c_char); + + pub fn crocksdb_get_properties_of_all_tables_cf(db: *mut DBInstance, + cf: *mut DBCFHandle, + props: *mut DBTablePropertiesCollection, + errptr: *mut *mut c_char); + + pub fn crocksdb_get_properties_of_tables_in_range(db: *mut DBInstance, + cf: *mut DBCFHandle, + num_ranges: c_int, + start_keys: *const *const uint8_t, + start_keys_lens: *const size_t, + limit_keys: *const *const uint8_t, + limit_keys_lens: *const size_t, + props: *mut DBTablePropertiesCollection, + errptr: *mut *mut c_char); + } #[cfg(test)] diff --git a/src/lib.rs b/src/lib.rs index 2d41db60d..7e4d0b787 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,14 +26,20 @@ pub mod merge_operator; pub mod comparator; mod compaction_filter; mod slice_transform; +mod table_properties; +mod table_properties_collector; +mod table_properties_collector_factory; pub use compaction_filter::CompactionFilter; pub use librocksdb_sys::{DBCompactionStyle, DBCompressionType, DBRecoveryMode, DBInfoLogLevel, DBStatisticsTickerType, DBStatisticsHistogramType, new_bloom_filter, - CompactionPriority, self as crocksdb_ffi}; + CompactionPriority, DBEntryType, self as crocksdb_ffi}; pub use merge_operator::MergeOperands; pub use rocksdb::{DB, DBIterator, DBVector, Kv, SeekKey, Writable, WriteBatch, CFHandle, Range, BackupEngine, SstFileWriter}; pub use rocksdb_options::{BlockBasedOptions, Options, ReadOptions, WriteOptions, RestoreOptions, IngestExternalFileOptions, EnvOptions, HistogramData, CompactOptions}; pub use slice_transform::SliceTransform; +pub use table_properties::{TableProperties, TablePropertiesCollection}; +pub use table_properties_collector::TablePropertiesCollector; +pub use table_properties_collector_factory::TablePropertiesCollectorFactory; diff --git a/src/rocksdb.rs b/src/rocksdb.rs index bc398cf6d..70b0daa0a 100644 --- a/src/rocksdb.rs +++ b/src/rocksdb.rs @@ -26,6 +26,7 @@ use std::fmt::{self, Debug, Formatter}; use std::ops::Deref; use std::path::Path; use std::str::from_utf8; +use table_properties::{TablePropertiesCollection, new_table_properties_collection}; const DEFAULT_COLUMN_FAMILY: &'static str = "default"; @@ -33,6 +34,12 @@ pub struct CFHandle { inner: *mut DBCFHandle, } +impl CFHandle { + pub fn id(&self) -> u32 { + unsafe { crocksdb_ffi::crocksdb_column_family_handle_id(self.inner) } + } +} + impl Drop for CFHandle { fn drop(&mut self) { unsafe { @@ -1037,6 +1044,46 @@ impl DB { pub fn get_block_cache_usage_cf(&self, cf: &CFHandle) -> u64 { self.get_options_cf(cf).get_block_cache_usage() } + + pub fn get_properties_of_all_tables(&self) -> Result { + unsafe { + let props = new_table_properties_collection(); + ffi_try!(crocksdb_get_properties_of_all_tables(self.inner, props.inner)); + Ok(props) + } + } + + pub fn get_properties_of_all_tables_cf(&self, + cf: &CFHandle) + -> Result { + unsafe { + let props = new_table_properties_collection(); + ffi_try!(crocksdb_get_properties_of_all_tables_cf(self.inner, cf.inner, props.inner)); + Ok(props) + } + } + + pub fn get_properties_of_tables_in_range(&self, + cf: &CFHandle, + ranges: &[Range]) + -> Result { + let start_keys: Vec<*const u8> = ranges.iter().map(|x| x.start_key.as_ptr()).collect(); + let start_keys_lens: Vec<_> = ranges.iter().map(|x| x.start_key.len()).collect(); + let limit_keys: Vec<*const u8> = ranges.iter().map(|x| x.end_key.as_ptr()).collect(); + let limit_keys_lens: Vec<_> = ranges.iter().map(|x| x.end_key.len()).collect(); + unsafe { + let props = new_table_properties_collection(); + ffi_try!(crocksdb_get_properties_of_tables_in_range(self.inner, + cf.inner, + ranges.len() as i32, + start_keys.as_ptr(), + start_keys_lens.as_ptr(), + limit_keys.as_ptr(), + limit_keys_lens.as_ptr(), + props.inner)); + Ok(props) + } + } } impl Writable for DB { diff --git a/src/rocksdb_options.rs b/src/rocksdb_options.rs index ab75b8cf4..7672cb274 100644 --- a/src/rocksdb_options.rs +++ b/src/rocksdb_options.rs @@ -26,6 +26,8 @@ use merge_operator::MergeFn; use slice_transform::{SliceTransform, new_slice_transform}; use std::ffi::{CStr, CString}; use std::mem; +use table_properties_collector_factory::{TablePropertiesCollectorFactory, + new_table_properties_collector_factory}; #[derive(Default, Debug)] pub struct HistogramData { @@ -378,6 +380,15 @@ impl Options { } } + pub fn add_table_properties_collector_factory(&mut self, + fname: &str, + factory: Box) { + unsafe { + let f = new_table_properties_collector_factory(fname, factory); + crocksdb_ffi::crocksdb_options_add_table_properties_collector_factory(self.inner, f); + } + } + pub fn create_if_missing(&mut self, create_if_missing: bool) { unsafe { crocksdb_ffi::crocksdb_options_set_create_if_missing(self.inner, create_if_missing); diff --git a/src/table_properties.rs b/src/table_properties.rs new file mode 100644 index 000000000..0c0354e02 --- /dev/null +++ b/src/table_properties.rs @@ -0,0 +1,270 @@ +// Copyright 2017 PingCAP, Inc. +// +// 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, +// See the License for the specific language governing permissions and +// limitations under the License. + +use crocksdb_ffi::{self, DBTableProperties, DBTableProperty, DBUserCollectedPropertiesIterator, + DBTablePropertiesCollection, DBTablePropertiesCollectionIterator}; +use libc::size_t; +use std::collections::HashMap; +use std::marker::PhantomData; +use std::slice; +use std::str; + +pub fn new_table_properties_collection() -> TablePropertiesCollection { + TablePropertiesCollection::new() +} + +pub struct TablePropertiesCollection { + pub inner: *mut DBTablePropertiesCollection, +} + +impl Drop for TablePropertiesCollection { + fn drop(&mut self) { + unsafe { + crocksdb_ffi::crocksdb_table_properties_collection_destroy(self.inner); + } + } +} + +impl TablePropertiesCollection { + fn new() -> TablePropertiesCollection { + unsafe { + TablePropertiesCollection { + inner: crocksdb_ffi::crocksdb_table_properties_collection_create(), + } + } + } + + pub fn collect(&self) -> HashMap<&str, TableProperties> { + let mut res = HashMap::new(); + let mut iter = TablePropertiesCollectionIter::new(self); + while iter.valid() { + res.insert(iter.key(), iter.value()); + iter.next(); + } + res + } +} + +pub struct TablePropertiesCollectionIter<'a> { + props: PhantomData<&'a TablePropertiesCollection>, + inner: *mut DBTablePropertiesCollectionIterator, +} + +impl<'a> Drop for TablePropertiesCollectionIter<'a> { + fn drop(&mut self) { + unsafe { + crocksdb_ffi::crocksdb_table_properties_collection_iter_destroy(self.inner); + } + } +} + +impl<'a> TablePropertiesCollectionIter<'a> { + fn new(props: &'a TablePropertiesCollection) -> TablePropertiesCollectionIter<'a> { + unsafe { + TablePropertiesCollectionIter { + props: PhantomData, + inner: crocksdb_ffi::crocksdb_table_properties_collection_iter_create(props.inner), + } + } + } + + pub fn valid(&self) -> bool { + unsafe { crocksdb_ffi::crocksdb_table_properties_collection_iter_valid(self.inner) } + } + + pub fn next(&mut self) { + unsafe { + crocksdb_ffi::crocksdb_table_properties_collection_iter_next(self.inner); + } + } + + pub fn key(&self) -> &'a str { + unsafe { + let mut klen: size_t = 0; + let k = crocksdb_ffi::crocksdb_table_properties_collection_iter_key(self.inner, + &mut klen); + let bytes = slice::from_raw_parts(k, klen); + str::from_utf8(bytes).unwrap() + } + } + + pub fn value(&self) -> TableProperties { + unsafe { + let props = TableProperties::new(); + crocksdb_ffi::crocksdb_table_properties_collection_iter_value(self.inner, props.inner); + props + } + } +} + +pub struct TableProperties { + inner: *mut DBTableProperties, +} + +impl Drop for TableProperties { + fn drop(&mut self) { + unsafe { + crocksdb_ffi::crocksdb_table_properties_destroy(self.inner); + } + } +} + +impl TableProperties { + fn new() -> TableProperties { + unsafe { TableProperties { inner: crocksdb_ffi::crocksdb_table_properties_create() } } + } + + fn get_u64(&self, prop: DBTableProperty) -> u64 { + unsafe { crocksdb_ffi::crocksdb_table_properties_get_u64(self.inner, prop) } + } + + fn get_str(&self, prop: DBTableProperty) -> &str { + unsafe { + let mut slen: size_t = 0; + let s = crocksdb_ffi::crocksdb_table_properties_get_str(self.inner, prop, &mut slen); + let bytes = slice::from_raw_parts(s, slen); + str::from_utf8(bytes).unwrap() + } + } + + pub fn data_size(&self) -> u64 { + self.get_u64(DBTableProperty::DataSize) + } + + pub fn index_size(&self) -> u64 { + self.get_u64(DBTableProperty::IndexSize) + } + + pub fn filter_size(&self) -> u64 { + self.get_u64(DBTableProperty::FilterSize) + } + + pub fn raw_key_size(&self) -> u64 { + self.get_u64(DBTableProperty::RawKeySize) + } + + pub fn raw_value_size(&self) -> u64 { + self.get_u64(DBTableProperty::RawValueSize) + } + + pub fn num_data_blocks(&self) -> u64 { + self.get_u64(DBTableProperty::NumDataBlocks) + } + + pub fn num_entries(&self) -> u64 { + self.get_u64(DBTableProperty::NumEntries) + } + + pub fn format_version(&self) -> u64 { + self.get_u64(DBTableProperty::FormatVersion) + } + + pub fn fixed_key_len(&self) -> u64 { + self.get_u64(DBTableProperty::FixedKeyLen) + } + + pub fn column_family_id(&self) -> u64 { + self.get_u64(DBTableProperty::ColumnFamilyId) + } + + pub fn column_family_name(&self) -> &str { + self.get_str(DBTableProperty::ColumnFamilyName) + } + + pub fn filter_policy_name(&self) -> &str { + self.get_str(DBTableProperty::FilterPolicyName) + } + + pub fn comparator_name(&self) -> &str { + self.get_str(DBTableProperty::ComparatorName) + } + + pub fn merge_operator_name(&self) -> &str { + self.get_str(DBTableProperty::MergeOperatorName) + } + + pub fn prefix_extractor_name(&self) -> &str { + self.get_str(DBTableProperty::PrefixExtractorName) + } + + pub fn property_collectors_names(&self) -> &str { + self.get_str(DBTableProperty::PropertyCollectorsNames) + } + + pub fn compression_name(&self) -> &str { + self.get_str(DBTableProperty::CompressionName) + } + + pub fn user_collected_properties(&self) -> HashMap<&[u8], &[u8]> { + let mut res = HashMap::new(); + let mut iter = UserCollectedPropertiesIter::new(self); + while iter.valid() { + res.insert(iter.key(), iter.value()); + iter.next(); + } + res + } +} + +struct UserCollectedPropertiesIter<'a> { + props: PhantomData<&'a TableProperties>, + inner: *mut DBUserCollectedPropertiesIterator, +} + +impl<'a> Drop for UserCollectedPropertiesIter<'a> { + fn drop(&mut self) { + unsafe { + crocksdb_ffi::crocksdb_user_collected_properties_iter_destroy(self.inner); + } + } +} + +impl<'a> UserCollectedPropertiesIter<'a> { + fn new(props: &'a TableProperties) -> UserCollectedPropertiesIter<'a> { + unsafe { + let inner = crocksdb_ffi::crocksdb_table_properties_get_user_properties(props.inner); + UserCollectedPropertiesIter { + props: PhantomData, + inner: crocksdb_ffi::crocksdb_user_collected_properties_iter_create(inner), + } + } + } + + fn valid(&self) -> bool { + unsafe { crocksdb_ffi::crocksdb_user_collected_properties_iter_valid(self.inner) } + } + + fn next(&mut self) { + unsafe { + crocksdb_ffi::crocksdb_user_collected_properties_iter_next(self.inner); + } + } + + fn key(&self) -> &'a [u8] { + unsafe { + let mut klen: size_t = 0; + let k = crocksdb_ffi::crocksdb_user_collected_properties_iter_key(self.inner, + &mut klen); + slice::from_raw_parts(k, klen) + } + } + + fn value(&self) -> &'a [u8] { + unsafe { + let mut vlen: size_t = 0; + let v = crocksdb_ffi::crocksdb_user_collected_properties_iter_value(self.inner, + &mut vlen); + slice::from_raw_parts(v, vlen) + } + } +} diff --git a/src/table_properties_collector.rs b/src/table_properties_collector.rs new file mode 100644 index 000000000..36cbc20d7 --- /dev/null +++ b/src/table_properties_collector.rs @@ -0,0 +1,108 @@ +// Copyright 2017 PingCAP, Inc. +// +// 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, +// See the License for the specific language governing permissions and +// limitations under the License. + +use crocksdb_ffi::{self, DBEntryType, DBUserCollectedProperties, DBTablePropertiesCollector}; +use libc::{c_void, c_char, c_int, uint8_t, uint64_t, size_t}; +use std::collections::HashMap; +use std::ffi::CString; +use std::mem; +use std::slice; + +/// `TablePropertiesCollector` provides the mechanism for users to collect +/// their own properties that they are interested in. This class is essentially +/// a collection of callback functions that will be invoked during table +/// building. It is constructed with TablePropertiesCollectorFactory. The methods +/// don't need to be thread-safe, as we will create exactly one +/// TablePropertiesCollector object per table and then call it sequentially +pub trait TablePropertiesCollector { + /// Will be called when a new key/value pair is inserted into the table. + fn add(&mut self, + key: &[u8], + value: &[u8], + entry_type: DBEntryType, + seq: u64, + file_size: u64); + + /// Will be called when a table has already been built and is ready for + /// writing the properties block. + fn finish(&mut self) -> HashMap, Vec>; +} + +struct TablePropertiesCollectorHandle { + name: CString, + rep: Box, +} + +impl TablePropertiesCollectorHandle { + fn new(name: &str, rep: Box) -> TablePropertiesCollectorHandle { + TablePropertiesCollectorHandle { + name: CString::new(name).unwrap(), + rep: rep, + } + } +} + +extern "C" fn name(handle: *mut c_void) -> *const c_char { + unsafe { + let handle = &mut *(handle as *mut TablePropertiesCollectorHandle); + handle.name.as_ptr() + } +} + +extern "C" fn destruct(handle: *mut c_void) { + unsafe { + Box::from_raw(handle as *mut TablePropertiesCollectorHandle); + } +} + +pub extern "C" fn add(handle: *mut c_void, + key: *const uint8_t, + key_len: size_t, + value: *const uint8_t, + value_len: size_t, + entry_type: c_int, + seq: uint64_t, + file_size: uint64_t) { + unsafe { + let handle = &mut *(handle as *mut TablePropertiesCollectorHandle); + let key = slice::from_raw_parts(key, key_len); + let value = slice::from_raw_parts(value, value_len); + handle.rep.add(key, value, mem::transmute(entry_type), seq, file_size); + } +} + +pub extern "C" fn finish(handle: *mut c_void, props: *mut DBUserCollectedProperties) { + unsafe { + let handle = &mut *(handle as *mut TablePropertiesCollectorHandle); + for (key, value) in handle.rep.finish() { + crocksdb_ffi::crocksdb_user_collected_properties_add(props, + key.as_ptr(), + key.len(), + value.as_ptr(), + value.len()); + } + } +} + +pub unsafe fn new_table_properties_collector(cname: &str, + collector: Box) + -> *mut DBTablePropertiesCollector { + let handle = TablePropertiesCollectorHandle::new(cname, collector); + crocksdb_ffi::crocksdb_table_properties_collector_create( + Box::into_raw(Box::new(handle)) as *mut c_void, + name, + destruct, + add, + finish, + ) +} diff --git a/src/table_properties_collector_factory.rs b/src/table_properties_collector_factory.rs new file mode 100644 index 000000000..527e44da5 --- /dev/null +++ b/src/table_properties_collector_factory.rs @@ -0,0 +1,75 @@ +// Copyright 2017 PingCAP, Inc. +// +// 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, +// See the License for the specific language governing permissions and +// limitations under the License. + +use crocksdb_ffi::{self, DBTablePropertiesCollector, DBTablePropertiesCollectorFactory}; +use libc::{c_void, c_char, uint32_t}; +use std::ffi::CString; +use table_properties_collector::{TablePropertiesCollector, new_table_properties_collector}; + +/// Constructs `TablePropertiesCollector`. +/// Internals create a new `TablePropertiesCollector` for each new table. +pub trait TablePropertiesCollectorFactory { + /// Has to be thread-safe. + fn create_table_properties_collector(&mut self, cf: u32) -> Box; +} + +struct TablePropertiesCollectorFactoryHandle { + name: CString, + rep: Box, +} + +impl TablePropertiesCollectorFactoryHandle { + fn new(name: &str, + rep: Box) + -> TablePropertiesCollectorFactoryHandle { + TablePropertiesCollectorFactoryHandle { + name: CString::new(name).unwrap(), + rep: rep, + } + } +} + +extern "C" fn name(handle: *mut c_void) -> *const c_char { + unsafe { + let handle = &mut *(handle as *mut TablePropertiesCollectorFactoryHandle); + handle.name.as_ptr() + } +} + +extern "C" fn destruct(handle: *mut c_void) { + unsafe { + Box::from_raw(handle as *mut TablePropertiesCollectorFactoryHandle); + } +} + +extern "C" fn create_table_properties_collector(handle: *mut c_void, + cf: uint32_t) + -> *mut DBTablePropertiesCollector { + unsafe { + let handle = &mut *(handle as *mut TablePropertiesCollectorFactoryHandle); + let collector = handle.rep.create_table_properties_collector(cf); + new_table_properties_collector(handle.name.to_str().unwrap(), collector) + } +} + +pub unsafe fn new_table_properties_collector_factory + (fname: &str, factory: Box) + -> *mut DBTablePropertiesCollectorFactory { + let handle = TablePropertiesCollectorFactoryHandle::new(fname, factory); + crocksdb_ffi::crocksdb_table_properties_collector_factory_create( + Box::into_raw(Box::new(handle)) as *mut c_void, + name, + destruct, + create_table_properties_collector, + ) +} diff --git a/tests/test.rs b/tests/test.rs index 1e84d1abe..54f50d1ee 100644 --- a/tests/test.rs +++ b/tests/test.rs @@ -1,5 +1,6 @@ extern crate rocksdb; extern crate tempdir; +extern crate byteorder; mod test_iterator; mod test_multithreaded; @@ -11,3 +12,4 @@ mod test_ingest_external_file; mod test_slice_transform; mod test_prefix_extractor; mod test_statistics; +mod test_table_properties; diff --git a/tests/test_table_properties.rs b/tests/test_table_properties.rs new file mode 100644 index 000000000..8a3de3f06 --- /dev/null +++ b/tests/test_table_properties.rs @@ -0,0 +1,191 @@ +// Copyright 2017 PingCAP, Inc. +// +// 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, +// See the License for the specific language governing permissions and +// limitations under the License. + +use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt}; +use rocksdb::{DB, Range, Options, Writable, DBEntryType, TablePropertiesCollection, + TablePropertiesCollector, TablePropertiesCollectorFactory}; +use std::collections::HashMap; +use std::fmt; +use tempdir::TempDir; + +enum Props { + NumKeys = 0, + NumPuts, + NumMerges, + NumDeletes, +} + +fn encode_u32(x: u32) -> Vec { + let mut w = Vec::new(); + w.write_u32::(x).unwrap(); + w +} + +fn decode_u32(mut x: &[u8]) -> u32 { + x.read_u32::().unwrap() +} + +struct ExampleCollector { + num_keys: u32, + num_puts: u32, + num_merges: u32, + num_deletes: u32, + last_key: Vec, +} + +impl ExampleCollector { + fn new() -> ExampleCollector { + ExampleCollector { + num_keys: 0, + num_puts: 0, + num_merges: 0, + num_deletes: 0, + last_key: Vec::new(), + } + } + + fn add(&mut self, other: &ExampleCollector) { + self.num_keys += other.num_keys; + self.num_puts += other.num_puts; + self.num_merges += other.num_merges; + self.num_deletes += other.num_deletes; + } + + fn encode(&self) -> HashMap, Vec> { + let mut props = HashMap::new(); + props.insert(vec![Props::NumKeys as u8], encode_u32(self.num_keys)); + props.insert(vec![Props::NumPuts as u8], encode_u32(self.num_puts)); + props.insert(vec![Props::NumMerges as u8], encode_u32(self.num_merges)); + props.insert(vec![Props::NumDeletes as u8], encode_u32(self.num_deletes)); + props + } + + fn decode(props: HashMap<&[u8], &[u8]>) -> ExampleCollector { + let mut c = ExampleCollector::new(); + c.num_keys = decode_u32(props.get(&[Props::NumKeys as u8].as_ref()).unwrap()); + c.num_puts = decode_u32(props.get(&[Props::NumPuts as u8].as_ref()).unwrap()); + c.num_merges = decode_u32(props.get(&[Props::NumMerges as u8].as_ref()).unwrap()); + c.num_deletes = decode_u32(props.get(&[Props::NumDeletes as u8].as_ref()).unwrap()); + c + } +} + +impl fmt::Display for ExampleCollector { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, + "keys={}, puts={}, merges={}, deletes={}", + self.num_keys, + self.num_puts, + self.num_merges, + self.num_deletes) + } +} + +impl TablePropertiesCollector for ExampleCollector { + fn add(&mut self, key: &[u8], _: &[u8], entry_type: DBEntryType, _: u64, _: u64) { + if key != self.last_key.as_slice() { + self.num_keys += 1; + self.last_key.clear(); + self.last_key.extend_from_slice(key); + } + match entry_type { + DBEntryType::Put => self.num_puts += 1, + DBEntryType::Merge => self.num_merges += 1, + DBEntryType::Delete | + DBEntryType::SingleDelete => self.num_deletes += 1, + DBEntryType::Other => {} + } + } + + fn finish(&mut self) -> HashMap, Vec> { + self.encode() + } +} + +struct ExampleFactory {} + +impl ExampleFactory { + fn new() -> ExampleFactory { + ExampleFactory {} + } +} + +impl TablePropertiesCollectorFactory for ExampleFactory { + fn create_table_properties_collector(&mut self, _: u32) -> Box { + Box::new(ExampleCollector::new()) + } +} + +fn check_collection(collection: &TablePropertiesCollection, + num_files: usize, + num_keys: u32, + num_puts: u32, + num_merges: u32, + num_deletes: u32) { + let mut res = ExampleCollector::new(); + let props = collection.collect(); + for (k, v) in &props { + assert!(k.ends_with(".sst")); + assert_eq!(v.property_collectors_names(), "[example-collector]"); + res.add(&ExampleCollector::decode(v.user_collected_properties())); + } + assert_eq!(props.len(), num_files); + assert_eq!(res.num_keys, num_keys); + assert_eq!(res.num_puts, num_puts); + assert_eq!(res.num_merges, num_merges); + assert_eq!(res.num_deletes, num_deletes); +} + +#[test] +fn test_table_properties_collector_factory() { + let f = ExampleFactory::new(); + let mut opts = Options::new(); + opts.create_if_missing(true); + opts.add_table_properties_collector_factory("example-collector", Box::new(f)); + + let path = TempDir::new("_rust_rocksdb_collectortest").expect(""); + let db = DB::open(opts, path.path().to_str().unwrap()).unwrap(); + + let samples = vec![(b"key1".to_vec(), b"value1".to_vec()), + (b"key2".to_vec(), b"value2".to_vec()), + (b"key3".to_vec(), b"value3".to_vec()), + (b"key4".to_vec(), b"value4".to_vec())]; + + // Put 4 keys. + for &(ref k, ref v) in &samples { + db.put(k, v).unwrap(); + assert_eq!(v.as_slice(), &*db.get(k).unwrap().unwrap()); + } + db.flush(true).unwrap(); + let collection = db.get_properties_of_all_tables().unwrap(); + check_collection(&collection, 1, 4, 4, 0, 0); + + // Delete 2 keys. + let cf = db.cf_handle("default").unwrap(); + for &(ref k, _) in &samples[0..2] { + db.delete_cf(cf, k).unwrap(); + } + db.flush_cf(cf, true).unwrap(); + let collection = db.get_properties_of_all_tables_cf(cf).unwrap(); + check_collection(&collection, 2, 6, 4, 0, 2); + + // ["key2", "key3") covers two sst files. + let range = Range::new(b"key2", b"key3"); + let collection = db.get_properties_of_tables_in_range(cf, &[range]).unwrap(); + check_collection(&collection, 2, 6, 4, 0, 2); + + // ["key3", "key4") covers only the first sst file. + let range = Range::new(b"key3", b"key4"); + let collection = db.get_properties_of_tables_in_range(cf, &[range]).unwrap(); + check_collection(&collection, 1, 4, 4, 0, 0); +}