From 06a8a6724b36342b9e4d4e5d7d3a2b9ca734a597 Mon Sep 17 00:00:00 2001
From: Calvin Nix <6693710+Calvinnix@users.noreply.github.com>
Date: Mon, 11 Mar 2024 11:56:36 -0400
Subject: [PATCH] BAAS-28597: Add counter for limiter wait calls (#117)

---
 runtime.go | 13 +++++++++++++
 vm.go      | 42 ++++++++++++++++++++++++++++++++----------
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/runtime.go b/runtime.go
index a8a3ecc9..726866f3 100644
--- a/runtime.go
+++ b/runtime.go
@@ -205,6 +205,9 @@ type Runtime struct {
 	limiterTicksLeft int
 	ticks            uint64
 
+	limiterWaitCount     int64
+	limiterWaitTotalTime time.Duration
+
 	tickMetricTrackingEnabled bool
 	tickMetrics               map[string]uint64
 }
@@ -213,6 +216,16 @@ func (self *Runtime) Ticks() uint64 {
 	return self.ticks
 }
 
+// LimiterWaitCount tracks the amount of times the rate limiter throttles the execution.
+func (self *Runtime) LimiterWaitCount() int64 {
+	return self.limiterWaitCount
+}
+
+// LimiterWaitTotalTime tracks the total amount of time that the execution was throttled.
+func (self *Runtime) LimiterWaitTotalTime() time.Duration {
+	return self.limiterWaitTotalTime
+}
+
 // SetStackTraceLimit sets an upper limit to the number of stack frames that
 // goja will use when formatting an error's stack trace. By default, the limit
 // is 10. This is consistent with V8 and SpiderMonkey.
diff --git a/vm.go b/vm.go
index cc974c78..e1661fbb 100644
--- a/vm.go
+++ b/vm.go
@@ -9,6 +9,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/dop251/goja/unistring"
 )
@@ -617,17 +618,26 @@ func (vm *vm) run() {
 				ctx = context.Background()
 			}
 
-			if waitErr := vm.r.limiter.WaitN(ctx, vm.r.limiterTicksLeft); waitErr != nil {
-				if vm.r.vm.ctx == nil {
-					panic(waitErr)
-				}
-				if ctxErr := vm.r.vm.ctx.Err(); ctxErr != nil {
-					panic(ctxErr)
-				}
-				if strings.Contains(waitErr.Error(), "would exceed") {
-					panic(context.DeadlineExceeded)
+			select {
+			case <-ctx.Done():
+				panic(ctx.Err())
+			default:
+			}
+
+			now := time.Now()
+			r := vm.r.limiter.ReserveN(now, vm.r.limiterTicksLeft)
+			if !r.OK() {
+				panic(context.DeadlineExceeded)
+			}
+
+			if delay := r.DelayFrom(now); delay > 0 {
+				vm.r.limiterWaitCount++
+				vm.r.limiterWaitTotalTime += delay
+
+				if err := sleep(ctx, delay); err != nil {
+					r.Cancel()
+					panic(err)
 				}
-				panic(waitErr)
 			}
 		}
 
@@ -5606,3 +5616,15 @@ func (s valueStack) MemUsage(ctx *MemUsageContext) (memUsage uint64, newMemUsage
 
 	return memUsage, newMemUsage, nil
 }
+
+// sleep is equivalent to time.Sleep, but is interruptible via context cancelation.
+func sleep(ctx context.Context, duration time.Duration) error {
+	timer := time.NewTimer(duration)
+	defer timer.Stop()
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}