diff --git a/memory_usage.go b/memory_usage.go index 7b9a6ad..a2fad4e 100644 --- a/memory_usage.go +++ b/memory_usage.go @@ -20,16 +20,38 @@ type MemoryUsage struct { CacheTotal uint64 } +type NativeDB interface { + getNativeDB() *C.rocksdb_t +} + +func (db *DB) getNativeDB() *C.rocksdb_t { + return db.c +} + +func (db *TransactionDB) getNativeDB() *C.rocksdb_t { + return (*C.rocksdb_t)(db.c) +} + // GetApproximateMemoryUsageByType returns summary // memory usage stats for given databases and caches. func GetApproximateMemoryUsageByType(dbs []*DB, caches []*Cache) (*MemoryUsage, error) { + nativeDBs := make([]NativeDB, 0, len(dbs)) + for _, db := range dbs { + nativeDBs = append(nativeDBs, db) + } + return GetApproximateMemoryUsageByTypeNativeDB(nativeDBs, caches) +} + +// GetApproximateMemoryUsageByTypeNativeDB returns summary +// memory usage stats for given databases and caches. +func GetApproximateMemoryUsageByTypeNativeDB(dbs []NativeDB, caches []*Cache) (*MemoryUsage, error) { // register memory consumers consumers := C.rocksdb_memory_consumers_create() defer C.rocksdb_memory_consumers_destroy(consumers) for _, db := range dbs { if db != nil { - C.rocksdb_memory_consumers_add_db(consumers, db.c) + C.rocksdb_memory_consumers_add_db(consumers, (db.getNativeDB())) } } for _, cache := range caches { diff --git a/memory_usage_test.go b/memory_usage_test.go index 3f15be9..fd541d6 100644 --- a/memory_usage_test.go +++ b/memory_usage_test.go @@ -60,3 +60,55 @@ func TestMemoryUsage(t *testing.T) { assert.True(t, mu2.CacheTotal >= mu1.CacheTotal) assert.True(t, mu2.MemTableReadersTotal >= mu1.MemTableReadersTotal) } + +func TestMemoryUsageTransactionDB(t *testing.T) { + // create database with cache + cache := NewLRUCache(8 * 1024 * 1024) + bbto := NewDefaultBlockBasedTableOptions() + bbto.SetBlockCache(cache) + defer bbto.Destroy() + defer cache.Destroy() + + applyOpts := func(opts *Options, transactionDBOpts *TransactionDBOptions) { + opts.SetBlockBasedTableFactory(bbto) + } + + db := newTestTransactionDB(t, "TestMemoryUsage", applyOpts) + defer db.Close() + + // take first memory usage snapshot + mu1, err := GetApproximateMemoryUsageByTypeNativeDB([]NativeDB{db}, []*Cache{cache}) + ensure.Nil(t, err) + + // perforx`m IO operations that will affect in-memory tables (and maybe cache as well) + wo := NewDefaultWriteOptions() + defer wo.Destroy() + ro := NewDefaultReadOptions() + defer ro.Destroy() + + key := []byte("key") + value := make([]byte, 1024) + _, err = rand.Read(value) + ensure.Nil(t, err) + + err = db.Put(wo, key, value) + ensure.Nil(t, err) + + // A single Put is not enough to increase approximate memtable usage. + err = db.Put(wo, key, value) + ensure.Nil(t, err) + + _, err = db.Get(ro, key) + ensure.Nil(t, err) + + // take second memory usage snapshot + mu2, err := GetApproximateMemoryUsageByTypeNativeDB([]NativeDB{db}, []*Cache{cache}) + ensure.Nil(t, err) + + // the amount of memory used by memtables should increase after write/read; + // cache memory usage is not likely to be changed, perhaps because requested key is kept by memtable + assert.True(t, mu2.MemTableTotal > mu1.MemTableTotal) + assert.True(t, mu2.MemTableUnflushed > mu1.MemTableUnflushed) + assert.True(t, mu2.CacheTotal >= mu1.CacheTotal) + assert.True(t, mu2.MemTableReadersTotal >= mu1.MemTableReadersTotal) +}