Memory Usage (#341)

SERVER MEMORY

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

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
+	)
+}

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

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=

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) {

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 {

size/readme.md

@@ -0,0 +1,2 @@
+
+Credits to DmitriyVTitov on his package https://github.com/DmitriyVTitov/size

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
+}

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)
+			}
+		})
+	}
+}