diff --git a/comparator.go b/comparator.go index e17376e..66ba90b 100644 --- a/comparator.go +++ b/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() {} diff --git a/comparator_test.go b/comparator_test.go index a478237..679cbbe 100644 --- a/comparator_test.go +++ b/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")} diff --git a/db_test.go b/db_test.go index c86e301..38f9643 100755 --- a/db_test.go +++ b/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) diff --git a/options.go b/options.go index ecdf913..fd057e1 100644 --- a/options.go +++ b/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) } diff --git a/options_test.go b/options_test.go index 30a0313..3fc7a69 100644 --- a/options_test.go +++ b/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)) diff --git a/sst_file_writer.go b/sst_file_writer.go index 6d84e94..a96d03d 100644 --- a/sst_file_writer.go +++ b/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} }