Change btree implementation for collections

internal/controller/controller.go

@@ -16,12 +16,12 @@ import (
 	"sync"
 	"time"
 
-	"github.com/tidwall/btree"
 	"github.com/tidwall/buntdb"
 	"github.com/tidwall/geojson"
 	"github.com/tidwall/resp"
 	"github.com/tidwall/tile38/core"
 	"github.com/tidwall/tile38/internal/collection"
+	"github.com/tidwall/tile38/internal/ds"
 	"github.com/tidwall/tile38/internal/endpoint"
 	"github.com/tidwall/tile38/internal/expire"
 	"github.com/tidwall/tile38/internal/log"
@@ -34,11 +34,6 @@ const goingLive = "going live"
 
 const hookLogPrefix = "hook:log:"
 
-type collectionT struct {
-	Key        string
-	Collection *collection.Collection
-}
-
 type commandDetailsT struct {
 	command   string
 	key, id   string
@@ -57,10 +52,6 @@ type commandDetailsT struct {
 	children []*commandDetailsT // for multi actions such as "PDEL"
 }
 
-func (col *collectionT) Less(item btree.Item, ctx interface{}) bool {
-	return col.Key < item.(*collectionT).Key
-}
-
 // Controller is a tile38 controller
 type Controller struct {
 	// static values
@@ -95,7 +86,7 @@ type Controller struct {
 	aofsz   int                             // active size of the aof file
 	qdb     *buntdb.DB                      // hook queue log
 	qidx    uint64                          // hook queue log last idx
-	cols    *btree.BTree                    // data collections
+	cols    ds.BTree                        // data collections
 	expires map[string]map[string]time.Time // synced with cols
 
 	follows    map[*bytes.Buffer]bool
@@ -136,7 +127,6 @@ func ListenAndServeEx(host string, port int, dir string, ln *net.Listener, http
 		host:     host,
 		port:     port,
 		dir:      dir,
-		cols:     btree.New(16, 0),
 		follows:  make(map[*bytes.Buffer]bool),
 		fcond:    sync.NewCond(&sync.Mutex{}),
 		lives:    make(map[*liveBuffer]bool),
@@ -353,30 +343,29 @@ func (c *Controller) watchLuaStatePool() {
 }
 
 func (c *Controller) setCol(key string, col *collection.Collection) {
-	c.cols.ReplaceOrInsert(&collectionT{Key: key, Collection: col})
+	c.cols.Set(key, col)
 }
 
 func (c *Controller) getCol(key string) *collection.Collection {
-	item := c.cols.Get(&collectionT{Key: key})
+	if value, ok := c.cols.Get(key); ok {
-	if item == nil {
+		return value.(*collection.Collection)
-		return nil
 	}
-	return item.(*collectionT).Collection
+	return nil
 }
 
-func (c *Controller) scanGreaterOrEqual(key string, iterator func(key string, col *collection.Collection) bool) {
+func (c *Controller) scanGreaterOrEqual(
-	c.cols.AscendGreaterOrEqual(&collectionT{Key: key}, func(item btree.Item) bool {
+	key string, iterator func(key string, col *collection.Collection) bool,
-		col := item.(*collectionT)
+) {
-		return iterator(col.Key, col.Collection)
+	c.cols.Ascend(key, func(ikey string, ivalue interface{}) bool {
+		return iterator(ikey, ivalue.(*collection.Collection))
 	})
 }
 
 func (c *Controller) deleteCol(key string) *collection.Collection {
-	i := c.cols.Delete(&collectionT{Key: key})
+	if prev, ok := c.cols.Delete(key); ok {
-	if i == nil {
+		return prev.(*collection.Collection)
-		return nil
 	}
-	return i.(*collectionT).Collection
+	return nil
 }
 
 func isReservedFieldName(field string) bool {
@@ -625,7 +614,7 @@ func randomKey(n int) string {
 
 func (c *Controller) reset() {
 	c.aofsz = 0
-	c.cols = btree.New(16, 0)
+	c.cols = ds.BTree{}
 	c.exlistmu.Lock()
 	c.exlist = nil
 	c.exlistmu.Unlock()

internal/controller/crud.go

@@ -8,11 +8,11 @@ import (
 	"time"
 
 	"github.com/mmcloughlin/geohash"
-	"github.com/tidwall/btree"
 	"github.com/tidwall/geojson"
 	"github.com/tidwall/geojson/geometry"
 	"github.com/tidwall/resp"
 	"github.com/tidwall/tile38/internal/collection"
+	"github.com/tidwall/tile38/internal/ds"
 	"github.com/tidwall/tile38/internal/glob"
 	"github.com/tidwall/tile38/internal/server"
 )
@@ -461,7 +461,7 @@ func (c *Controller) cmdFlushDB(msg *server.Message) (res resp.Value, d commandD
 		err = errInvalidNumberOfArguments
 		return
 	}
-	c.cols = btree.New(16, 0)
+	c.cols = ds.BTree{}
 	c.exlistmu.Lock()
 	c.exlist = nil
 	c.exlistmu.Unlock()

internal/controller/keys.go

@@ -5,7 +5,6 @@ import (
 	"strings"
 	"time"
 
-	"github.com/tidwall/btree"
 	"github.com/tidwall/resp"
 	"github.com/tidwall/tile38/internal/glob"
 	"github.com/tidwall/tile38/internal/server"
@@ -34,8 +33,7 @@ func (c *Controller) cmdKeys(msg *server.Message) (res resp.Value, err error) {
 	var greaterPivot string
 	var vals []resp.Value
 
-	iterator := func(item btree.Item) bool {
+	iterator := func(key string, value interface{}) bool {
-		key := item.(*collectionT).Key
 		var match bool
 		if everything {
 			match = true
@@ -66,24 +64,24 @@ func (c *Controller) cmdKeys(msg *server.Message) (res resp.Value, err error) {
 	}
 	if pattern == "*" {
 		everything = true
-		c.cols.Ascend(iterator)
+		c.cols.Scan(iterator)
 	} else {
 		if strings.HasSuffix(pattern, "*") {
 			greaterPivot = pattern[:len(pattern)-1]
 			if glob.IsGlob(greaterPivot) {
 				greater = false
-				c.cols.Ascend(iterator)
+				c.cols.Scan(iterator)
 			} else {
 				greater = true
-				c.cols.AscendGreaterOrEqual(&collectionT{Key: greaterPivot}, iterator)
+				c.cols.Ascend(greaterPivot, iterator)
 			}
 		} else if glob.IsGlob(pattern) {
 			greater = false
-			c.cols.Ascend(iterator)
+			c.cols.Scan(iterator)
 		} else {
 			greater = true
 			greaterPivot = pattern
-			c.cols.AscendGreaterOrEqual(&collectionT{Key: greaterPivot}, iterator)
+			c.cols.Ascend(greaterPivot, iterator)
 		}
 	}
 	if msg.OutputType == server.JSON {

internal/controller/stats.go

@@ -10,9 +10,9 @@ import (
 	"strings"
 	"time"
 
-	"github.com/tidwall/btree"
 	"github.com/tidwall/resp"
 	"github.com/tidwall/tile38/core"
+	"github.com/tidwall/tile38/internal/collection"
 	"github.com/tidwall/tile38/internal/server"
 )
 
@@ -86,8 +86,8 @@ func (c *Controller) cmdServer(msg *server.Message) (res resp.Value, err error)
 	m["num_collections"] = c.cols.Len()
 	m["num_hooks"] = len(c.hooks)
 	sz := 0
-	c.cols.Ascend(func(item btree.Item) bool {
+	c.cols.Scan(func(key string, value interface{}) bool {
-		col := item.(*collectionT).Collection
+		col := value.(*collection.Collection)
 		sz += col.TotalWeight()
 		return true
 	})
@@ -95,8 +95,8 @@ func (c *Controller) cmdServer(msg *server.Message) (res resp.Value, err error)
 	points := 0
 	objects := 0
 	strings := 0
-	c.cols.Ascend(func(item btree.Item) bool {
+	c.cols.Scan(func(key string, value interface{}) bool {
-		col := item.(*collectionT).Collection
+		col := value.(*collection.Collection)
 		points += col.PointCount()
 		objects += col.Count()
 		strings += col.StringCount()