BAAS-22402: add early exits when possible memusage (#87)

mongodb-forks · Aug 30, 2023 · b504431 · b504431
1 parent f5cc0e5
commit b504431
Show file tree

Hide file tree

Showing 12 changed files with 341 additions and 17 deletions.
diff --git a/.evg.yml b/.evg.yml
@@ -25,7 +25,11 @@ tasks:
             export PATH=$ROOT_DIR/go/bin:$PATH
 
             cd goja
-            go test -v ./...
+            go test -v ./... >> ${workdir}/goja.suite
+      - command: gotest.parse_files
+        params:
+          files:
+            - "${workdir}/goja.suite"
 
 buildvariants:
   - name: linux-64

diff --git a/array.go b/array.go
@@ -554,7 +554,6 @@ func toIdx(v valueInt) uint32 {
 }
 
 var (
-	errMemUsageExceedsLimitNil     = errors.New("error checking mem usage limit")
 	errArrayLenExceedsThresholdNil = errors.New("error checking array len threshold")
 )
 
@@ -582,11 +581,8 @@ func (a *arrayObject) estimateMemUsage(ctx *MemUsageContext) (uint64, error) {
 		if err != nil {
 			return runningEstimate, err
 		}
-		if ctx.MemUsageExceedsLimit == nil {
-			return runningEstimate, errMemUsageExceedsLimitNil
-		}
-		if ctx.MemUsageExceedsLimit(total) {
-			return runningEstimate, nil
+		if exceeded := ctx.MemUsageLimitExceeded(total); exceeded {
+			return total, nil
 		}
 	}
 
@@ -634,12 +630,9 @@ func (a *arrayObject) MemUsage(ctx *MemUsageContext) (uint64, error) {
 		if err != nil {
 			return total, err
 		}
-		if ctx.MemUsageExceedsLimit == nil {
-			return total, errMemUsageExceedsLimitNil
-		}
 		// This is an early exit in case we reach the mem usage
 		// limit before we get to scan the whole array.
-		if ctx.MemUsageExceedsLimit(total) {
+		if exceeded := ctx.MemUsageLimitExceeded(total); exceeded {
 			return total, nil
 		}
 	}

diff --git a/array_sparse.go b/array_sparse.go
@@ -507,6 +507,9 @@ func (a *sparseArrayObject) MemUsage(ctx *MemUsageContext) (uint64, error) {
 			if err != nil {
 				return total, err
 			}
+			if exceeded := ctx.MemUsageLimitExceeded(total); exceeded {
+				return total, nil
+			}
 		}
 	}
 

diff --git a/array_sparse_test.go b/array_sparse_test.go
@@ -262,3 +262,77 @@ func TestSparseArrayExportToSlice(t *testing.T) {
 		}
 	}
 }
+
+func TestSparseArrayObjectMemUsage(t *testing.T) {
+	tests := []struct {
+		name        string
+		mu          *MemUsageContext
+		sao         *sparseArrayObject
+		expected    uint64
+		errExpected error
+	}{
+		{
+			name: "mem below threshold",
+			mu:   NewMemUsageContext(New(), 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			sao: &sparseArrayObject{
+				items: []sparseArrayItem{
+					{
+						idx:   1,
+						value: New()._newString(newStringValue("key"), nil),
+					},
+				},
+			},
+			expected:    45,
+			errExpected: nil,
+		},
+		{
+			name: "mem way above threshold returns first crossing of threshold",
+			mu:   NewMemUsageContext(New(), 88, 100, 50, 50, TestNativeMemUsageChecker{}),
+			sao: &sparseArrayObject{
+				items: []sparseArrayItem{
+					{
+						idx:   1,
+						value: New()._newString(newStringValue("key"), nil),
+					},
+					{
+						idx:   2,
+						value: New()._newString(newStringValue("key1"), nil),
+					},
+					{
+						idx:   3,
+						value: New()._newString(newStringValue("key2"), nil),
+					},
+					{
+						idx:   4,
+						value: New()._newString(newStringValue("key3"), nil),
+					},
+					{
+						idx:   5,
+						value: New()._newString(newStringValue("key4"), nil),
+					},
+					{
+						idx:   6,
+						value: New()._newString(newStringValue("key5"), nil),
+					},
+				},
+			},
+			expected:    127,
+			errExpected: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			total, err := tc.sao.MemUsage(tc.mu)
+			if err == nil && tc.errExpected != nil || err != nil && tc.errExpected == nil {
+				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.expected {
+				t.Fatalf("Unexpected memory return. Actual: %v Expected: %v", total, tc.expected)
+			}
+		})
+	}
+}
diff --git a/array_test.go b/array_test.go
@@ -131,3 +131,79 @@ func BenchmarkArraySetEmpty(b *testing.B) {
 		a.self.setOwnIdx(0, valueTrue, true)
 	}
 }
