Merge pull request tecbot#147 from vitalyisaev2/UCS-5061

Support MemoryUsage API

.gitignore

@@ -0,0 +1 @@
+.idea

.travis.yml

@@ -1,6 +1,6 @@
 language: go
 go:
-  - 1.7
+  - 1.11
   - tip
 
 before_install:
@@ -18,7 +18,7 @@ install:
   - git clone https://github.com/facebook/rocksdb.git /tmp/rocksdb
   - pushd /tmp/rocksdb
   - make clean
-  - make shared_lib
+  - make shared_lib -j`nproc`
   - sudo cp --preserve=links ./librocksdb.* /usr/lib/
   - sudo cp -r ./include/rocksdb/ /usr/include/
   - popd

README.md

@@ -4,7 +4,7 @@
 
 ## Install
 
-You'll need to build [RocksDB](https://github.com/facebook/rocksdb) v5.5+ on your machine.
+You'll need to build [RocksDB](https://github.com/facebook/rocksdb) v5.16+ on your machine.
 
 After that, you can install gorocksdb using the following command:
 

memory_usage.go

@@ -0,0 +1,58 @@
+package gorocksdb
+
+// #include <stdlib.h>
+// #include "rocksdb/c.h"
+import "C"
+import (
+	"errors"
+	"unsafe"
+)
+
+// MemoryUsage contains memory usage statistics provided by RocksDB
+type MemoryUsage struct {
+	// MemTableTotal estimates memory usage of all mem-tables
+	MemTableTotal uint64
+	// MemTableUnflushed estimates memory usage of unflushed mem-tables
+	MemTableUnflushed uint64
+	// MemTableReadersTotal memory usage of table readers (indexes and bloom filters)
+	MemTableReadersTotal uint64
+	// CacheTotal memory usage of cache
+	CacheTotal uint64
+}
+
+// GetApproximateMemoryUsageByType returns summary
+// memory usage stats for given databases and caches.
+func GetApproximateMemoryUsageByType(dbs []*DB, 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)
+		}
+	}
+	for _, cache := range caches {
+		if cache != nil {
+			C.rocksdb_memory_consumers_add_cache(consumers, cache.c)
+		}
+	}
+
+	// obtain memory usage stats
+	var cErr *C.char
+	memoryUsage := C.rocksdb_approximate_memory_usage_create(consumers, &cErr)
+	if cErr != nil {
+		defer C.free(unsafe.Pointer(cErr))
+		return nil, errors.New(C.GoString(cErr))
+	}
+
+	defer C.rocksdb_approximate_memory_usage_destroy(memoryUsage)
+
+	result := &MemoryUsage{
+		MemTableTotal:        uint64(C.rocksdb_approximate_memory_usage_get_mem_table_total(memoryUsage)),
+		MemTableUnflushed:    uint64(C.rocksdb_approximate_memory_usage_get_mem_table_unflushed(memoryUsage)),
+		MemTableReadersTotal: uint64(C.rocksdb_approximate_memory_usage_get_mem_table_readers_total(memoryUsage)),
+		CacheTotal:           uint64(C.rocksdb_approximate_memory_usage_get_cache_total(memoryUsage)),
+	}
+	return result, nil
+}

memory_usage_test.go

@@ -0,0 +1,56 @@
+package gorocksdb
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/facebookgo/ensure"
+)
+
+func TestMemoryUsage(t *testing.T) {
+	// create database with cache
+	cache := NewLRUCache(8 * 1024 * 1024)
+	bbto := NewDefaultBlockBasedTableOptions()
+	bbto.SetBlockCache(cache)
+	defer cache.Destroy()
+
+	applyOpts := func(opts *Options) {
+		opts.SetBlockBasedTableFactory(bbto)
+	}
+
+	db := newTestDB(t, "TestMemoryUsage", applyOpts)
+	defer db.Close()
+
+	// take first memory usage snapshot
+	mu1, err := GetApproximateMemoryUsageByType([]*DB{db}, []*Cache{cache})
+	ensure.Nil(t, err)
+
+	// perform 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)
+	_, err = db.Get(ro, key)
+	ensure.Nil(t, err)
+
+	// take second memory usage snapshot
+	mu2, err := GetApproximateMemoryUsageByType([]*DB{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)
+}