diff --git a/cmd_server.go b/cmd_server.go index 6e51727b..b1c81569 100644 --- a/cmd_server.go +++ b/cmd_server.go @@ -3,10 +3,12 @@ package miniredis import ( + "fmt" "strconv" "strings" "github.com/alicebob/miniredis/v2/server" + "github.com/alicebob/miniredis/v2/size" ) func commandsServer(m *Miniredis) { @@ -16,6 +18,69 @@ func commandsServer(m *Miniredis) { m.srv.Register("FLUSHDB", m.cmdFlushdb) m.srv.Register("INFO", m.cmdInfo) m.srv.Register("TIME", m.cmdTime) + m.srv.Register("MEMORY", m.cmdMemory) +} + +// MEMORY +func (m *Miniredis) cmdMemory(c *server.Peer, cmd string, args []string) { + if len(args) == 0 { + setDirty(c) + c.WriteError(errWrongNumber(cmd)) + return + } + if !m.handleAuth(c) { + return + } + if m.checkPubsub(c, cmd) { + return + } + + withTx(m, c, func(c *server.Peer, ctx *connCtx) { + db := m.db(ctx.selectedDB) + + cmd, args := args[0], args[1:] + switch cmd { + case "USAGE": + if len(args) < 1 { + setDirty(c) + c.WriteError(errWrongNumber("memory|usage")) + return + } + if len(args) > 1 { + setDirty(c) + c.WriteError(msgSyntaxError) + return + } + + var ( + value interface{} + ok bool + ) + switch db.keys[args[0]] { + case "string": + value, ok = db.stringKeys[args[0]] + case "set": + value, ok = db.setKeys[args[0]] + case "hash": + value, ok = db.hashKeys[args[0]] + case "list": + value, ok = db.listKeys[args[0]] + case "hll": + value, ok = db.hllKeys[args[0]] + case "zset": + value, ok = db.sortedsetKeys[args[0]] + case "stream": + value, ok = db.streamKeys[args[0]] + } + if !ok { + c.WriteNull() + return + } + c.WriteInt(size.Of(value)) + default: + c.WriteError(fmt.Sprintf(msgMemorySubcommand, strings.ToUpper(cmd))) + } + }) } // DBSIZE diff --git a/cmd_server_test.go b/cmd_server_test.go index f95bd356..bc534968 100644 --- a/cmd_server_test.go +++ b/cmd_server_test.go @@ -127,3 +127,30 @@ func TestCmdServerTime(t *testing.T) { proto.Error(errWrongNumber("time")), ) } + +// Test Memory Usage +func TestCmdServerMemoryUsage(t *testing.T) { + s, err := Run() + ok(t, err) + defer s.Close() + c, err := proto.Dial(s.Addr()) + ok(t, err) + defer c.Close() + + c.Do("SET", "foo", "bar") + mustDo(t, c, + "PFADD", "h", "aap", "noot", "mies", + proto.Int(1), + ) + + // Intended only for having metrics not to be 1:1 Redis + mustDo(t, c, + "MEMORY", "USAGE", "foo", + proto.Int(19), // normally, with Redis it should be 56 but we don't have the same overhead as Redis + ) + // Intended only for having metrics not to be 1:1 Redis + mustDo(t, c, + "MEMORY", "USAGE", "h", + proto.Int(124), // normally, with Redis it should be 56 but we don't have the same overhead as Redis + ) +} diff --git a/go.mod b/go.mod index bbaffb47..a239f335 100644 --- a/go.mod +++ b/go.mod @@ -5,4 +5,6 @@ require ( github.com/yuin/gopher-lua v1.1.0 ) +require github.com/DmitriyVTitov/size v1.5.0 + go 1.14 diff --git a/go.sum b/go.sum index 22f21d81..b52694e5 100644 --- a/go.sum +++ b/go.sum @@ -1,12 +1,12 @@ +github.com/DmitriyVTitov/size v1.5.0 h1:/PzqxYrOyOUX1BXj6J9OuVRVGe+66VL4D9FlUaW515g= +github.com/DmitriyVTitov/size v1.5.0/go.mod h1:le6rNI4CoLQV1b9gzp1+3d7hMAD/uu2QcJ+aYbNgiU0= github.com/alicebob/gopher-json v0.0.0-20200520072559-a9ecdc9d1d3a h1:HbKu58rmZpUGpz5+4FfNmIU+FmZg2P3Xaj2v2bfNWmk= github.com/alicebob/gopher-json v0.0.0-20200520072559-a9ecdc9d1d3a/go.mod h1:SGnFV6hVsYE877CKEZ6tDNTjaSXYUk6QqoIK6PrAtcc= github.com/chzyer/logex v1.1.10/go.mod h1:+Ywpsq7O8HXn0nuIou7OrIPyXbp3wmkHB+jjWRnGsAI= github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e/go.mod h1:nSuG5e5PlCu98SY8svDHJxuZscDgtXS6KTTbou5AhLI= github.com/chzyer/test v0.0.0-20180213035817-a1ea475d72b1/go.mod h1:Q3SI9o4m/ZMnBNeIyt5eFwwo7qiLfzFZmjNmxjkiQlU= -github.com/yuin/gopher-lua v0.0.0-20210529063254-f4c35e4016d9 h1:k/gmLsJDWwWqbLCur2yWnJzwQEKRcAHXo6seXGuSwWw= -github.com/yuin/gopher-lua v0.0.0-20210529063254-f4c35e4016d9/go.mod h1:E1AXubJBdNmFERAOucpDIxNzeGfLzg0mYh+UfMWdChA= -github.com/yuin/gopher-lua v0.0.0-20220504180219-658193537a64 h1:5mLPGnFdSsevFRFc9q3yYbBkB6tsm4aCwwQV/j1JQAQ= -github.com/yuin/gopher-lua v0.0.0-20220504180219-658193537a64/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= +github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE= +github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= github.com/yuin/gopher-lua v1.1.0 h1:BojcDhfyDWgU2f2TOzYK/g5p2gxMrku8oupLDqlnSqE= github.com/yuin/gopher-lua v1.1.0/go.mod h1:GBR0iDaNXjAgGg9zfCvksxSRnQx76gclCIb7kdAd1Pw= golang.org/x/sys v0.0.0-20190204203706-41f3e6584952/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= diff --git a/integration/server_test.go b/integration/server_test.go index bc70b832..491079bd 100644 --- a/integration/server_test.go +++ b/integration/server_test.go @@ -35,6 +35,26 @@ func TestServer(t *testing.T) { c.Error("syntax error", "FLUSHDB", "ASYNC", "ASYNC") c.Error("syntax error", "FLUSHALL", "ASYNC", "foo") }) + + testRaw(t, func(c *client) { + c.Do("SET", "plain", "hello") + c.DoLoosely("MEMORY", "USAGE", "plain") + c.Do("LPUSH", "alist", "hello", "42") + c.DoLoosely("MEMORY", "USAGE", "alist") + c.Do("HSET", "ahash", "key", "value") + c.DoLoosely("MEMORY", "USAGE", "ahash") + c.Do("ZADD", "asset", "0", "line") + c.DoLoosely("MEMORY", "USAGE", "asset") + c.Do("PFADD", "ahll", "123") + c.DoLoosely("MEMORY", "USAGE", "ahll") + c.Do("XADD", "astream", "0-1", "name", "Mercury") + c.DoLoosely("MEMORY", "USAGE", "astream") + c.DoLoosely("MEMORY", "USAGE", "nosuch") + + c.Error("Try MEMORY HELP", "MEMORY", "FOO") + c.Error("wrong number of arguments", "MEMORY", "USAGE") + c.Error("syntax error", "MEMORY", "USAGE", "too", "many") + }) } func TestServerTLS(t *testing.T) { diff --git a/redis.go b/redis.go index 6cc52db0..3fdc7813 100644 --- a/redis.go +++ b/redis.go @@ -54,6 +54,7 @@ const ( msgRankIsZero = "ERR RANK can't be zero: use 1 to start from the first match, 2 from the second ... or use negative to start from the end of the list" msgCountIsNegative = "ERR COUNT can't be negative" msgMaxLengthIsNegative = "ERR MAXLEN can't be negative" + msgMemorySubcommand = "ERR unknown subcommand '%s'. Try MEMORY HELP." ) func errWrongNumber(cmd string) string { diff --git a/size/readme.md b/size/readme.md new file mode 100644 index 00000000..89220e45 --- /dev/null +++ b/size/readme.md @@ -0,0 +1,2 @@ + +Credits to DmitriyVTitov on his package https://github.com/DmitriyVTitov/size diff --git a/size/size.go b/size/size.go new file mode 100644 index 00000000..43fee6e2 --- /dev/null +++ b/size/size.go @@ -0,0 +1,138 @@ +package size + +import ( + "reflect" + "unsafe" +) + +// Of returns the size of 'v' in bytes. +// If there is an error during calculation, Of returns -1. +func Of(v interface{}) int { + // Cache with every visited pointer so we don't count two pointers + // to the same memory twice. + cache := make(map[uintptr]bool) + return sizeOf(reflect.Indirect(reflect.ValueOf(v)), cache) +} + +// sizeOf returns the number of bytes the actual data represented by v occupies in memory. +// If there is an error, sizeOf returns -1. +func sizeOf(v reflect.Value, cache map[uintptr]bool) int { + switch v.Kind() { + + case reflect.Array: + sum := 0 + for i := 0; i < v.Len(); i++ { + s := sizeOf(v.Index(i), cache) + if s < 0 { + return -1 + } + sum += s + } + + return sum + (v.Cap()-v.Len())*int(v.Type().Elem().Size()) + + case reflect.Slice: + // return 0 if this node has been visited already + if cache[v.Pointer()] { + return 0 + } + cache[v.Pointer()] = true + + sum := 0 + for i := 0; i < v.Len(); i++ { + s := sizeOf(v.Index(i), cache) + if s < 0 { + return -1 + } + sum += s + } + + sum += (v.Cap() - v.Len()) * int(v.Type().Elem().Size()) + + return sum + int(v.Type().Size()) + + case reflect.Struct: + sum := 0 + for i, n := 0, v.NumField(); i < n; i++ { + s := sizeOf(v.Field(i), cache) + if s < 0 { + return -1 + } + sum += s + } + + // Look for struct padding. + padding := int(v.Type().Size()) + for i, n := 0, v.NumField(); i < n; i++ { + padding -= int(v.Field(i).Type().Size()) + } + + return sum + padding + + case reflect.String: + s := v.String() + hdr := (*reflect.StringHeader)(unsafe.Pointer(&s)) + if cache[hdr.Data] { + return int(v.Type().Size()) + } + cache[hdr.Data] = true + return len(s) + int(v.Type().Size()) + + case reflect.Ptr: + // return Ptr size if this node has been visited already (infinite recursion) + if cache[v.Pointer()] { + return int(v.Type().Size()) + } + cache[v.Pointer()] = true + if v.IsNil() { + return int(reflect.New(v.Type()).Type().Size()) + } + s := sizeOf(reflect.Indirect(v), cache) + if s < 0 { + return -1 + } + return s + int(v.Type().Size()) + + case reflect.Bool, + reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64, + reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64, + reflect.Int, reflect.Uint, + reflect.Chan, + reflect.Uintptr, + reflect.Float32, reflect.Float64, reflect.Complex64, reflect.Complex128, + reflect.Func: + return int(v.Type().Size()) + + case reflect.Map: + // return 0 if this node has been visited already (infinite recursion) + if cache[v.Pointer()] { + return 0 + } + cache[v.Pointer()] = true + sum := 0 + keys := v.MapKeys() + for i := range keys { + val := v.MapIndex(keys[i]) + // calculate size of key and value separately + sv := sizeOf(val, cache) + if sv < 0 { + return -1 + } + sum += sv + sk := sizeOf(keys[i], cache) + if sk < 0 { + return -1 + } + sum += sk + } + // Include overhead due to unused map buckets. 10.79 comes + // from https://golang.org/src/runtime/map.go. + return sum + int(v.Type().Size()) + int(float64(len(keys))*10.79) + + case reflect.Interface: + return sizeOf(v.Elem(), cache) + int(v.Type().Size()) + + } + + return -1 +} diff --git a/size/size_test.go b/size/size_test.go new file mode 100644 index 00000000..049febb4 --- /dev/null +++ b/size/size_test.go @@ -0,0 +1,65 @@ +package size + +import ( + "testing" +) + +func TestOf(t *testing.T) { + tests := []struct { + name string + v interface{} + want int + }{ + { + name: "Array", + v: [3]int32{1, 2, 3}, // 3 * 4 = 12 + want: 12, + }, + { + name: "Slice", + v: make([]int64, 2, 5), // 5 * 8 + 24 = 64 + want: 64, + }, + { + name: "String", + v: "ABCdef", // 6 + 16 = 22 + want: 22, + }, + { + name: "Map", + // (8 + 3 + 16) + (8 + 4 + 16) = 55 + // 55 + 8 + 10.79 * 2 = 84 + v: map[int64]string{0: "ABC", 1: "DEFG"}, + want: 84, + }, + { + name: "Struct", + v: struct { + slice []int64 + array [2]bool + structure struct { + i int8 + s string + } + }{ + slice: []int64{12345, 67890}, // 2 * 8 + 24 = 40 + array: [2]bool{true, false}, // 2 * 1 = 2 + structure: struct { + i int8 + s string + }{ + i: 5, // 1 + s: "abc", // 3 * 1 + 16 = 19 + }, // 20 + 7 (padding) = 27 + }, // 40 + 2 + 27 = 69 + 6 (padding) = 75 + want: 75, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := Of(tt.v); got != tt.want { + t.Errorf("Of() = %v, want %v", got, tt.want) + } + }) + } +}