Update comparator api (#67)

comparator.go

@@ -1,45 +1,39 @@
 package grocksdb
 
 // #include "rocksdb/c.h"
+// #include "grocksdb.h"
 import "C"
+
 import (
 	"bytes"
 )
 
-// A Comparator object provides a total order across slices that are
-// used as keys in an sstable or a database.
-type Comparator interface {
-	// Three-way comparison. Returns value:
-	//   < 0 iff "a" < "b",
-	//   == 0 iff "a" == "b",
-	//   > 0 iff "a" > "b"
-	Compare(a, b []byte) int
-
-	// The name of the comparator.
-	Name() string
-
-	// Return native comparator.
-	Native() *C.rocksdb_comparator_t
+// Comparing functor.
+//
+// Three-way comparison. Returns value:
+//   < 0 iff "a" < "b",
+//   == 0 iff "a" == "b",
+//   > 0 iff "a" > "b"
+type Comparing = func(a, b []byte) int
 
-	// Destroy comparator.
-	Destroy()
+// NewComparator creates a Comparator object which contains native c-comparator pointer.
+func NewComparator(name string, compare Comparing) *Comparator {
+	cmp := &Comparator{name: name, compare: compare}
+	idx := registerComperator(cmp)
+	cmp.c = C.gorocksdb_comparator_create(C.uintptr_t(idx))
+	return cmp
 }
 
-// NewNativeComparator creates a Comparator object.
-func NewNativeComparator(c *C.rocksdb_comparator_t) Comparator {
-	return &nativeComparator{c}
+// NativeComparator wraps c-comparator pointer.
+type Comparator struct {
+	c       *C.rocksdb_comparator_t
+	compare Comparing
+	name    string
 }
 
-type nativeComparator struct {
-	c *C.rocksdb_comparator_t
-}
-
-func (c *nativeComparator) Compare(a, b []byte) int { return 0 }
-func (c *nativeComparator) Name() string            { return "" }
-func (c *nativeComparator) Native() *C.rocksdb_comparator_t {
-	return c.c
-}
-func (c *nativeComparator) Destroy() {
+func (c *Comparator) Compare(a, b []byte) int { return c.compare(a, b) }
+func (c *Comparator) Name() string            { return c.name }
+func (c *Comparator) Destroy() {
 	C.rocksdb_comparator_destroy(c.c)
 	c.c = nil
 }
@@ -49,10 +43,10 @@ var comperators = NewCOWList()
 
 type comperatorWrapper struct {
 	name       *C.char
-	comparator Comparator
+	comparator *Comparator
 }
 
-func registerComperator(cmp Comparator) int {
+func registerComperator(cmp *Comparator) int {
 	return comperators.Append(comperatorWrapper{C.CString(cmp.Name()), cmp})
 }
 
@@ -75,5 +69,4 @@ func (cmp *testBytesReverseComparator) Name() string { return "grocksdb.bytes-re
 func (cmp *testBytesReverseComparator) Compare(a, b []byte) int {
 	return bytes.Compare(a, b) * -1
 }
-func (cmp *testBytesReverseComparator) Native() *C.rocksdb_comparator_t { return nil }
-func (cmp *testBytesReverseComparator) Destroy()                        {}
+func (cmp *testBytesReverseComparator) Destroy() {}

comparator_test.go

@@ -1,16 +1,22 @@
 package grocksdb
 
 import (
+	"bytes"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
 func TestComparator(t *testing.T) {
-	db := newTestDB(t, "TestComparator", func(opts *Options) {
-		opts.SetComparator(&testBytesReverseComparator{})
+	db, opts := newTestDBAndOpts(t, "TestComparator", func(opts *Options) {
+		opts.SetComparator(NewComparator("rev", func(a, b []byte) int {
+			return bytes.Compare(a, b) * -1
+		}))
 	})
-	defer db.Close()
+	defer func() {
+		db.Close()
+		opts.Destroy()
+	}()
 
 	// insert keys
 	givenKeys := [][]byte{[]byte("key1"), []byte("key2"), []byte("key3")}

db_test.go

@@ -166,6 +166,25 @@ func newTestDB(t *testing.T, name string, applyOpts func(opts *Options)) *DB {
 	return db
 }
 
+func newTestDBAndOpts(t *testing.T, name string, applyOpts func(opts *Options)) (*DB, *Options) {
+	dir, err := ioutil.TempDir("", "gorocksdb-"+name)
+	require.Nil(t, err)
+
+	opts := NewDefaultOptions()
+	// test the ratelimiter
+	rateLimiter := NewRateLimiter(1024, 100*1000, 10)
+	opts.SetRateLimiter(rateLimiter)
+	opts.SetCreateIfMissing(true)
+	opts.SetCompression(ZSTDCompression)
+	if applyOpts != nil {
+		applyOpts(opts)
+	}
+	db, err := OpenDb(opts, dir)
+	require.Nil(t, err)
+
+	return db, opts
+}
+
 func newTestDBMultiCF(t *testing.T, name string, columns []string, applyOpts func(opts *Options)) (db *DB, cfh []*ColumnFamilyHandle, cleanup func()) {
 	dir, err := ioutil.TempDir("", "gorocksdb-"+name)
 	require.Nil(t, err)

options.go

@@ -166,18 +166,22 @@ func (opts *Options) SetCompactionFilter(value CompactionFilter) {
 }
 
 // SetComparator sets the comparator which define the order of keys in the table.
+// This operation is `move`, thus underlying native c-pointer is owned by Options.
+// `cmp` is no longer usable.
 //
 // Default: a comparator that uses lexicographic byte-wise ordering
-func (opts *Options) SetComparator(value Comparator) {
-	C.rocksdb_comparator_destroy(opts.ccmp)
-
-	if nc, ok := value.(*nativeComparator); ok {
-		opts.ccmp = nc.c
-	} else {
-		idx := registerComperator(value)
-		opts.ccmp = C.gorocksdb_comparator_create(C.uintptr_t(idx))
-	}
+func (opts *Options) SetComparator(cmp *Comparator) {
+	cmp_ := unsafe.Pointer(cmp.c)
+	opts.SetNativeComparator(cmp_)
+	cmp.c = nil
+}
 
+// SetNativeComparator sets the comparator which define the order of keys in the table.
+//
+// Default: a comparator that uses lexicographic byte-wise ordering
+func (opts *Options) SetNativeComparator(cmp unsafe.Pointer) {
+	C.rocksdb_comparator_destroy(opts.ccmp)
+	opts.ccmp = (*C.rocksdb_comparator_t)(cmp)
 	C.rocksdb_options_set_comparator(opts.c, opts.ccmp)
 }
 

options_test.go

@@ -350,9 +350,6 @@ func TestOptions(t *testing.T) {
 	// set compaction filter
 	opts.SetCompactionFilter(NewNativeCompactionFilter(nil))
 
-	// set comparator
-	opts.SetComparator(NewNativeComparator(nil))
-
 	// set merge operator
 	opts.SetMergeOperator(NewNativeMergeOperator(nil))
 

sst_file_writer.go

@@ -21,8 +21,15 @@ func NewSSTFileWriter(opts *EnvOptions, dbOpts *Options) *SSTFileWriter {
 }
 
 // NewSSTFileWriterWithComparator creates an SSTFileWriter object with comparator.
-func NewSSTFileWriterWithComparator(opts *EnvOptions, dbOpts *Options, cmp Comparator) *SSTFileWriter {
-	c := C.rocksdb_sstfilewriter_create_with_comparator(opts.c, dbOpts.c, cmp.Native())
+func NewSSTFileWriterWithComparator(opts *EnvOptions, dbOpts *Options, cmp *Comparator) *SSTFileWriter {
+	cmp_ := unsafe.Pointer(cmp.c)
+	return NewSSTFileWriterWithNativeComparator(opts, dbOpts, cmp_)
+}
+
+// NewSSTFileWriterWithNativeComparator creates an SSTFileWriter object with native comparator.
+func NewSSTFileWriterWithNativeComparator(opts *EnvOptions, dbOpts *Options, cmp unsafe.Pointer) *SSTFileWriter {
+	cmp_ := (*C.rocksdb_comparator_t)(cmp)
+	c := C.rocksdb_sstfilewriter_create_with_comparator(opts.c, dbOpts.c, cmp_)
 	return &SSTFileWriter{c: c}
 }