BAAS-11692: add mem usage for setObject and func objects (#121)

builtin_set.go

@@ -344,3 +344,95 @@ func (r *Runtime) getSet() *Object {
 	}
 	return ret
 }
+
+// estimateMemUsage helps calculating mem usage for large objects.
+// It will sample the object and use those samples to estimate the
+// mem usage.
+func (so *setObject) estimateMemUsage(ctx *MemUsageContext) (estimate uint64, err error) {
+	var samplesVisited, memUsage uint64
+	totalItems := so.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 := so.m.iterFirst; item != nil; item = item.iterNext {
+		if samplesVisited%uint64(sampleSize) != 0 {
+			continue
+		}
+		samplesVisited += 1
+
+		// We still want to account for both key and value if we return a non-zero value on error.
+		// This could otherwise skew the estimate when in reality key/value pairs contribute to
+		// mem usage together.
+		inc, incErr := item.key.MemUsage(ctx)
+		memUsage += inc
+		inc, valErr := item.value.MemUsage(ctx)
+		memUsage += inc
+		if valErr != nil {
+			return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), valErr
+		}
+		if incErr != nil {
+			return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), incErr
+		}
+	}
+
+	return computeMemUsageEstimate(memUsage, samplesVisited, totalItems), nil
+}
+
+func (so *setObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
+	if so == nil || ctx.IsObjVisited(so) {
+		return SizeEmptyStruct, nil
+	}
+	ctx.VisitObj(so)
+
+	if err := ctx.Descend(); err != nil {
+		return memUsage, err
+	}
+
+	memUsage, err = so.baseObject.MemUsage(ctx)
+	if err != nil {
+		return memUsage, err
+	}
+
+	if so.m != nil {
+		if ctx.ObjectPropsLenExceedsThreshold(so.m.size) {
+			inc, err := so.estimateMemUsage(ctx)
+			memUsage += inc
+			if err != nil {
+				return memUsage, err
+			}
+		} else {
+			for _, entry := range so.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.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
+				}
+			}
+		}
+	}
+
+	ctx.Ascend()
+
+	return memUsage, nil
+}

builtin_set_test.go

@@ -192,3 +192,109 @@ func TestSetHasFloatVsInt(t *testing.T) {
 
 	testScript(SCRIPT, valueTrue, t)
 }
