refactor: add shared function to compute sample step in mem usage est…

…imates
mongodb-forks · Jun 11, 2024 · 2d49ae4 · 2d49ae4
1 parent c8c0371
commit 2d49ae4
Show file tree

Hide file tree

Showing 31 changed files with 189 additions and 143 deletions.
diff --git a/added_values_test.go b/added_values_test.go
@@ -77,7 +77,7 @@ func TestAddedValuesMemUsage(t *testing.T) {
 		},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
-			mem, err := tc.val.MemUsage(NewMemUsageContext(vm, 100, 100, 100, 100, nil))
+			mem, err := tc.val.MemUsage(NewMemUsageContext(vm, 100, 100, 100, 100, 0.1, nil))
 			if err != nil {
 				t.Fatalf("Unexpected error. Actual: %v Expected: nil", err)
 			}

diff --git a/array.go b/array.go
@@ -580,7 +580,7 @@ func (a *arrayObject) estimateMemUsage(ctx *MemUsageContext) (estimate uint64, e
 	if arrayLen == 0 {
 		return memUsage, nil
 	}
-	sampleSize := int(math.Floor(float64(arrayLen) * SampleRate))
+	sampleSize := ctx.ComputeSampleStep(arrayLen)
 
 	// grabbing one sample every "sampleSize" to provide consistent
 	// memory usage across function executions

diff --git a/array_sparse.go b/array_sparse.go
@@ -508,7 +508,7 @@ func (a *sparseArrayObject) estimateMemUsage(ctx *MemUsageContext) (estimate uin
 	if totalItems == 0 {
 		return memUsage, nil
 	}
-	sampleSize := totalItems / 10
+	sampleSize := ctx.ComputeSampleStep(totalItems)
 
 	// grabbing one sample every "sampleSize" to provide consistent
 	// memory usage across function executions

diff --git a/array_sparse_test.go b/array_sparse_test.go
@@ -274,7 +274,7 @@ func TestSparseArrayObjectMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao: &sparseArrayObject{
 				items: []sparseArrayItem{
 					{
@@ -289,14 +289,14 @@ func TestSparseArrayObjectMemUsage(t *testing.T) {
 		},
 		{
 			name:        "mem is SizeEmptyStruct for nil sparse array",
-			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao:         nil,
 			expected:    SizeEmptyStruct,
 			errExpected: nil,
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao: &sparseArrayObject{
 				items: []sparseArrayItem{
 					{
@@ -331,7 +331,7 @@ func TestSparseArrayObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem above estimate threshold and within memory limit returns correct usage",
-			mu:   NewMemUsageContext(vm, 88, 100, 5, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 100, 5, 50, 0.1, TestNativeMemUsageChecker{}),
 			sao: &sparseArrayObject{
 				items: []sparseArrayItem{
 					{

diff --git a/array_test.go b/array_test.go
@@ -144,23 +144,23 @@ func TestArrayObjectMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold given a nil slice of values",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao:   &arrayObject{},
 			// array overhead + array baseObject
 			expectedMem: SizeEmptyStruct + SizeEmptyStruct,
 			errExpected: nil,
 		},
 		{
 			name: "mem below threshold given empty slice of values",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao:   &arrayObject{values: []Value{}},
 			// array overhead + array baseObject + values slice overhead
 			expectedMem: SizeEmptyStruct + SizeEmptyStruct + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao: &arrayObject{
 				values: []Value{
 					vm.ToValue("key0"),
@@ -179,7 +179,7 @@ func TestArrayObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "empty array with negative threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, -1, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 100, -1, 50, 0.1, TestNativeMemUsageChecker{}),
 			ao: &arrayObject{
 				values: []Value{},
 			},

diff --git a/builtin_map.go b/builtin_map.go
@@ -1,7 +1,6 @@
 package goja
 
 import (
-	"math"
 	"reflect"
 )
 
@@ -104,11 +103,12 @@ func (mo *mapObject) estimateMemUsage(ctx *MemUsageContext) (estimate uint64, er
 	if totalItems == 0 {
 		return memUsage, nil
 	}
-	sampleSize := int(math.Floor(float64(totalItems) * SampleRate))
+	sampleSize := ctx.ComputeSampleStep(totalItems)
 
-	i := 0
-	for item := mo.m.iterFirst; item != nil && i < totalItems; item = item.iterNext {
-		if i%sampleSize != 0 {
+	// We can use samplesVisited instead of an index since we iterate using
+	// iterNext
+	for item := mo.m.iterFirst; item != nil && samplesVisited < uint64(totalItems); item = item.iterNext {
+		if samplesVisited%uint64(sampleSize) != 0 {
 			continue
 		}
 		samplesVisited += 1
@@ -123,7 +123,6 @@ func (mo *mapObject) estimateMemUsage(ctx *MemUsageContext) (estimate uint64, er
 		if err != nil {
 			return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), err
 		}
-		i += 1
 	}
 
 	return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), nil

diff --git a/builtin_map_test.go b/builtin_map_test.go
@@ -3,7 +3,6 @@ package goja
 import (
 	"fmt"
 	"hash/maphash"
-	"math"
 	"testing"
 )
 
@@ -252,7 +251,7 @@ func createOrderedMap(vm *Runtime, size int) *orderedMap {
 		// 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 {
+		if i%(computeSampleStep(size, 0.1)) == 1 {
 			value = vm.ToValue("verylongstring")
 		}
 
@@ -287,7 +286,7 @@ func TestMapObjectMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold",
-			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			mo: &mapObject{
 				m: &orderedMap{
 					hashTable: map[uint64]*mapEntry{
@@ -304,14 +303,14 @@ func TestMapObjectMemUsage(t *testing.T) {
 		},
 		{
 			name:        "mem is SizeEmptyStruct given a nil map object",
-			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			mo:          nil,
 			expectedMem: SizeEmptyStruct,
 			errExpected: nil,
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			mo: &mapObject{
 				m: &orderedMap{
 					hashTable: map[uint64]*mapEntry{
@@ -344,7 +343,7 @@ func TestMapObjectMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem above estimate threshold and within memory limit returns correct mem usage",
-			mu:   NewMemUsageContext(vm, 88, 100, 50, 5, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(vm, 88, 100, 50, 5, 0.1, TestNativeMemUsageChecker{}),
 			mo: &mapObject{
 				m: createOrderedMap(vm, 20),
 			},

diff --git a/builtin_proxy_test.go b/builtin_proxy_test.go
@@ -1290,7 +1290,7 @@ func TestBuiltinProxyMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
 			if err != nil {
 				t.Fatalf("Unexpected error. Actual: %v Expected: nil", err)
 			}

diff --git a/compiler_test.go b/compiler_test.go
@@ -5657,7 +5657,7 @@ func TestProgramMemUsage(t *testing.T) {
 	}{
 		{
 			name: "mem below threshold",
-			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			p: &Program{
 				values: []Value{
 					New().newDateObject(time.Now(), true, nil),
@@ -5669,7 +5669,7 @@ func TestProgramMemUsage(t *testing.T) {
 		},
 		{
 			name: "mem way above threshold returns first crossing of threshold",
-			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, TestNativeMemUsageChecker{}),
+			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, 0.1, TestNativeMemUsageChecker{}),
 			p: &Program{
 				values: []Value{
 					New().newDateObject(time.Now(), true, nil),

diff --git a/date_test.go b/date_test.go
@@ -349,7 +349,7 @@ func TestDateMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}

diff --git a/destruct_test.go b/destruct_test.go
@@ -27,7 +27,7 @@ func TestDestructMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}

diff --git a/func_test.go b/func_test.go
@@ -184,7 +184,7 @@ func TestNativeFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}
@@ -234,7 +234,7 @@ func TestFuncObjectMemUsage(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, nil))
+			total, err := tc.val.MemUsage(NewMemUsageContext(New(), 100, 100, 100, 100, 0.1, nil))
 			if err != tc.errExpected {
 				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
 			}

diff --git a/mem_context.go b/mem_context.go
@@ -2,6 +2,7 @@ package goja
 
 import (
 	"errors"
+	"math"
 )
 
 type visitTracker struct {
@@ -61,6 +62,7 @@ type MemUsageContext struct {
 	NativeMemUsageChecker
 	ArrayLenExceedsThreshold       func(arrayLen int) bool
 	ObjectPropsLenExceedsThreshold func(objPropsLen int) bool
+	ComputeSampleStep              func(totalItems int) int
 	memoryLimit                    uint64
 }
 
@@ -69,6 +71,7 @@ func NewMemUsageContext(
 	maxDepth int,
 	memLimit uint64,
 	arrayLenThreshold, objPropsLenThreshold int,
+	sampleRate float64,
 	nativeChecker NativeMemUsageChecker,
 ) *MemUsageContext {
 	return &MemUsageContext{
@@ -84,6 +87,9 @@ func NewMemUsageContext(
 			// number of obj props beyond which we should estimate mem usage
 			return objPropsLen > objPropsLenThreshold
 		},
+		ComputeSampleStep: func(totalItems int) int {
+			return computeSampleStep(totalItems, sampleRate)
+		},
 	}
 }
 
@@ -93,10 +99,31 @@ func (m *MemUsageContext) MemUsageLimitExceeded(memUsage uint64) bool {
 	return memUsage > m.memoryLimit
 }
 
+func computeMemUsageEstimate(memUsage, samplesVisited uint64, totalProps int) uint64 {
+	// averageMemUsage * total object props
+	return uint64(float32(memUsage) / float32(samplesVisited) * float32(totalProps))
+}
+
+// computeSampleStep will take the total items we want to sample from and a sample rate.
+// It will use this value to determine the sample step, which indicates how often we need
+// to grab a sample. For example, with 100 total items and a 0.2 sample rate, it means
+// we want to sample 20% of 100 items, in order to do so we need to pick an item ever 5
+// (5 * 20 == 100)
+func computeSampleStep(totalItems int, sampleRate float64) int {
+	if sampleRate == 0 || totalItems == 0 {
+		return 1
+	}
+	if sampleRate >= 0.5 {
+		// We allow a max sample size half of the total items
+		sampleRate = 0.5
+	}
+
+	totalSamples := float64(totalItems) * sampleRate
+	return int(math.Floor(float64(totalItems) / totalSamples))
+}
+
 var (
 	ErrMaxDepth = errors.New("reached max depth")
-	// SampleRate represents the percentage of samples we pick up to estimate mem usage
-	SampleRate = 0.1
 )
 
 type MemUsageReporter interface {

diff --git a/mem_context_test.go b/mem_context_test.go
@@ -36,3 +36,30 @@ func TestMemUsageLimitExceeded(t *testing.T) {
 		})
 	}
 }
+
+func TestComputeSampleStep(t *testing.T) {
+	tests := []struct {
+		name       string
+		totalItems int
+		sampleRate float64
+		expected   int
+	}{
+		{"should compute sample size given 10 items and 10% rate", 100, 0.1, 10},
+		{"should compute sample size given 10 items and 20% rate", 100, 0.2, 5},
+		{"should compute sample size given 10 items and 50% rate", 100, 0.5, 2},
+		{"should compute sample size of 5 given 10 items and 70% rate", 100, 0.7, 2},
+		{"should compute sample size of 5 given 10 items and 100% rate", 100, 1, 2},
+		{"should compute sample size of 5 given 10 items and 150% rate", 100, 1.5, 2},
+		{"should compute sample size of 1 given 0 items and 50% rate", 0, 0.5, 1},
+		{"should compute sample size of 1 given 10 items and 0% rate", 100, 0, 1},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			actual := computeSampleStep(tc.totalItems, tc.sampleRate)
+			if actual != tc.expected {
+				t.Fatalf("ACTUAL: %v EXPECTED: %v", actual, tc.expected)
+			}
+		})
+	}
+}