From c8c03715e548ce4165a13a77f08a649d72a95e5e Mon Sep 17 00:00:00 2001
From: Gabriele Cimato <gabriele.cimato@gmail.com>
Date: Mon, 10 Jun 2024 14:59:37 -0500
Subject: [PATCH] feat: add mem usage estimation for map objects

---
 builtin_map.go      | 77 +++++++++++++++++++++++++++++++++++----------
 builtin_map_test.go | 47 +++++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 17 deletions(-)

diff --git a/builtin_map.go b/builtin_map.go
index 42017685..9bc1ebb7 100644
--- a/builtin_map.go
+++ b/builtin_map.go
@@ -1,6 +1,7 @@
 package goja
 
 import (
+	"math"
 	"reflect"
 )
 
@@ -94,6 +95,40 @@ func (mo *mapObject) exportToMap(dst reflect.Value, typ reflect.Type, ctx *objec
 	return nil
 }
 
+// estimateMemUsage helps calculating mem usage for large objects.
+// It will sample the object and use those samples to estimate the
+// mem usage.
+func (mo *mapObject) estimateMemUsage(ctx *MemUsageContext) (estimate uint64, err error) {
+	var samplesVisited, memUsage uint64
+	totalItems := mo.m.size
+	if totalItems == 0 {
+		return memUsage, nil
+	}
+	sampleSize := int(math.Floor(float64(totalItems) * SampleRate))
+
+	i := 0
+	for item := mo.m.iterFirst; item != nil && i < totalItems; item = item.iterNext {
+		if i%sampleSize != 0 {
+			continue
+		}
+		samplesVisited += 1
+		inc, err := item.key.MemUsage(ctx)
+		memUsage += inc
+		if err != nil {
+			return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), err
+		}
+
+		inc, err = item.value.MemUsage(ctx)
+		memUsage += inc
+		if err != nil {
+			return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), err
+		}
+		i += 1
+	}
+
+	return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), nil
+}
+
 func (mo *mapObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
 	if mo == nil || ctx.IsObjVisited(mo) {
 		return SizeEmptyStruct, nil
@@ -109,28 +144,36 @@ func (mo *mapObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error)
 		return memUsage, err
 	}
 
-	for _, entry := range mo.m.hashTable {
-		if entry == nil {
-			continue
+	if ctx.ObjectPropsLenExceedsThreshold(mo.m.size) {
+		inc, err := mo.estimateMemUsage(ctx)
+		memUsage += inc
+		if err != nil {
+			return memUsage, err
 		}
+	} else {
+		for _, entry := range mo.m.hashTable {
+			if entry == nil {
+				continue
+			}
 
-		if entry.key != nil {
-			inc, err := entry.key.MemUsage(ctx)
-			memUsage += inc
-			if err != nil {
-				return memUsage, err
+			if entry.key != nil {
+				inc, err := entry.key.MemUsage(ctx)
+				memUsage += inc
+				if err != nil {
+					return memUsage, err
+				}
 			}
-		}
 
-		if entry.value != nil {
-			inc, err := entry.value.MemUsage(ctx)
-			memUsage += inc
-			if err != nil {
-				return memUsage, err
+			if entry.value != nil {
+				inc, err := entry.value.MemUsage(ctx)
+				memUsage += inc
+				if err != nil {
+					return memUsage, err
+				}
+			}
+			if exceeded := ctx.MemUsageLimitExceeded(memUsage); exceeded {
+				return memUsage, nil
 			}
-		}
-		if exceeded := ctx.MemUsageLimitExceeded(memUsage); exceeded {
-			return memUsage, nil
 		}
 	}
 
diff --git a/builtin_map_test.go b/builtin_map_test.go
index ec4f8433..8d931ca2 100644
--- a/builtin_map_test.go
+++ b/builtin_map_test.go
@@ -3,6 +3,7 @@ package goja
 import (
 	"fmt"
 	"hash/maphash"
+	"math"
 	"testing"
 )
 
@@ -243,8 +244,40 @@ func BenchmarkMapDeleteJS(b *testing.B) {
 	}
 }
 
+func createOrderedMap(vm *Runtime, size int) *orderedMap {
+	ht := make(map[uint64]*mapEntry, 0)
+	for i := 0; i < size; i += 1 {
+		value := vm.ToValue("value")
+		// This is leveraging the fact that we sample only 10% of the data.
+		// We intentionally set the non-sampled items to something different
+		// so that we can show in our test that we are correctly using
+		// the samples to estimate mem usage and nothing else.
+		if i%(int(math.Floor(float64(size)*SampleRate))) == 1 {
+			value = vm.ToValue("verylongstring")
+		}
+
+		ht[uint64(i)] = &mapEntry{
+			key:   vm.ToValue("key"),
+			value: value,
+		}
+		// These iter items are necessary for testing the mem usage
+		// estimation since that's who we iterate through the map.
+		if i > 0 {
+			ht[uint64(i)].iterPrev = ht[uint64(i-1)]
+			ht[uint64(i-1)].iterNext = ht[uint64(i)]
+		}
+	}
+	return &orderedMap{
+		size:      size,
+		iterFirst: ht[uint64(0)],
+		iterLast:  ht[uint64(size-1)],
+		hashTable: ht,
+	}
+}
+
 func TestMapObjectMemUsage(t *testing.T) {
 	vm := New()
+
 	tests := []struct {
 		name        string
 		mu          *MemUsageContext
@@ -309,6 +342,20 @@ func TestMapObjectMemUsage(t *testing.T) {
 				(5+SizeString)*3,
 			errExpected: nil,
 		},
+		{
+			name: "mem above estimate threshold and within memory limit returns correct mem usage",
+			mu:   NewMemUsageContext(vm, 88, 100, 50, 5, TestNativeMemUsageChecker{}),
+			mo: &mapObject{
+				m: createOrderedMap(vm, 20),
+			},
+			// baseObject
+			expectedMem: SizeEmptyStruct +
+				// len(key) + overhead (we reach the limit after 3)
+				(3+SizeString)*20 +
+				// len(value) + overhead (we reach the limit after 3)
+				(5+SizeString)*20,
+			errExpected: nil,
+		},
 	}
 
 	for _, tc := range tests {