BAAS-22176: add mem usage estimation for map objects

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

array.go

@@ -559,7 +559,7 @@
 }
 
 func toIdx(v valueInt) uint32 {
 	if v >= 0 && v < math.MaxUint32 {
 		return uint32(v)
 	}
 	return math.MaxUint32
@@ -580,7 +580,7 @@
 	if arrayLen == 0 {
 		return memUsage, nil
 	}
-	sampleSize := arrayLen / 10
+	sampleSize := ctx.ComputeSampleStep(arrayLen)
 
 	// grabbing one sample every "sampleSize" to provide consistent
 	// memory usage across function executions

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

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

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{},
 			},

builtin_map.go

@@ -94,6 +94,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 := ctx.ComputeSampleStep(totalItems)
+
+	// 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
+		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
+		}
+	}
+
+	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 +143,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
 		}
 	}
 

builtin_map_test.go

@@ -243,8 +243,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%(computeSampleStep(size, 0.1)) == 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
@@ -254,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{
@@ -271,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{
@@ -309,6 +341,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, 0.1, 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 {

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

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

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

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

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