+
+func TestArrayObjectMemUsage(t *testing.T) {
+	tests := []struct {
+		name        string
+		mu          *MemUsageContext
+		ao          *arrayObject
+		expected    uint64
+		errExpected error
+	}{
+		{
+			name: "mem below threshold",
+			mu:   NewMemUsageContext(New(), 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			ao: &arrayObject{
+				values: []Value{
+					New()._newString(newStringValue("key"), nil),
+				},
+			},
+			expected:    41,
+			errExpected: nil,
+		},
+		{
+			name: "mem way above threshold returns first crossing of threshold",
+			mu:   NewMemUsageContext(New(), 88, 100, 50, 50, TestNativeMemUsageChecker{}),
+			ao: &arrayObject{
+				values: []Value{
+					New()._newString(newStringValue("key"), nil),
+					New()._newString(newStringValue("key1"), nil),
+					New()._newString(newStringValue("key2"), nil),
+					New()._newString(newStringValue("key3"), nil),
+					New()._newString(newStringValue("key4"), nil),
+					New()._newString(newStringValue("key5"), nil),
+				},
+			},
+			expected:    119,
+			errExpected: nil,
+		},
+		{
+			name: "array limit function undefined throws error",
+			mu: &MemUsageContext{
+				visitTracker: visitTracker{
+					objsVisited:    map[objectImpl]bool{},
+					stashesVisited: map[*stash]bool{}},
+				depthTracker: &depthTracker{
+					curDepth: 0,
+					maxDepth: 50,
+				},
+				NativeMemUsageChecker: &TestNativeMemUsageChecker{},
+				memoryLimit:           50,
+				ObjectPropsLenExceedsThreshold: func(objPropsLen int) bool {
+					// number of obj props beyond which we should estimate mem usage
+					return objPropsLen > 50
+				},
+			},
+			ao: &arrayObject{
+				values: []Value{New()._newString(newStringValue("key"), nil)},
+			},
+			expected:    16,
+			errExpected: errArrayLenExceedsThresholdNil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			total, err := tc.ao.MemUsage(tc.mu)
+			if err == nil && tc.errExpected != nil || err != nil && tc.errExpected == nil {
+				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.expected {
+				t.Fatalf("Unexpected memory return. Actual: %v Expected: %v", total, tc.expected)
+			}
+		})
+	}
+}
diff --git a/builtin_map.go b/builtin_map.go
@@ -129,6 +129,9 @@ func (mo *mapObject) MemUsage(ctx *MemUsageContext) (uint64, error) {
 				return total, err
 			}
 		}
+		if exceeded := ctx.MemUsageLimitExceeded(total); exceeded {
+			return total, nil
+		}
 	}
 
 	ctx.Ascend()

diff --git a/builtin_map_test.go b/builtin_map_test.go
@@ -242,3 +242,73 @@ func BenchmarkMapDeleteJS(b *testing.B) {
 		}
 	}
 }
