added maxmemory setting fixes #15

controller/config.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"strconv"
 	"strings"
 	"time"
 
@@ -12,7 +13,7 @@ import (
 	"github.com/tidwall/tile38/controller/server"
 )
 
-var validProperties = []string{"requirepass", "leaderauth", "protected-mode"}
+var validProperties = []string{"requirepass", "leaderauth", "protected-mode", "maxmemory"}
 
 // Config is a tile38 config
 type Config struct {
@@ -30,6 +31,8 @@ type Config struct {
 	LeaderAuth     string `json:"-"`
 	ProtectedModeP string `json:"protected-mode,omitempty"`
 	ProtectedMode  string `json:"-"`
+	MaxMemoryP     string `json:"maxmemory,omitempty"`
+	MaxMemory      int    `json:"-"`
 }
 
 func (c *Controller) loadConfig() error {
@@ -54,9 +57,58 @@ func (c *Controller) loadConfig() error {
 	if err := c.setConfigProperty("protected-mode", c.config.ProtectedModeP, true); err != nil {
 		return err
 	}
+	if err := c.setConfigProperty("maxmemory", c.config.MaxMemoryP, true); err != nil {
+		return err
+	}
 	return nil
 }
 
+func parseMemSize(s string) (bytes int, ok bool) {
+	if s == "" {
+		return 0, true
+	}
+	s = strings.ToLower(s)
+	var n uint64
+	var sz int
+	var err error
+	if strings.HasSuffix(s, "gb") {
+		n, err = strconv.ParseUint(s[:len(s)-2], 10, 64)
+		sz = int(n * 1024 * 1024 * 1024)
+	} else if strings.HasSuffix(s, "mb") {
+		n, err = strconv.ParseUint(s[:len(s)-2], 10, 64)
+		sz = int(n * 1024 * 1024)
+	} else if strings.HasSuffix(s, "kb") {
+		n, err = strconv.ParseUint(s[:len(s)-2], 10, 64)
+		sz = int(n * 1024)
+	} else {
+		n, err = strconv.ParseUint(s, 10, 64)
+		sz = int(n)
+	}
+	if err != nil {
+		return 0, false
+	}
+	return sz, true
+}
+
+func formatMemSize(sz int) string {
+	if sz <= 0 {
+		return ""
+	}
+	if sz < 1024 {
+		return strconv.FormatInt(int64(sz), 10)
+	}
+	sz /= 1024
+	if sz < 1024 {
+		return strconv.FormatInt(int64(sz), 10) + "kb"
+	}
+	sz /= 1024
+	if sz < 1024 {
+		return strconv.FormatInt(int64(sz), 10) + "mb"
+	}
+	sz /= 1024
+	return strconv.FormatInt(int64(sz), 10) + "gb"
+}
+
 func (c *Controller) setConfigProperty(name, value string, fromLoad bool) error {
 	var invalid bool
 	switch name {
@@ -66,6 +118,12 @@ func (c *Controller) setConfigProperty(name, value string, fromLoad bool) error
 		c.config.RequirePass = value
 	case "leaderauth":
 		c.config.LeaderAuth = value
+	case "maxmemory":
+		sz, ok := parseMemSize(value)
+		if !ok {
+			return fmt.Errorf("Invalid argument '%s' for CONFIG SET '%s'", value, name)
+		}
+		c.config.MaxMemory = sz
 	case "protected-mode":
 		switch strings.ToLower(value) {
 		case "":
@@ -106,6 +164,8 @@ func (c *Controller) getConfigProperty(name string) string {
 		return c.config.LeaderAuth
 	case "protected-mode":
 		return c.config.ProtectedMode
+	case "maxmemory":
+		return formatMemSize(c.config.MaxMemory)
 	}
 }
 
@@ -128,6 +188,7 @@ func (c *Controller) writeConfig(writeProperties bool) error {
 		c.config.RequirePassP = c.config.RequirePass
 		c.config.LeaderAuthP = c.config.LeaderAuth
 		c.config.ProtectedModeP = c.config.ProtectedMode
+		c.config.MaxMemoryP = formatMemSize(c.config.MaxMemory)
 	}
 	var data []byte
 	data, err = json.MarshalIndent(c.config, "", "\t")

controller/controller.go

@@ -23,6 +23,8 @@ import (
 	"github.com/tidwall/tile38/geojson"
 )
 
+var errOOM = errors.New("OOM command not allowed when used memory > 'maxmemory'")
+
 type collectionT struct {
 	Key        string
 	Collection *collection.Collection
@@ -68,6 +70,9 @@ type Controller struct {
 	hooks     map[string]*Hook            // hook name
 	hookcols  map[string]map[string]*Hook // col key
 	aofconnM  map[net.Conn]bool
+
+	stopWatchingMemory bool
+	outOfMemory        bool
 }
 
 // ListenAndServe starts a new tile38 server
@@ -114,6 +119,12 @@ func ListenAndServe(host string, port int, dir string) error {
 		c.mu.Unlock()
 	}()
 	go c.processLives()
+	go c.watchMemory()
+	defer func() {
+		c.mu.Lock()
+		c.stopWatchingMemory = true
+		c.mu.Unlock()
+	}()
 	handler := func(conn *server.Conn, msg *server.Message, rd *server.AnyReaderWriter, w io.Writer, websocket bool) error {
 		err := c.handleInputCommand(conn, msg, w)
 		if err != nil {
@@ -141,6 +152,39 @@ func ListenAndServe(host string, port int, dir string) error {
 	return server.ListenAndServe(host, port, protected, handler)
 }
 
+func (c *Controller) watchMemory() {
+	t := time.NewTicker(time.Second * 2)
+	defer t.Stop()
+	var mem runtime.MemStats
+	for range t.C {
+		func() {
+			c.mu.RLock()
+			if c.stopWatchingMemory {
+				c.mu.RUnlock()
+				return
+			}
+			maxmem := c.config.MaxMemory
+			oom := c.outOfMemory
+			c.mu.RUnlock()
+			if maxmem == 0 {
+				if oom {
+					c.mu.Lock()
+					c.outOfMemory = false
+					c.mu.Unlock()
+				}
+				return
+			}
+			if oom {
+				runtime.GC()
+			}
+			runtime.ReadMemStats(&mem)
+			c.mu.Lock()
+			c.outOfMemory = int(mem.HeapAlloc) > maxmem
+			c.mu.Unlock()
+		}()
+	}
+}
+
 func (c *Controller) setCol(key string, col *collection.Collection) {
 	c.cols.ReplaceOrInsert(&collectionT{Key: key, Collection: col})
 }

controller/crud.go

@@ -506,6 +506,10 @@ func (c *Controller) parseSetArgs(vs []resp.Value) (d commandDetailsT, fields []
 }
 
 func (c *Controller) cmdSet(msg *server.Message) (res string, d commandDetailsT, err error) {
+	if c.config.MaxMemory > 0 && c.outOfMemory {
+		err = errOOM
+		return
+	}
 	start := time.Now()
 	vs := msg.Values[1:]
 	var fields []string

controller/stats.go

@@ -104,6 +104,7 @@ func (c *Controller) cmdServer(msg *server.Message) (res string, err error) {
 		avgsz = int(mem.HeapAlloc) / points
 	}
 	m["heap_size"] = mem.HeapAlloc
+	m["max_heap_size"] = c.config.MaxMemory
 	m["avg_item_size"] = avgsz
 	m["pointer_size"] = (32 << uintptr(uint64(^uintptr(0))>>63)) / 8
 	m["read_only"] = c.config.ReadOnly

vendor/github.com/google/btree/btree.go

@@ -64,6 +64,39 @@ type Item interface {
 	Less(than Item) bool
 }
 
+const (
+	DefaultFreeListSize = 32
+)
+
+// FreeList represents a free list of btree nodes. By default each
+// BTree has its own FreeList, but multiple BTrees can share the same
+// FreeList.
+// Two Btrees using the same freelist are not safe for concurrent write access.
+type FreeList struct {
+	freelist []*node
+}
+
+// NewFreeList creates a new free list.
+// size is the maximum size of the returned free list.
+func NewFreeList(size int) *FreeList {
+	return &FreeList{freelist: make([]*node, 0, size)}
+}
+
+func (f *FreeList) newNode() (n *node) {
+	index := len(f.freelist) - 1
+	if index < 0 {
+		return new(node)
+	}
+	f.freelist, n = f.freelist[:index], f.freelist[index]
+	return
+}
+
+func (f *FreeList) freeNode(n *node) {
+	if len(f.freelist) < cap(f.freelist) {
+		f.freelist = append(f.freelist, n)
+	}
+}
+
 // ItemIterator allows callers of Ascend* to iterate in-order over portions of
 // the tree.  When this function returns false, iteration will stop and the
 // associated Ascend* function will immediately return.
@@ -74,12 +107,17 @@ type ItemIterator func(i Item) bool
 // New(2), for example, will create a 2-3-4 tree (each node contains 1-3 items
 // and 2-4 children).
 func New(degree int) *BTree {
+	return NewWithFreeList(degree, NewFreeList(DefaultFreeListSize))
+}
+
+// NewWithFreeList creates a new B-Tree that uses the given node free list.
+func NewWithFreeList(degree int, f *FreeList) *BTree {
 	if degree <= 1 {
 		panic("bad degree")
 	}
 	return &BTree{
 		degree:   degree,
-		freelist: make([]*node, 0, 32),
+		freelist: f,
 	}
 }
 
@@ -100,6 +138,7 @@ func (s *items) insertAt(index int, item Item) {
 // back.
 func (s *items) removeAt(index int) Item {
 	item := (*s)[index]
+	(*s)[index] = nil
 	copy((*s)[index:], (*s)[index+1:])
 	*s = (*s)[:len(*s)-1]
 	return item
@@ -108,7 +147,9 @@ func (s *items) removeAt(index int) Item {
 // pop removes and returns the last element in the list.
 func (s *items) pop() (out Item) {
 	index := len(*s) - 1
-	out, *s = (*s)[index], (*s)[:index]
+	out = (*s)[index]
+	(*s)[index] = nil
+	*s = (*s)[:index]
 	return
 }
 
@@ -142,6 +183,7 @@ func (s *children) insertAt(index int, n *node) {
 // back.
 func (s *children) removeAt(index int) *node {
 	n := (*s)[index]
+	(*s)[index] = nil
 	copy((*s)[index:], (*s)[index+1:])
 	*s = (*s)[:len(*s)-1]
 	return n
@@ -150,7 +192,9 @@ func (s *children) removeAt(index int) *node {
 // pop removes and returns the last element in the list.
 func (s *children) pop() (out *node) {
 	index := len(*s) - 1
-	out, *s = (*s)[index], (*s)[:index]
+	out = (*s)[index]
+	(*s)[index] = nil
+	*s = (*s)[:index]
 	return
 }
 
@@ -234,6 +278,34 @@ func (n *node) get(key Item) Item {
 	return nil
 }
 
+// min returns the first item in the subtree.
+func min(n *node) Item {
+	if n == nil {
+		return nil
+	}
+	for len(n.children) > 0 {
+		n = n.children[0]
+	}
+	if len(n.items) == 0 {
+		return nil
+	}
+	return n.items[0]
+}
+
+// max returns the last item in the subtree.
+func max(n *node) Item {
+	if n == nil {
+		return nil
+	}
+	for len(n.children) > 0 {
+		n = n.children[len(n.children)-1]
+	}
+	if len(n.items) == 0 {
+		return nil
+	}
+	return n.items[len(n.items)-1]
+}
+
 // toRemove details what item to remove in a node.remove call.
 type toRemove int
 
@@ -396,7 +468,7 @@ type BTree struct {
 	degree   int
 	length   int
 	root     *node
-	freelist []*node
+	freelist *FreeList
 }
 
 // maxItems returns the max number of items to allow per node.
@@ -411,16 +483,12 @@ func (t *BTree) minItems() int {
 }
 
 func (t *BTree) newNode() (n *node) {
-	index := len(t.freelist) - 1
+	n = t.freelist.newNode()
-	if index < 0 {
+	n.t = t
-		return &node{t: t}
-	}
-	t.freelist, n = t.freelist[:index], t.freelist[index]
 	return
 }
 
 func (t *BTree) freeNode(n *node) {
-	if len(t.freelist) < cap(t.freelist) {
 	for i := range n.items {
 		n.items[i] = nil // clear to allow GC
 	}
@@ -429,8 +497,8 @@ func (t *BTree) freeNode(n *node) {
 		n.children[i] = nil // clear to allow GC
 	}
 	n.children = n.children[:0]
-		t.freelist = append(t.freelist, n)
+	n.t = nil // clear to allow GC
-	}
+	t.freelist.freeNode(n)
 }
 
 // ReplaceOrInsert adds the given item to the tree.  If an item in the tree
@@ -552,6 +620,16 @@ func (t *BTree) Get(key Item) Item {
 	return t.root.get(key)
 }
 
+// Min returns the smallest item in the tree, or nil if the tree is empty.
+func (t *BTree) Min() Item {
+	return min(t.root)
+}
+
+// Max returns the largest item in the tree, or nil if the tree is empty.
+func (t *BTree) Max() Item {
+	return max(t.root)
+}
+
 // Has returns true if the given key is in the tree.
 func (t *BTree) Has(key Item) bool {
 	return t.Get(key) != nil

vendor/github.com/google/btree/btree_test.go

@@ -60,6 +60,12 @@ func TestBTree(t *testing.T) {
 	tr := New(*btreeDegree)
 	const treeSize = 10000
 	for i := 0; i < 10; i++ {
+		if min := tr.Min(); min != nil {
+			t.Fatalf("empty min, got %+v", min)
+		}
+		if max := tr.Max(); max != nil {
+			t.Fatalf("empty max, got %+v", max)
+		}
 		for _, item := range perm(treeSize) {
 			if x := tr.ReplaceOrInsert(item); x != nil {
 				t.Fatal("insert found item", item)
@@ -70,6 +76,12 @@ func TestBTree(t *testing.T) {
 				t.Fatal("insert didn't find item", item)
 			}
 		}
+		if min, want := tr.Min(), Item(Int(0)); min != want {
+			t.Fatalf("min: want %+v, got %+v", want, min)
+		}
+		if max, want := tr.Max(), Item(Int(treeSize-1)); max != want {
+			t.Fatalf("max: want %+v, got %+v", want, max)
+		}
 		got := all(tr)
 		want := rang(treeSize)
 		if !reflect.DeepEqual(got, want) {
@@ -98,7 +110,9 @@ func ExampleBTree() {
 	fmt.Println("del100:    ", tr.Delete(Int(100)))
 	fmt.Println("replace5:  ", tr.ReplaceOrInsert(Int(5)))
 	fmt.Println("replace100:", tr.ReplaceOrInsert(Int(100)))
+	fmt.Println("min:       ", tr.Min())
 	fmt.Println("delmin:    ", tr.DeleteMin())
+	fmt.Println("max:       ", tr.Max())
 	fmt.Println("delmax:    ", tr.DeleteMax())
 	fmt.Println("len:       ", tr.Len())
 	// Output:
@@ -109,7 +123,9 @@ func ExampleBTree() {
 	// del100:     <nil>
 	// replace5:   5
 	// replace100: <nil>
+	// min:        0
 	// delmin:     0
+	// max:        100
 	// delmax:     100
 	// len:        8
 }