feat: add early exit to global object in runtime

mongodb-forks · Jan 9, 2024 · 165184e · 165184e
1 parent a4d8244
commit 165184e
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/runtime.go b/runtime.go
@@ -501,6 +501,9 @@ func (r *Runtime) MemUsage(ctx *MemUsageContext) (memUsage uint64, newMemUsage u
 		if err != nil {
 			return memUsage, newMemUsage, err
 		}
+		if exceeded := ctx.MemUsageLimitExceeded(memUsage); exceeded {
+			return memUsage, newMemUsage, nil
+		}
 	}
 
 	if r.vm == nil {

diff --git a/runtime_test.go b/runtime_test.go
@@ -3087,14 +3087,34 @@ func TestRuntimeMemUsage(t *testing.T) {
 			name: "should account for globalObject given a runtime with non-empty globalObject",
 			val: &Runtime{
 				globalObject: &Object{
-					self: &baseObject{propNames: []unistring.String{"test"}, values: map[unistring.String]Value{"test": valueInt(99)}},
+					self: &baseObject{
+						propNames: []unistring.String{"test0", "test1", "test2"},
+						values:    map[unistring.String]Value{"test0": valueInt(99), "test1": valueInt(99), "test2": valueInt(99)},
+					},
 				},
 			},
 			memLimit: 100,
 			// baseObject overhead + key/value pair with string overhead
-			expectedMem: SizeEmptyStruct + (4 + SizeString + SizeInt),
+			expectedMem: SizeEmptyStruct + (5+SizeString+SizeInt)*3,
 			// baseObject overhead + key/value pair with string overhead
-			expectedNewMem: SizeEmptyStruct + (4 + SizeString + SizeInt),
+			expectedNewMem: SizeEmptyStruct + (5+SizeString+SizeInt)*3,
+			errExpected:    nil,
+		},
+		{
+			name: "should exit early given a runtime with non-empty globalObject exceeding the memory limit",
+			val: &Runtime{
+				globalObject: &Object{
+					self: &baseObject{
+						propNames: []unistring.String{"test0", "test1", "test2"},
+						values:    map[unistring.String]Value{"test0": valueInt(99), "test1": valueInt(99), "test2": valueInt(99)},
+					},
+				},
+			},
+			memLimit: 0,
+			// baseObject overhead + key/value pair with string overhead
+			expectedMem: SizeEmptyStruct + (5 + SizeString + SizeInt),
+			// baseObject overhead + key/value pair with string overhead
+			expectedNewMem: SizeEmptyStruct + (5 + SizeString + SizeInt),
 			errExpected:    nil,
 		},
 		{