diff --git a/controller/config.go b/controller/config.go index 2b1fd8d7..cd96e7e9 100644 --- a/controller/config.go +++ b/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") diff --git a/controller/controller.go b/controller/controller.go index e777ca53..a9a95cf4 100644 --- a/controller/controller.go +++ b/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}) } diff --git a/controller/crud.go b/controller/crud.go index d99f390b..98a3bdc0 100644 --- a/controller/crud.go +++ b/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 diff --git a/controller/stats.go b/controller/stats.go index 99c20dc8..ea0240b0 100644 --- a/controller/stats.go +++ b/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 diff --git a/vendor/github.com/google/btree/btree.go b/vendor/github.com/google/btree/btree.go index eb1f75cd..fc5aaaa1 100644 --- a/vendor/github.com/google/btree/btree.go +++ b/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,26 +483,22 @@ func (t *BTree) minItems() int { } func (t *BTree) newNode() (n *node) { - index := len(t.freelist) - 1 - if index < 0 { - return &node{t: t} - } - t.freelist, n = t.freelist[:index], t.freelist[index] + n = t.freelist.newNode() + n.t = t 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 - } - n.items = n.items[:0] - for i := range n.children { - n.children[i] = nil // clear to allow GC - } - n.children = n.children[:0] - t.freelist = append(t.freelist, n) + for i := range n.items { + n.items[i] = nil // clear to allow GC } + n.items = n.items[:0] + for i := range n.children { + n.children[i] = nil // clear to allow GC + } + n.children = n.children[:0] + 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 diff --git a/vendor/github.com/google/btree/btree_test.go b/vendor/github.com/google/btree/btree_test.go index 43dcea63..0a2fdde1 100644 --- a/vendor/github.com/google/btree/btree_test.go +++ b/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: // replace5: 5 // replace100: + // min: 0 // delmin: 0 + // max: 100 // delmax: 100 // len: 8 }