[patch] add Real Map memory size function (#138)

Signed-off-by: kpango <kpango@vdaas.org>

example/main.go

@@ -2,13 +2,58 @@ package main
 
 import (
 	"context"
+	"encoding/gob"
+	"encoding/json"
+	"math/rand"
 	"os"
+	"runtime"
+	"strconv"
 	"time"
+	"unsafe"
 
 	"github.com/kpango/gache/v2"
 	"github.com/kpango/glg"
 )
 
+var (
+	bigData      = map[string]string{}
+	bigDataLen   = 2 << 10
+	bigDataCount = 2 << 11
+)
+
+func init() {
+	for i := 0; i < bigDataCount; i++ {
+		bigData[randStr(bigDataLen)] = randStr(bigDataLen)
+	}
+}
+
+var randSrc = rand.NewSource(time.Now().UnixNano())
+
+const (
+	rs6Letters       = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
+	rs6LetterIdxBits = 6
+	rs6LetterIdxMask = 1<<rs6LetterIdxBits - 1
+	rs6LetterIdxMax  = 63 / rs6LetterIdxBits
+)
+
+func randStr(n int) string {
+	b := make([]byte, n)
+	cache, remain := randSrc.Int63(), rs6LetterIdxMax
+	for i := n - 1; i >= 0; {
+		if remain == 0 {
+			cache, remain = randSrc.Int63(), rs6LetterIdxMax
+		}
+		idx := int(cache & rs6LetterIdxMask)
+		if idx < len(rs6Letters) {
+			b[i] = rs6Letters[idx]
+			i--
+		}
+		cache >>= rs6LetterIdxBits
+		remain--
+	}
+	return *(*string)(unsafe.Pointer(&b))
+}
+
 func main() {
 	var (
 		key1   = "key1"
@@ -45,18 +90,30 @@ func main() {
 		return true
 	})
 
-	file, err := os.OpenFile("/tmp/gache-sample.gdb", os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0o755)
+	var m runtime.MemStats
+	runtime.ReadMemStats(&m)
+	mbody, err := json.Marshal(m)
+	if err == nil {
+		glg.Debugf("memory size: %d, lenght: %d, mem stats: %v", gc.Size(), gc.Len(), string(mbody))
+	}
+	path := "/tmp/gache-sample.gdb"
+
+	file, err := os.OpenFile(path, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0o755)
 	if err != nil {
 		glg.Error(err)
 		return
 	}
-	gc.Write(context.Background(), file)
-
+	gob.Register(struct{}{})
+	err = gc.Write(context.Background(), file)
 	gc.Stop()
 	file.Close()
+	if err != nil {
+		glg.Error(err)
+		return
+	}
 
 	gcn := gache.New[any]().SetDefaultExpire(time.Minute)
-	file, err = os.OpenFile("/tmp/gache-sample.gdb", os.O_RDONLY, 0o755)
+	file, err = os.OpenFile(path, os.O_RDONLY, 0o755)
 	if err != nil {
 		glg.Error(err)
 		return
@@ -88,4 +145,53 @@ func main() {
 		glg.Debugf("key:\t%v\nval:\t%d", k, v)
 		return true
 	})
+
+	runtime.GC()
+	gcs := gache.New[string]()
+	maxCnt := 10000000
+	digitLen :=  len(strconv.Itoa(maxCnt))
+	for i := 0; i < maxCnt; i++ {
+		if i%1000 == 0 {
+			// runtime.ReadMemStats(&m)
+			// mbody, err := json.Marshal(m)
+			if err == nil {
+				// glg.Debugf("before set memory size: %d, lenght: %d, mem stats: %v", gcs.Size(), gcs.Len(), string(mbody))
+				glg.Debugf("Execution No.%-*d:\tbefore set memory size: %d, lenght: %d", digitLen, i, gcs.Size(), gcs.Len())
+			}
+		}
+		for k, v := range bigData {
+			gcs.Set(k, v)
+		}
+		if i%1000 == 0 {
+			// runtime.ReadMemStats(&m)
+			// mbody, err := json.Marshal(m)
+			if err == nil {
+				glg.Debugf("Execution No.%-*d:\tafter set memory size: %d, lenght: %d", digitLen, i, gcs.Size(), gcs.Len())
+				// glg.Debugf("after set memory size: %d, lenght: %d, mem stats: %v", gcs.Size(), gcs.Len(), string(mbody))
+			}
+		}
+
+		for k := range bigData {
+			gcs.Get(k)
+		}
+		for k := range bigData {
+			gcs.Delete(k)
+		}
+		if i%1000 == 0 {
+			// runtime.ReadMemStats(&m)
+			// mbody, err := json.Marshal(m)
+			if err == nil {
+				glg.Debugf("Execution No.%-*d:\tafter delete memory size: %d, lenght: %d", digitLen, i, gcs.Size(), gcs.Len())
+				// glg.Debugf("after delete memory size: %d, lenght: %d, mem stats: %v", gcs.Size(), gcs.Len(), string(mbody))
+			}
+			runtime.GC()
+			// runtime.ReadMemStats(&m)
+			// mbody, err = json.Marshal(m)
+			if err == nil {
+				glg.Debugf("Execution No.%-*d:\tafter gc memory size: %d, lenght: %d", digitLen, i, gcs.Size(), gcs.Len())
+				// glg.Debugf("after gc memory size: %d, lenght: %d, mem stats: %v", gcs.Size(), gcs.Len(), string(mbody))
+			}
+		}
+
+	}
 }

gache.go

@@ -32,6 +32,7 @@ type (
 		SetWithExpire(string, V, time.Duration)
 		StartExpired(context.Context, time.Duration) Gache[V]
 		Len() int
+		Size() uintptr
 		ToMap(context.Context) *sync.Map
 		ToRawMap(context.Context) map[string]V
 		Write(context.Context, io.Writer) error
@@ -58,18 +59,18 @@ type (
 
 	// gache is base instance type
 	gache[V any] struct {
-		expFuncEnabled bool
-		expire         int64
-		l              uint64
+		shards         [slen]*Map[string, *value[V]]
 		cancel         atomic.Pointer[context.CancelFunc]
 		expChan        chan string
 		expFunc        func(context.Context, string)
-		shards         [slen]*Map[string, *value[V]]
+		expFuncEnabled bool
+		expire         int64
+		l              uint64
 	}
 
 	value[V any] struct {
-		expire int64
 		val    V
+		expire int64
 	}
 )
 
@@ -83,6 +84,8 @@ const (
 
 	// NoTTL can be use for disabling ttl cache expiration
 	NoTTL time.Duration = -1
+
+	maxHashKeyLength = 256
 )
 
 // New returns Gache (*gache) instance
@@ -103,8 +106,8 @@ func newMap[V any]() (m *Map[string, *value[V]]) {
 }
 
 func getShardID(key string) (id uint64) {
-	if len(key) > 128 {
-		return xxh3.HashString(key[:128]) & mask
+	if len(key) > maxHashKeyLength {
+		return xxh3.HashString(key[:maxHashKeyLength]) & mask
 	}
 	return xxh3.HashString(key) & mask
 }
@@ -142,7 +145,8 @@ func (g *gache[V]) SetExpiredHook(f func(context.Context, string)) Gache[V] {
 func (g *gache[V]) StartExpired(ctx context.Context, dur time.Duration) Gache[V] {
 	go func() {
 		tick := time.NewTicker(dur)
-		ctx, cancel := context.WithCancel(ctx)
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithCancel(ctx)
 		g.cancel.Store(&cancel)
 		for {
 			select {
@@ -193,7 +197,8 @@ func (g *gache[V]) ToRawMap(ctx context.Context) map[string]V {
 // get returns value & exists from key
 func (g *gache[V]) get(key string) (v V, expire int64, ok bool) {
 	var val *value[V]
-	val, ok = g.shards[getShardID(key)].Load(key)
+	shard := g.shards[getShardID(key)]
+	val, ok = shard.Load(key)
 	if !ok {
 		return v, 0, false
 	}
@@ -222,7 +227,8 @@ func (g *gache[V]) set(key string, val V, expire int64) {
 	if expire > 0 {
 		expire = fastime.UnixNanoNow() + expire
 	}
-	_, loaded := g.shards[getShardID(key)].Swap(key, &value[V]{
+	shard := g.shards[getShardID(key)]
+	_, loaded := shard.Swap(key, &value[V]{
 		expire: expire,
 		val:    val,
 	})
@@ -313,6 +319,19 @@ func (g *gache[V]) Len() int {
 	return *(*int)(unsafe.Pointer(&l))
 }
 
+func (g *gache[V]) Size() (size uintptr) {
+	size += unsafe.Sizeof(g.expFuncEnabled) // bool
+	size += unsafe.Sizeof(g.expire)         // int64
+	size += unsafe.Sizeof(g.l)              // uint64
+	size += unsafe.Sizeof(g.cancel)         // atomic.Pointer[context.CancelFunc]
+	size += unsafe.Sizeof(g.expChan)        // chan string
+	size += unsafe.Sizeof(g.expFunc)        // func(context.Context, string)
+	for _, shard := range g.shards {
+		size += shard.Size()
+	}
+	return size
+}
+
 // Write writes all cached data to writer
 func (g *gache[V]) Write(ctx context.Context, w io.Writer) error {
 	m := g.ToRawMap(ctx)
@@ -329,7 +348,7 @@ func (g *gache[V]) Read(r io.Reader) error {
 		return err
 	}
 	for k, v := range m {
-		g.Set(k, v)
+		go g.Set(k, v)
 	}
 	return nil
 }
@@ -348,3 +367,12 @@ func (g *gache[V]) Clear() {
 		g.shards[i] = newMap[V]()
 	}
 }
+
+func (v *value[V]) Size() uintptr {
+	var size uintptr
+
+	size += unsafe.Sizeof(v.expire) // int64
+	size += unsafe.Sizeof(v.val)    // V size
+
+	return size
+}

go.mod

@@ -1,6 +1,6 @@
 module github.com/kpango/gache/v2
 
-go 1.22.3
+go 1.23.1
 
 require (
 	github.com/kpango/fastime v1.1.9

hmap.go

@@ -0,0 +1,103 @@
+// Copyright (c) 2009 The Go Authors. All rights resered.
+// Modified <Yusuke Kato (kpango)>
+
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are
+// met:
+
+//    * Redistributions of source code must retain the above copyright
+// notice, this list of conditions and the following disclaimer.
+//    * Redistributions in binary form must reproduce the above
+// copyright notice, this list of conditions and the following disclaimer
+// in the documentation and/or other materials provided with the
+// distribution.
+//    * Neither the name of Google Inc. nor the names of its
+// contributors may be used to endorse or promote products derived from
+// this software without specific prior written permission.
+
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+package gache
+
+import "unsafe"
+
+// A header for a Go map.
+type hmap struct {
+	// Note: the format of the hmap is also encoded in cmd/compile/internal/reflectdata/reflect.go.
+	// Make sure this stays in sync with the compiler's definition.
+	count     int // # live cells == size of map.  Must be first (used by len() builtin)
+	flags     uint8
+	B         uint8  // log_2 of # of buckets (can hold up to loadFactor * 2^B items)
+	noverflow uint16 // approximate number of overflow buckets; see incrnoverflow for details
+	hash0     uint32 // hash seed
+
+	buckets    unsafe.Pointer // array of 2^B Buckets. may be nil if count==0.
+	oldbuckets unsafe.Pointer // previous bucket array of half the size, non-nil only when growing
+	nevacuate  uintptr        // progress counter for evacuation (buckets less than this have been evacuated)
+}
+
+const bucketCnt = 8
+
+// A bucket for a Go map.
+type bmap struct {
+	// tophash generally contains the top byte of the hash value
+	// for each key in this bucket. If tophash[0] < minTopHash,
+	// tophash[0] is a bucket evacuation state instead.
+	tophash [bucketCnt]uint8
+	// Followed by bucketCnt keys and then bucketCnt elems.
+	// NOTE: packing all the keys together and then all the elems together makes the
+	// code a bit more complicated than alternating key/elem/key/elem/... but it allows
+	// us to eliminate padding which would be needed for, e.g., map[int64]int8.
+	// Followed by an overflow pointer.
+}
+
+var singleBucketSize = unsafe.Sizeof(bmap{})
+
+func mapSize[K comparable, V any](m map[K]V) (size uintptr) {
+	h := (*hmap)(*(*unsafe.Pointer)(unsafe.Pointer(&m)))
+	if h == nil {
+		return 0
+	}
+	var (
+		zeroK K
+		zeroV V
+	)
+	return h.Size(unsafe.Sizeof(zeroK), unsafe.Sizeof(zeroV))
+}
+
+func (b *bmap) Size() (size uintptr) {
+	return unsafe.Sizeof(b.tophash)
+}
+
+func (h *hmap) Size(kSize, vSize uintptr) (size uintptr) {
+	size += unsafe.Sizeof(h.count)
+	size += unsafe.Sizeof(h.flags)
+	size += unsafe.Sizeof(h.B)
+	size += unsafe.Sizeof(h.noverflow)
+	size += unsafe.Sizeof(h.hash0)
+	size += unsafe.Sizeof(h.buckets)
+	size += unsafe.Sizeof(h.oldbuckets)
+	size += unsafe.Sizeof(h.nevacuate)
+
+	if h.B == 0 {
+		return size
+	}
+	bucketSize := singleBucketSize + (bucketCnt * (kSize + vSize))
+	if h.buckets != nil {
+		size += uintptr(1<<h.B) * bucketSize
+	}
+	if h.oldbuckets != nil && h.B > 1 {
+		size += uintptr(1<<(h.B-1)) * bucketSize
+	}
+	return size
+}