revert: fixed window ratelimiting

apps/agent/pkg/clock/real_clock.go

@@ -2,31 +2,15 @@ package clock
 
 import "time"
 
-type TestClock struct {
-	now time.Time
+type RealClock struct {
 }
 
-func NewTestClock(now ...time.Time) *TestClock {
-	if len(now) == 0 {
-		now = append(now, time.Now())
-	}
-	return &TestClock{now: now[0]}
+func New() *RealClock {
+	return &RealClock{}
 }
 
-var _ Clock = &TestClock{}
+var _ Clock = &RealClock{}
 
-func (c *TestClock) Now() time.Time {
-	return c.now
-}
-
-// Tick advances the clock by the given duration and returns the new time.
-func (c *TestClock) Tick(d time.Duration) time.Time {
-	c.now = c.now.Add(d)
-	return c.now
-}
-
-// Set sets the clock to the given time and returns the new time.
-func (c *TestClock) Set(t time.Time) time.Time {
-	c.now = t
-	return c.now
+func (c *RealClock) Now() time.Time {
+	return time.Now()
 }

apps/agent/pkg/clock/test_clock.go

@@ -2,15 +2,31 @@ package clock
 
 import "time"
 
-type RealClock struct {
+type TestClock struct {
+	now time.Time
 }
 
-func New() *RealClock {
-	return &RealClock{}
+func NewTestClock(now ...time.Time) *TestClock {
+	if len(now) == 0 {
+		now = append(now, time.Now())
+	}
+	return &TestClock{now: now[0]}
 }
 
-var _ Clock = &RealClock{}
+var _ Clock = &TestClock{}
 
-func (c *RealClock) Now() time.Time {
-	return time.Now()
+func (c *TestClock) Now() time.Time {
+	return c.now
+}
+
+// Tick advances the clock by the given duration and returns the new time.
+func (c *TestClock) Tick(d time.Duration) time.Time {
+	c.now = c.now.Add(d)
+	return c.now
+}
+
+// Set sets the clock to the given time and returns the new time.
+func (c *TestClock) Set(t time.Time) time.Time {
+	c.now = t
+	return c.now
 }

apps/agent/services/ratelimit/mitigate.go

@@ -20,9 +20,8 @@ func (s *service) Mitigate(ctx context.Context, req *ratelimitv1.MitigateRequest
 	duration := time.Duration(req.Duration) * time.Millisecond
 	bucket, _ := s.getBucket(bucketKey{req.Identifier, req.Limit, duration})
 	bucket.Lock()
-	defer bucket.Unlock()
-
 	bucket.windows[req.Window.GetSequence()] = req.Window
+	bucket.Unlock()
 
 	return &ratelimitv1.MitigateResponse{}, nil
 }
@@ -51,12 +50,14 @@ func (s *service) broadcastMitigation(req mitigateWindowRequest) {
 		return
 	}
 	for _, peer := range peers {
-		_, err := peer.client.Mitigate(ctx, connect.NewRequest(&ratelimitv1.MitigateRequest{
-			Identifier: req.identifier,
-			Limit:      req.limit,
-			Duration:   req.duration.Milliseconds(),
-			Window:     req.window,
-		}))
+		_, err := s.mitigateCircuitBreaker.Do(ctx, func(innerCtx context.Context) (*connect.Response[ratelimitv1.MitigateResponse], error) {
+			return peer.client.Mitigate(innerCtx, connect.NewRequest(&ratelimitv1.MitigateRequest{
+				Identifier: req.identifier,
+				Limit:      req.limit,
+				Duration:   req.duration.Milliseconds(),
+				Window:     req.window,
+			}))
+		})
 		if err != nil {
 			s.logger.Err(err).Msg("failed to call mitigate")
 		} else {

apps/agent/services/ratelimit/ratelimit_mitigation_test.go

@@ -23,8 +23,8 @@ import (
 )
 
 func TestExceedingTheLimitShouldNotifyAllNodes(t *testing.T) {
-	t.Skip()
-	for _, clusterSize := range []int{1, 3, 5} {
+
+	for _, clusterSize := range []int{1, 3, 5, 9, 27} {
 		t.Run(fmt.Sprintf("Cluster Size %d", clusterSize), func(t *testing.T) {
 			logger := logging.New(nil)
 			clusters := []cluster.Cluster{}
@@ -94,23 +94,25 @@ func TestExceedingTheLimitShouldNotifyAllNodes(t *testing.T) {
 			ctx := context.Background()
 
 			// Saturate the window
-			for i := int64(0); i <= limit; i++ {
+			for i := int64(0); i < limit; i++ {
 				rl := util.RandomElement(ratelimiters)
 				res, err := rl.Ratelimit(ctx, req)
 				require.NoError(t, err)
 				t.Logf("saturate res: %+v", res)
 				require.True(t, res.Success)
+
 			}
 
 			time.Sleep(time.Second * 5)
 
 			// Let's hit everry node again
 			// They should all be mitigated
 			for i, rl := range ratelimiters {
+
 				res, err := rl.Ratelimit(ctx, req)
 				require.NoError(t, err)
 				t.Logf("res from %d: %+v", i, res)
-				// require.False(t, res.Success)
+				require.False(t, res.Success)
 			}
 
 		})

apps/agent/services/ratelimit/ratelimit_replication_test.go

@@ -24,8 +24,7 @@ import (
 	"github.com/unkeyed/unkey/apps/agent/pkg/util"
 )
 
-func TestReplication(t *testing.T) {
-	t.Skip()
+func TestSync(t *testing.T) {
 	type Node struct {
 		srv     *service
 		cluster cluster.Cluster