+
+func TestSetObjectMemUsage(t *testing.T) {
+	vm := New()
+
+	tests := []struct {
+		name        string
+		mu          *MemUsageContext
+		so          *setObject
+		expectedMem uint64
+		errExpected error
+	}{
+		{
+			name: "mem below threshold",
+			mu:   NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			so: &setObject{
+				m: &orderedMap{
+					hashTable: map[uint64]*mapEntry{
+						1: {
+							key:   vm.ToValue("key"),
+							value: vm.ToValue("value"),
+						},
+					},
+				},
+			},
+			// baseObject + (len(key) + overhead)  + (len(value) + overhead)
+			expectedMem: SizeEmptyStruct + (3 + SizeString) + (5 + SizeString),
+			errExpected: nil,
+		},
+		{
+			name:        "mem is SizeEmptyStruct given a nil map object",
+			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			so:          nil,
+			expectedMem: SizeEmptyStruct,
+			errExpected: nil,
+		},
+		{
+			name: "mem way above threshold returns first crossing of threshold",
+			mu:   NewMemUsageContext(vm, 88, 100, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			so: &setObject{
+				m: &orderedMap{
+					hashTable: map[uint64]*mapEntry{
+						1: {
+							key:   vm.ToValue("key"),
+							value: vm.ToValue("value"),
+						},
+						2: {
+							key:   vm.ToValue("key"),
+							value: vm.ToValue("value"),
+						},
+						3: {
+							key:   vm.ToValue("key"),
+							value: vm.ToValue("value"),
+						},
+						4: {
+							key:   vm.ToValue("key"),
+							value: vm.ToValue("value"),
+						},
+					},
+				},
+			},
+			// baseObject
+			expectedMem: SizeEmptyStruct +
+				// len(key) + overhead (we reach the limit after 3)
+				(3+SizeString)*3 +
+				// len(value) + overhead (we reach the limit after 3)
+				(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{}),
+			so: &setObject{
+				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,
+		},
+		{
+			name:        "mem is SizeEmptyStruct given a nil orderedMap object",
+			mu:          NewMemUsageContext(vm, 88, 5000, 50, 50, 0.1, TestNativeMemUsageChecker{}),
+			so:          &setObject{},
+			expectedMem: SizeEmptyStruct,
+			errExpected: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			total, err := tc.so.MemUsage(tc.mu)
+			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.go

@@ -505,13 +505,22 @@ func (f *funcObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error)
 	}
 	ctx.VisitObj(f)
 
-	memUsage, err = f.baseObject.MemUsage(ctx)
+	return f.baseJsFuncObject.MemUsage(ctx)
+}
+
+func (b *baseJsFuncObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
+	if b == nil || ctx.IsObjVisited(b) {
+		return SizeEmptyStruct, nil
+	}
+	ctx.VisitObj(b)
+
+	memUsage, err = b.baseFuncObject.MemUsage(ctx)
 	if err != nil {
 		return memUsage, err
 	}
 
-	if f.stash != nil {
-		inc, err := f.stash.MemUsage(ctx)
+	if b.stash != nil {
+		inc, err := b.stash.MemUsage(ctx)
 		memUsage += inc
 		if err != nil {
 			return memUsage, err
@@ -520,3 +529,90 @@ func (f *funcObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error)
 
 	return memUsage, nil
 }
+
+func (c *classFuncObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
+	if c == nil || ctx.IsObjVisited(c) {
+		return SizeEmptyStruct, nil
+	}
+	ctx.VisitObj(c)
+
+	memUsage, err = c.baseJsFuncObject.MemUsage(ctx)
+	if err != nil {
+		return memUsage, err
+	}
+
+	if c.initFields != nil {
+		inc, err := c.initFields.MemUsage(ctx)
+		memUsage += inc
+		if err != nil {
+			return memUsage, err
+		}
+	}
+
+	for _, v := range c.computedKeys {
+		inc, err := v.MemUsage(ctx)
+		memUsage += inc
+		if err != nil {
+			return memUsage, err
+		}
+	}
+
+	for _, v := range c.privateMethods {
+		inc, err := v.MemUsage(ctx)
+		memUsage += inc
+		if err != nil {
+			return memUsage, err
+		}
+	}
+
+	return memUsage, err
+}
+
+func (m *methodFuncObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
+	if m == nil || ctx.IsObjVisited(m) {
+		return SizeEmptyStruct, nil
+	}
+	ctx.VisitObj(m)
+
+	memUsage, err = m.baseJsFuncObject.MemUsage(ctx)
+	if err != nil {
+		return memUsage, err
+	}
+
+	inc, err := m.homeObject.MemUsage(ctx)
+	memUsage += inc
+	if err != nil {
+		return memUsage, err
+	}
+
+	return memUsage, err
+}
+
+func (a *arrowFuncObject) MemUsage(ctx *MemUsageContext) (memUsage uint64, err error) {
+	if a == nil || ctx.IsObjVisited(a) {
+		return SizeEmptyStruct, nil
+	}
+	ctx.VisitObj(a)
+
+	memUsage, err = a.baseJsFuncObject.MemUsage(ctx)
+	if err != nil {
+		return memUsage, err
+	}
+
+	inc, err := a.funcObj.MemUsage(ctx)
+	memUsage += inc
+	if err != nil {
+		return memUsage, err
+	}
+
+	if a.newTarget != nil {
+		inc, err = a.newTarget.MemUsage(ctx)
+		memUsage += inc
+		if err != nil {
+			return memUsage, err
+		}
+
+	}
+
+	return memUsage, err
+}