BAAS-32231: Track memory of temporary value created in Array.prototyp…

…e.map calls (#123)

builtin_array.go

@@ -786,6 +786,8 @@ func (r *Runtime) arrayproto_map(call FunctionCall) Value {
 	if _, stdSrc := o.self.(*arrayObject); stdSrc {
 		if arr, ok := a.self.(*arrayObject); ok {
 			values := make([]Value, length)
+			// Point tmpValues to the values slice so that we can track the memory usage
+			r.vm.tmpValues = values
 			for k := int64(0); k < length; k++ {
 				idx := valueInt(k)
 				if val := o.self.getIdx(idx, nil); val != nil {
@@ -794,6 +796,8 @@ func (r *Runtime) arrayproto_map(call FunctionCall) Value {
 					values[k] = callbackFn(fc)
 				}
 			}
+			// Set tmpValues to nil because we no longer need to track the memory usage
+			r.vm.tmpValues = nil
 			setArrayValues(arr, values)
 			return a
 		}

func_test.go

@@ -226,8 +226,8 @@ func TestFuncObjectMemUsage(t *testing.T) {
 					},
 				},
 			},
-			// baseJsFuncObject + value in baseJsFuncObject stash
-			expectedMem: SizeEmptyStruct + SizeInt,
+			// baseJsFuncObject + value in baseJsFuncObject stash + []Value
+			expectedMem: SizeEmptyStruct + SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 	}
@@ -274,8 +274,8 @@ func TestBaseJsFuncObjectMemUsage(t *testing.T) {
 					values: []Value{valueInt(0)},
 				},
 			},
-			// baseJsFuncObject + value in baseJsFuncObject stash
-			expectedMem: SizeEmptyStruct + SizeInt,
+			// baseJsFuncObject + value in baseJsFuncObject stash + []Value
+			expectedMem: SizeEmptyStruct + SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 	}
@@ -324,8 +324,8 @@ func TestClassFuncObjectMemUsage(t *testing.T) {
 					},
 				},
 			},
-			// baseJsFuncObject + value baseJsFuncObject in stash
-			expectedMem: SizeEmptyStruct + SizeInt,
+			// baseJsFuncObject + value baseJsFuncObject in stash + []Value
+			expectedMem: SizeEmptyStruct + SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
@@ -404,8 +404,8 @@ func TestMethodFuncObjectMemUsage(t *testing.T) {
 					},
 				},
 			},
-			// methodFuncObject + nil Object + value in baseJsFuncObject stash
-			expectedMem: SizeEmptyStruct + SizeEmptyStruct + SizeInt,
+			// methodFuncObject + nil Object + value in baseJsFuncObject stash + []Value
+			expectedMem: SizeEmptyStruct + SizeEmptyStruct + SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
@@ -464,8 +464,8 @@ func TestArrowFuncObjectMemUsage(t *testing.T) {
 					},
 				},
 			},
-			// arrowFuncObject + nil Object + value in baseJsFuncObject stash
-			expectedMem: SizeEmptyStruct + SizeEmptyStruct + SizeInt,
+			// arrowFuncObject + nil Object + value in baseJsFuncObject stash + []Value
+			expectedMem: SizeEmptyStruct + SizeEmptyStruct + SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{

memory_test.go

@@ -395,6 +395,18 @@ func TestMemCheck(t *testing.T) {
 			expectedSizeDiff: testNativeValueMemUsage +
 				(2 + SizeString), // "nv",
 		},
+		{
+			desc: "array.map function",
+			script: `x = [1,2,3,4,5]
+			x.map(_ => {
+				checkMem();
+				return "a"
+			})`,
+			// We are calculating memory of 4 "a" strings that are created while this map operation executes. Note that
+			// we don't calculate 5 because the last "a" string isn't created when the last checkMem() is called. This
+			// is expected because we only care that memory is being tracked within the map operation.
+			expectedSizeDiff: 4 * (1 + SizeString),
+		},
 	} {
 		t.Run(fmt.Sprintf(tc.desc), func(t *testing.T) {
 			memChecks := []uint64{}

runtime.go

@@ -569,6 +569,14 @@ func (r *Runtime) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
 		}
 	}
 
+	if r.vm.tmpValues != nil {
+		inc, err := valuesMemUsage(r.vm.tmpValues, ctx)
+		memUsage += inc
+		if err != nil {
+			return memUsage, err
+		}
+	}
+
 	return memUsage, nil
 }
 

runtime_test.go

@@ -3164,8 +3164,8 @@ func TestRuntimeMemUsage(t *testing.T) {
 				vm: &vm{stash: &stash{values: []Value{valueInt(99)}}},
 			},
 			memLimit: 100,
-			// stash value
-			expectedMem: SizeInt,
+			// stash value + []Value
+			expectedMem: SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
@@ -3174,8 +3174,8 @@ func TestRuntimeMemUsage(t *testing.T) {
 				vm: &vm{stack: []Value{valueInt(99)}},
 			},
 			memLimit: 100,
