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,9 @@
 	}
 	r.vm.init()
 
+	r.tickMetrics = make(map[string]uint64)
+
 	r.SetRateLimiter(rate.NewLimiter(rate.Inf, maxInt))
 }
 
 func (r *Runtime) typeErrorResult(throw bool, args ...interface{}) {
@@ -1528,6 +1533,7 @@
 		vm.clearStack()
 	} else {
 		vm.stack = nil
+		vm.profileTicks()
 		vm.prg = nil
 		vm.setFuncName("")
 		r.leave()
@@ -3229,3 +3235,17 @@
 	}
 	return defProto
 }
+
+func (self *Runtime) EnableTickMetricTracking() {
+	self.tickMetricTrackingEnabled = true
+}
+
+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,88 @@ func TestStashMemUsage(t *testing.T) {
 		})
 	}
 }
+
+func TestTickTracking(t *testing.T) {
+	tests := []struct {
+		name                        string
+		script                      string
+		functionTickTrackingEnabled 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()
+			`,
+			functionTickTrackingEnabled: 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()
+			`,
+			functionTickTrackingEnabled: 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);
+			`,
+			functionTickTrackingEnabled: 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()
+			`,
+			functionTickTrackingEnabled: false,
+			expectedTickMetricsKeys:     []string{},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			vm := New()
+			if tc.functionTickTrackingEnabled {
+				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])
+				}
+			}
+		})
+	}
+}