+
+func TestMapObjectMemUsage(t *testing.T) {
+	tests := []struct {
+		name        string
+		mu          *MemUsageContext
+		mo          *mapObject
+		expected    uint64
+		errExpected error
+	}{
+		{
+			name: "mem below threshold",
+			mu:   NewMemUsageContext(New(), 88, 5000, 50, 50, TestNativeMemUsageChecker{}),
+			mo: &mapObject{
+				m: &orderedMap{
+					hashTable: map[uint64]*mapEntry{
+						1: {
+							key:   New()._newString(newStringValue("key"), nil),
+							value: New()._newString(newStringValue("value"), nil),
+						},
+					},
+				},
+			},
+			expected:    60,
+			errExpected: nil,
+		},
+		{
+			name: "mem way above threshold returns first crossing of threshold",
+			mu:   NewMemUsageContext(New(), 88, 100, 50, 50, TestNativeMemUsageChecker{}),
+			mo: &mapObject{
+				m: &orderedMap{
+					hashTable: map[uint64]*mapEntry{
+						1: {
+							key:   New()._newString(newStringValue("key"), nil),
+							value: New()._newString(newStringValue("value"), nil),
+						},
+						2: {
+							key:   New()._newString(newStringValue("key"), nil),
+							value: New()._newString(newStringValue("value"), nil),
+						},
+						3: {
+							key:   New()._newString(newStringValue("key"), nil),
+							value: New()._newString(newStringValue("value"), nil),
+						},
+						4: {
+							key:   New()._newString(newStringValue("key"), nil),
+							value: New()._newString(newStringValue("value"), nil),
+						},
+					},
+				},
+			},
+			expected:    112,
+			errExpected: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			total, err := tc.mo.MemUsage(tc.mu)
+			if err == nil && tc.errExpected != nil || err != nil && tc.errExpected == nil {
+				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.expected {
+				t.Fatalf("Unexpected memory return. Actual: %v Expected: %v", total, tc.expected)
+			}
+		})
+	}
+}
diff --git a/compiler.go b/compiler.go
@@ -469,6 +469,9 @@ func (p *Program) MemUsage(ctx *MemUsageContext) (uint64, error) {
 		if err != nil {
 			return total, err
 		}
+		if exceeded := ctx.MemUsageLimitExceeded(total); exceeded {
+			return total, nil
+		}
 	}
 
 	return total, nil

diff --git a/compiler_test.go b/compiler_test.go
@@ -4,6 +4,7 @@ import (
 	"os"
 	"sync"
 	"testing"
+	"time"
 )
 
 const TESTLIB = `
@@ -5645,3 +5646,59 @@ func BenchmarkCompile(b *testing.B) {
 		}
 	}
 }
+
+func TestProgramMemUsage(t *testing.T) {
+	tests := []struct {
+		name        string
+		mu          *MemUsageContext
+		p           *Program
+		expected    uint64
+		errExpected error
+	}{
+		{
+			name: "mem below threshold",
+			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, TestNativeMemUsageChecker{}),
+			p: &Program{
+				values: []Value{
+					New().newDateObject(time.Now(), true, nil),
+				},
+			},
+			expected:    16,
+			errExpected: nil,
+		},
+		{
+			name: "mem way above threshold returns first crossing of threshold",
+			mu:   NewMemUsageContext(New(), 88, 50, 50, 50, TestNativeMemUsageChecker{}),
+			p: &Program{
+				values: []Value{
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+					New().newDateObject(time.Now(), true, nil),
+				},
+			},
+			expected:    64,
+			errExpected: nil,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			total, err := tc.p.MemUsage(tc.mu)
+			if err == nil && tc.errExpected != nil || err != nil && tc.errExpected == nil {
+				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.expected {
+				t.Fatalf("Unexpected memory return. Actual: %v Expected: %v", total, tc.expected)
+			}
+		})
+	}
+}