BAAS-28253: Add optional tick tracking at the function level

runtime.go

@@ -204,6 +204,9 @@
 	limiter          *rate.Limiter
 	limiterTicksLeft int
 	ticks            uint64
+
+	tickMetricTrackingEnabled bool
+	tickMetrics               map[string]uint64
 }
 
 func (self *Runtime) Ticks() uint64 {
@@ -489,7 +492,13 @@
 	}
 	r.vm.init()
 
+	// Initializing tickMetrics by default to avoid complexities when trying to enable tick tracking with
+	// vm.EnableFunctionTickTracking(). The problem here is we need to have the vm object in order to enable tick tracking,
+	// but we check whether tick tracking is enabled when we create the vm object.
+	// It is simpler, and not costly, to just have this map always initialized.
+	r.tickMetrics = make(map[string]uint64)
+
 	r.SetRateLimiter(rate.NewLimiter(rate.Inf, maxInt))
 }
 
 func (r *Runtime) typeErrorResult(throw bool, args ...interface{}) {
@@ -1528,6 +1537,7 @@
 		vm.clearStack()
 	} else {
 		vm.stack = nil
+		vm.profileTicks()
 		vm.prg = nil
 		vm.setFuncName("")
 		r.leave()
@@ -3229,3 +3239,21 @@
 	}
 	return defProto
 }
+
+// EnableTickMetricTracking sets tickMetricTrackingEnabled to true. tickMetricTrackingEnabled is checked during every
+// profileTick call so this would effectively enable tick tracking immediately.
+func (self *Runtime) EnableTickMetricTracking() {
+	self.tickMetricTrackingEnabled = true
+}
+
+// DisableTickMetricTracking sets tickMetricTrackingEnabled to false. tickMetricTrackingEnabled is checked during every
+// profileTick call so this would effectively disable tick tracking immediately.
+func (self *Runtime) DisableTickMetricTracking() {
+	self.tickMetricTrackingEnabled = false
+}
+
+// TickMetrics returns a map of ticks used per function.
+// This function is not thread-safe and should only be called at the end of the function execution.
+func (self *Runtime) TickMetrics() map[string]uint64 {
+	return self.tickMetrics
+}

vm.go

@@ -306,12 +306,30 @@ type vm struct {
 	interrupted   uint32
 	interruptVal  interface{}
 	interruptLock sync.Mutex
+
+	lastFunctionTicks uint64
 }
 
 type instruction interface {
 	exec(*vm)
 }
 
+// profileTicks tracks the ticks for the current Program
+func (vm *vm) profileTicks() {
+	if !vm.r.tickMetricTrackingEnabled {
+		return
+	}
+	currentTicks := vm.r.Ticks()
+	if vm.prg != nil {
+		function := string(vm.prg.funcName)
+		if vm.prg.src != nil {
+			function = vm.prg.src.Name() + "::" + function
+		}
+		vm.r.tickMetrics[function] += currentTicks - vm.lastFunctionTicks
+	}
+	vm.lastFunctionTicks = currentTicks
+}
+
 func (vm *vm) getFuncName() unistring.String {
 	vm.funcNameLock.RLock()
 	s := vm.funcName
@@ -857,6 +875,7 @@ func (vm *vm) saveCtx(ctx *vmContext) {
 }
 
 func (vm *vm) pushCtx() {
+	vm.profileTicks()
 	if len(vm.callStack) > vm.maxCallStackSize {
 		panic(rangeError("Maximum call stack size exceeded"))
 	}
@@ -866,6 +885,7 @@ func (vm *vm) pushCtx() {
 }
 
 func (vm *vm) restoreCtx(ctx *vmContext) {
+	vm.profileTicks()
 	vm.prg, vm.stash, vm.privEnv, vm.newTarget, vm.result, vm.pc, vm.sb, vm.args =
 		ctx.prg, ctx.stash, ctx.privEnv, ctx.newTarget, ctx.result, ctx.pc, ctx.sb, ctx.args
 

vm_test.go

@@ -594,3 +594,146 @@ func TestStashMemUsage(t *testing.T) {
 		})
 	}
 }