-			// stack value
-			expectedMem: SizeInt,
+			// stack value + []Value
+			expectedMem: SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 	}

vm.go

@@ -286,6 +286,7 @@ type vm struct {
 	pc           int
 	stack        valueStack
 	sp, sb, args int
+	tmpValues    []Value
 
 	stash     *stash
 	privEnv   *privateEnv
@@ -5591,6 +5592,11 @@ func (r *getPrivateRefId) exec(vm *vm) {
 }
 
 func (s valueStack) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
+	return valuesMemUsage(s, ctx)
+}
+
+func valuesMemUsage(s []Value, ctx *MemUsageContext) (memUsage uint64, err error) {
+	memUsage += SizeEmptySlice
 	for _, val := range s {
 		if val == nil {
 			continue

vm_test.go

@@ -395,30 +395,30 @@ func TestValueStackMemUsage(t *testing.T) {
 			name:        "should account for no memory usage given an empty value stack",
 			val:         []Value{},
 			memLimit:    100,
-			expectedMem: 0,
+			expectedMem: SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
 			name:        "should account for no memory usage given a value stack with nil",
 			val:         []Value{nil},
 			memLimit:    100,
-			expectedMem: 0,
+			expectedMem: SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
 			name:     "should account for each value given a non-empty value stack",
 			val:      []Value{valueInt(99)},
 			memLimit: 100,
-			// value
-			expectedMem: SizeInt,
+			// value + []Value
+			expectedMem: SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
 			name:     "should exit early given value stack over the memory limit",
 			val:      []Value{valueInt(99), valueInt(99), valueInt(99), valueInt(99)},
 			memLimit: 0,
-			// value
-			expectedMem: SizeInt,
+			// value + []Value
+			expectedMem: SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 	}
@@ -468,8 +468,8 @@ func TestVMContextMemUsage(t *testing.T) {
 		{
 			name: "should account for stash given a vmContext with non-empty stash",
 			val:  &vmContext{stash: &stash{values: []Value{valueInt(99)}}},
-			// vmContext overhead + stash value
-			expectedMem: SizeEmptyStruct + SizeInt,
+			// vmContext overhead + stash value + []Value
+			expectedMem: SizeEmptyStruct + SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
@@ -532,15 +532,15 @@ func TestStashMemUsage(t *testing.T) {
 		{
 			name: "should account for values given a stash with non-empty values",
 			val:  &stash{values: []Value{valueInt(99)}},
-			// value
-			expectedMem: SizeInt,
+			// value + []Value
+			expectedMem: SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 		{
 			name: "should account for outer given a stash with non-empty outer",
 			val:  &stash{outer: &stash{values: []Value{valueInt(99)}}},
-			// outer stash value
-			expectedMem: SizeInt,
+			// outer stash value + []Value
+			expectedMem: SizeInt + SizeEmptySlice,
 			errExpected: nil,
 		},
 	}
@@ -561,6 +561,60 @@ func TestStashMemUsage(t *testing.T) {
 	}
 }
 
+func TestTmpValuesMemUsage(t *testing.T) {
+	tests := []struct {
+		name        string
+		val         []Value
+		memLimit    uint64
+		expectedMem uint64
+		errExpected error
+	}{
+		{
+			name:        "should account for no memory usage given an empty tmpValues",
+			val:         []Value{},
+			memLimit:    100,
+			expectedMem: SizeEmptySlice,
+			errExpected: nil,
+		},
+		{
+			name:        "should account for no memory usage given a tmpValues with nil",
+			val:         []Value{nil},
+			memLimit:    100,
+			expectedMem: SizeEmptySlice,
+			errExpected: nil,
+		},
+		{
+			name:        "should account for each value given a non-empty tmpValues",
+			val:         []Value{valueInt(99), valueInt(99), valueInt(99), valueInt(99)},
+			memLimit:    100,
+			expectedMem: 4*SizeInt + SizeEmptySlice,
+			errExpected: nil,
+		},
+		{
+			name:        "should exit early given tmpValues over the memory limit",
+			val:         []Value{valueInt(99), valueInt(99), valueInt(99), valueInt(99)},
+			memLimit:    0,
+			expectedMem: SizeInt + SizeEmptySlice,
+			errExpected: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			total, err := valuesMemUsage(tc.val, NewMemUsageContext(New(), 100, tc.memLimit, 100, 100, 0.1, nil))
+			if err != tc.errExpected {
+				t.Fatalf("Unexpected error. Actual: %v Expected: %v", err, tc.errExpected)
+			}
+			if err != nil && tc.errExpected != nil && err.Error() != tc.errExpected.Error() {
+				t.Fatalf("Errors do not match. Actual: %v Expected: %v", err, tc.errExpected)
+			}
+			if total != tc.expectedMem {
+				t.Fatalf("Unexpected memory return. Actual: %v Expected: %v", total, tc.expectedMem)
+			}
+		})
+	}
+}
+
 func TestTickTracking(t *testing.T) {
 	tests := []struct {
 		name                      string