BAAS-35403 use func over bool for runtime shouldTrackMaxMemOnStack

runtime.go

@@ -211,9 +211,9 @@ type Runtime struct {
 	tickMetricTrackingEnabled bool
 	tickMetrics               map[string]uint64
 
-	// context used for tracking object with largest memory in the value stack. Must be non-nil if shouldTrackMaxMemOnStack is true
+	// context used for tracking object with largest memory in the value stack. Must be non-nil if shouldTrackMaxMemOnStack is defined
 	stackMemUsageContext     *MemUsageContext
-	shouldTrackMaxMemOnStack bool
+	shouldTrackMaxMemOnStack func(funcName string) bool
 	maxStackObjectMem        uint64
 }
 
@@ -1410,9 +1410,9 @@ func New() *Runtime {
 
 // NewWithContext creates an instance of a Javascript runtime that can be used to run code. Multiple instances may be created and
 // used simultaneously, however it is not possible to pass JS values across runtimes. The 'shouldTrackMaxMemOnStack' parameter is a
-// safety net to track the maximum memory of objects on the vm's goja.Value stack as there will always be a Function or Arguments
-// object with memory usage totalling all referenced objects within the function. This covers edge cases missed by the memory poller.
-func NewWithContext(ctx context.Context, shouldTrackMaxMemOnStack bool, stackMemUsageContext *MemUsageContext) *Runtime {
+// safety net to track the maximum memory of objects on the vm's goja.Value stack for certain functions.
+// This covers edge cases in async functions missed by the memory poller but does not account for non-async functions.
+func NewWithContext(ctx context.Context, shouldTrackMaxMemOnStack func(funcName string) bool, stackMemUsageContext *MemUsageContext) *Runtime {
 	r := &Runtime{ctx: ctx, shouldTrackMaxMemOnStack: shouldTrackMaxMemOnStack, stackMemUsageContext: stackMemUsageContext}
 	r.init()
 	return r

vm.go

@@ -856,12 +856,16 @@ func (vm *vm) push(v Value) {
 	vm.stack[vm.sp] = v
 	vm.sp++
 
-	shouldTrackMaxMemOnStack := (vm.r != nil && vm.r.shouldTrackMaxMemOnStack)
-
-	if !shouldTrackMaxMemOnStack || v == nil || !v.IsObject() {
+	if vm.r == nil || v == nil {
+		return
+	}
+	if vm.r.shouldTrackMaxMemOnStack == nil {
 		return
 	}
 
+	if !vm.r.shouldTrackMaxMemOnStack(vm.funcName.String()) {
+		return
+	}
 	// clear the visitTracker so mem check is forced on paths that contain objects with updated mem usage
 	vm.r.stackMemUsageContext.visitTracker = visitTracker{objsVisited: make(map[objectImpl]struct{}), stashesVisited: make(map[*stash]struct{})}
 

vm_test.go

@@ -827,7 +827,9 @@ func BenchmarkStackPushMemTracking(b *testing.B) {
 		)
 		b.ResetTimer()
 		b.Run(bm.name, func(b *testing.B) {
-			vmRuntime.shouldTrackMaxMemOnStack = bm.shouldTrackMaxMemOnStack
+			vmRuntime.shouldTrackMaxMemOnStack = func(_ string) bool {
+				return bm.shouldTrackMaxMemOnStack
+			}
 
 			for i := 0; i < b.N; i++ {
 				vmRuntime.vm.push(res)
@@ -879,7 +881,9 @@ func BenchmarkVmMemTracking(b *testing.B) {
 		)
 		b.ResetTimer()
 		b.Run(bm.name, func(b *testing.B) {
-			vmRuntime.shouldTrackMaxMemOnStack = bm.shouldTrackMaxMemOnStack
+			vmRuntime.shouldTrackMaxMemOnStack = func(_ string) bool {
+				return bm.shouldTrackMaxMemOnStack
+			}
 
 			for i := 0; i < b.N; i++ {
 				res, err := vmRuntime.RunProgram(prg)