Allow usage of GetApproximateMemoryUsageByType with TransactionDB

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 {

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)
+}