+
+func TestTickTracking(t *testing.T) {
+	tests := []struct {
+		name                      string
+		script                    string
+		tickMetricTrackingEnabled bool
+		expectedTickMetricsKeys   []string
+	}{
+		{
+			name: "should track looping function ticks when tick tracking is enabled",
+			script: `
+				function f() {
+					for (var i = 0; i < 100; i++) {
+					}
+				}
+				f()
+			`,
+			tickMetricTrackingEnabled: true,
+			expectedTickMetricsKeys:   []string{"test.js::", "test.js::f"},
+		},
+		{
+			name: "should track larger looping function ticks when tick tracking is enabled",
+			script: `
+				function f() {
+					for (var i = 0; i < 1000; i++) {
+					}
+				}
+				f()
+			`,
+			tickMetricTrackingEnabled: true,
+			expectedTickMetricsKeys:   []string{"test.js::", "test.js::f"},
+		},
+		{
+			name: "should track fib function ticks when tick tracking is enabled",
+			script: `
+				function fib(n) {
+					if (n < 2) return n;
+					return fib(n - 2) + fib(n - 1);
+				}
+				fib(35);
+			`,
+			tickMetricTrackingEnabled: true,
+			expectedTickMetricsKeys:   []string{"test.js::", "test.js::fib"},
+		},
+		{
+			name: "should not track function ticks when tick tracking is disabled",
+			script: `
+				function f() {
+					for (var i = 0; i < 100; i++) {
+					}
+				}
+				f()
+			`,
+			tickMetricTrackingEnabled: false,
+			expectedTickMetricsKeys:   []string{},
+		},
+		{
+			name: "should track ticks from a class method",
+			script: `
+				class Car {
+					constructor(brand) {
+						this.brand = brand;
+					}
+
+					drive() {}
+					honk() {}
+				}
+				const car = new Car("Tesla")
+				car.drive()
+				car.honk()
+			`,
+			tickMetricTrackingEnabled: true,
+			expectedTickMetricsKeys:   []string{"test.js::", "test.js::drive", "test.js::honk"},
+		},
+		{
+			name: "should track ticks from nested functions",
+			script: `
+				function outerFunction() {
+					function firstNestedFunction() {
+						function secondNestedFunction() {
+							function thirdNestedFunction() {}
+							thirdNestedFunction();
+						}
+						secondNestedFunction();
+					}
+					firstNestedFunction();
+				}
+				outerFunction();
+			`,
+			tickMetricTrackingEnabled: true,
+			expectedTickMetricsKeys: []string{
+				"test.js::",
+				"test.js::outerFunction",
+				"test.js::firstNestedFunction",
+				"test.js::secondNestedFunction",
+				"test.js::thirdNestedFunction",
+			},
+		},
+		{
+			name: "should track ticks even when a function throws an error",
+			script: `
+				function explode() {
+					throw new Error("boom");
+				}
+				try {
+					explode()
+				} catch (err) {}
+			`,
+			tickMetricTrackingEnabled: true,
+			expectedTickMetricsKeys: []string{
+				"test.js::",
+				"test.js::explode",
+			},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			vm := New()
+			if tc.tickMetricTrackingEnabled {
+				vm.EnableTickMetricTracking()
+			}
+
+			prg := MustCompile("test.js", tc.script, false)
+			_, err := vm.RunProgram(prg)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			actualTickMetrics := vm.TickMetrics()
+
+			if len(actualTickMetrics) != len(tc.expectedTickMetricsKeys) {
+				t.Fatalf("Unexpected tickMetrics length: Actual: %v Expected length: %v", actualTickMetrics, len(tc.expectedTickMetricsKeys))
+			}
+
+			for _, key := range tc.expectedTickMetricsKeys {
+				if actualTickMetrics[key] <= 0 {
+					t.Fatalf("Unexpected tickMetrics for key: %v. Expected a positive tick value but received: %v", key, actualTickMetrics[key])
+				}
+			}
+		})
+	}
+}