Use memory footprint of sugardb.store to compare against max memory for eviction policies

internal/constants/const.go

@@ -67,3 +67,8 @@ const (
 	AllKeysRandom  = "allkeys-random"
 	VolatileRandom = "volatile-random"
 )
+
+// CompositeTypes are SugarDB KeyData Value types like set, sorted set, etc.
+type CompositeType interface {
+	GetMem() int64
+}

internal/modules/admin/commands_test.go

@@ -392,6 +392,7 @@ func Test_AdminCommands(t *testing.T) {
 		respConn := resp.NewConn(conn)
 
 		for i := 0; i < len(tests); i++ {
+			t.Log(tests[i].name)
 			if len(tests[i].wantExecRes) > 0 {
 				// If the length of execCommand is > 0, write the command to the connection.
 				if err := respConn.WriteArray(tests[i].execCommand); err != nil {

internal/modules/generic/commands_test.go

@@ -3353,7 +3353,7 @@ func Test_LFU_Generic(t *testing.T) {
 			DataDir:          "",
 			EvictionPolicy:   constants.AllKeysLFU,
 			EvictionInterval: duration,
-			MaxMemory:        4000000,
+			MaxMemory:        550,
 		}),
 	)
 	if err != nil {
@@ -3534,7 +3534,7 @@ func Test_LRU_Generic(t *testing.T) {
 			DataDir:          "",
 			EvictionPolicy:   constants.AllKeysLRU,
 			EvictionInterval: duration,
-			MaxMemory:        4000000,
+			MaxMemory:        550,
 		}),
 	)
 	if err != nil {

internal/modules/pubsub/pubsub.go

@@ -17,12 +17,13 @@ package pubsub
 import (
 	"context"
 	"fmt"
-	"github.com/gobwas/glob"
-	"github.com/tidwall/resp"
 	"log"
 	"net"
 	"slices"
 	"sync"
+
+	"github.com/gobwas/glob"
+	"github.com/tidwall/resp"
 )
 
 type PubSub struct {

internal/modules/set/set.go

@@ -15,16 +15,36 @@
 package set
 
 import (
-	"github.com/echovault/sugardb/internal"
 	"math/rand"
 	"slices"
+	"unsafe"
+
+	"github.com/echovault/sugardb/internal"
+	"github.com/echovault/sugardb/internal/constants"
 )
 
 type Set struct {
 	members map[string]interface{}
 	length  int
 }
 
+func (s *Set) GetMem() int64 {
+	var size int64
+	size += int64(unsafe.Sizeof(s))
+	// above only gives us the size of the pointer to the map, so we need to add it's headers and contents
+	size += int64(unsafe.Sizeof(s.members))
+	for k, v := range s.members {
+		size += int64(unsafe.Sizeof(k))
+		size += int64(len(k))
+		size += int64(unsafe.Sizeof(v))
+	}
+
+	return size
+}
+
+// compile time interface check
+var _ constants.CompositeType = (*Set)(nil)
+
 func NewSet(elems []string) *Set {
 	set := &Set{
 		members: make(map[string]interface{}),

internal/modules/sorted_set/sorted_set.go

@@ -17,11 +17,14 @@ package sorted_set
 import (
 	"cmp"
 	"errors"
-	"github.com/echovault/sugardb/internal"
 	"math"
 	"math/rand"
 	"slices"
 	"strings"
+	"unsafe"
+
+	"github.com/echovault/sugardb/internal"
+	"github.com/echovault/sugardb/internal/constants"
 )
 
 type Value string
@@ -45,6 +48,29 @@ type SortedSet struct {
 	members map[Value]MemberObject
 }
 
+func (s *SortedSet) GetMem() int64 {
+	var size int64
+	// map header
+	size += int64(unsafe.Sizeof(s))
+	// map contents
+	for k, v := range s.members {
+		// string header
+		size += int64(unsafe.Sizeof(k))
+		// string
+		size += int64(len(k))
+		// MemberObject
+		size += int64(unsafe.Sizeof(v))
+		// value field
+		size += int64(unsafe.Sizeof(v.Value))
+		size += int64(len(v.Value))
+	}
+
+	return size
+}
+
+// compile time interface check
+var _ constants.CompositeType = (*SortedSet)(nil)
+
 func NewSortedSet(members []MemberParam) *SortedSet {
 	s := &SortedSet{
 		members: make(map[Value]MemberObject),

internal/types.go

@@ -16,16 +16,82 @@ package internal
 
 import (
 	"context"
-	"github.com/echovault/sugardb/internal/clock"
+	"errors"
+	"fmt"
 	"net"
+	"reflect"
 	"time"
+	"unsafe"
+
+	"github.com/echovault/sugardb/internal/clock"
+	"github.com/echovault/sugardb/internal/constants"
 )
 
 type KeyData struct {
 	Value    interface{}
 	ExpireAt time.Time
 }
 
+func (k *KeyData) GetMem() (int64, error) {
+	var size int64
+	size = int64(unsafe.Sizeof(k.ExpireAt))
+
+	// check type of Value field
+	switch v := k.Value.(type) {
+	case nil:
+		size += 0
+	// AdaptType() will always ensure data type is of string, float64 or int.
+	case int:
+		size += int64(unsafe.Sizeof(v))
+	// int64 data type used with module.SET
+	case float64, int64:
+		size += 8
+	case string:
+		// Add the size of the header and the number of bytes of the string
+		size += int64(unsafe.Sizeof(v))
+		size += int64(len(v))
+
+	// handle hash
+	// AdaptType() will always ensure data type is of string, float64 or int.
+	case map[string]interface{}:
+		// Map headers
+		size += int64(unsafe.Sizeof(v))
+
+		for key, val := range v {
+			size += int64(unsafe.Sizeof(key))
+			size += int64(len(key))
+			switch vt := val.(type) {
+
+			case nil:
+				size += 0
+			case int:
+				size += int64(unsafe.Sizeof(vt))
+			case float64, int64:
+				size += 8
+			case string:
+				size += int64(unsafe.Sizeof(vt))
+				size += int64(len(vt))
+			}
+		}
+
+	// handle list
+	case []string:
+		for _, s := range v {
+			size += int64(unsafe.Sizeof(s))
+			size += int64(len(s))
+		}
+
+	// handle non primitive datatypes like set and sorted set
+	case constants.CompositeType:
+		size += k.Value.(constants.CompositeType).GetMem()
+
+	default:
+		return 0, errors.New(fmt.Sprintf("ERROR: type %v is not supported in method KeyData.GetMem()", reflect.TypeOf(v)))
+	}
+
+	return size, nil
+}
+
 type ContextServerID string
 type ContextConnID string
 
@@ -51,12 +117,14 @@ type SnapshotObject struct {
 
 // ServerInfo holds information about the server/node.
 type ServerInfo struct {
-	Server  string
-	Version string
-	Id      string
-	Mode    string
-	Role    string
-	Modules []string
+	Server     string
+	Version    string
+	Id         string
+	Mode       string
+	Role       string
+	Modules    []string
+	MemoryUsed int64
+	MaxMemory  uint64
 }
 
 // ConnectionInfo holds information about the connection

internal/utils.go

@@ -21,7 +21,6 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
-	"github.com/echovault/sugardb/internal/constants"
 	"io"
 	"log"
 	"math/big"
@@ -34,6 +33,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/echovault/sugardb/internal/constants"
 	"github.com/sethvargo/go-retry"
 	"github.com/tidwall/resp"
 )
@@ -187,27 +187,23 @@ func ParseMemory(memory string) (uint64, error) {
 }
 
 // IsMaxMemoryExceeded checks whether we have exceeded the current maximum memory limit.
-func IsMaxMemoryExceeded(maxMemory uint64) bool {
+func IsMaxMemoryExceeded(memUsed int64, maxMemory uint64) bool {
 	if maxMemory == 0 {
 		return false
 	}
 
-	var memStats runtime.MemStats
-	runtime.ReadMemStats(&memStats)
-
 	// If we're currently using less than the configured max memory, return false.
-	if memStats.HeapInuse < maxMemory {
+	if uint64(memUsed) < maxMemory {
 		return false
 	}
 
 	// If we're currently using more than max memory, force a garbage collection before we start deleting keys.
 	// This measure is to prevent deleting keys that may be important when some memory can be reclaimed
 	// by just collecting garbage.
 	runtime.GC()
-	runtime.ReadMemStats(&memStats)
 
 	// Return true when whe are above or equal to max memory.
-	return memStats.HeapInuse >= maxMemory
+	return uint64(memUsed) >= maxMemory
 }
 
 // FilterExpiredKeys filters out keys that are already expired, so they are not persisted.

sugardb/config.go

@@ -47,7 +47,9 @@ func (server *SugarDB) GetServerInfo() internal.ServerInfo {
 			}
 			return "replica"
 		}(),
-		Modules: server.ListModules(),
+		Modules:    server.ListModules(),
+		MemoryUsed: server.memUsed,
+		MaxMemory:  server.config.MaxMemory,
 	}
 }