chore(workflows): update axiom-deployment-annotation action to use la…

…test version

chore(ratelimit): change method signatures to accept context parameter
chore(ratelimit): remove unused tracing and logging statements

.github/workflows/job_deploy_agent_production.yaml

@@ -40,13 +40,13 @@ jobs:
 
       - name: Annotate axiom
         id: axiom-deployment-annotation
-        uses: axiomhq/annotation-action@v0.1.0
+        uses: axiomhq/annotation-action
         with:
-          type: "production deployment"
+          type: "production-release"
           datasets: agent
           axiomToken: ${{ secrets.AXIOM_TOKEN }}
-          title: 'Production deployment ${{ env.VERSION }}'
-          description: 'Commit ${{ github.event.head_commit.message }}' # optional
+          # title: 'Production deployment ${{ env.VERSION }}'
+          # description: 'Commit ${{ github.event.head_commit.message }}'
 
 
       - name: Deploy prod

apps/agent/pkg/ratelimit/fixed_window.go

@@ -1,11 +1,14 @@
 package ratelimit
 
 import (
+	"context"
 	"fmt"
 	"sync"
 	"time"
 
 	"github.com/unkeyed/unkey/apps/agent/pkg/logging"
+	"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
+	"go.opentelemetry.io/otel/attribute"
 )
 
 type identifierWindow struct {
@@ -53,17 +56,20 @@ func buildKey(identifier string, limit int64, duration int64) string {
 	return fmt.Sprintf("ratelimit:%s:%d:%d", identifier, limit, window)
 }
 
-func (r *fixedWindow) Take(req RatelimitRequest) RatelimitResponse {
-	start := time.Now()
+func (r *fixedWindow) Take(ctx context.Context, req RatelimitRequest) RatelimitResponse {
+	ctx, span := tracing.Start(ctx, "fixedWindow.Take")
+	defer span.End()
+
 	key := buildKey(req.Identifier, req.Max, req.RefillInterval)
-	defer func() {
-		r.logger.Info().Str("key", key).Int64("latency", time.Since(start).Milliseconds()).Msg("fixedWindow.Take")
-	}()
+	span.SetAttributes(attribute.String("key", key))
 
+	_, lockSpan := tracing.Start(ctx, "fixedWindow.Take.lock")
 	r.identifiersLock.Lock()
+	lockSpan.End()
 	defer r.identifiersLock.Unlock()
 
 	id, ok := r.identifiers[key]
+	span.SetAttributes(attribute.Bool("identifierExisted", ok))
 	if !ok {
 		id = &identifierWindow{id: key, current: 0, reset: time.Now().Add(time.Duration(req.RefillInterval) * time.Millisecond)}
 		r.identifiers[key] = id
@@ -77,13 +83,17 @@ func (r *fixedWindow) Take(req RatelimitRequest) RatelimitResponse {
 	return RatelimitResponse{Pass: true, Remaining: req.Max - id.current, Reset: id.reset.UnixMilli(), Limit: req.Max, Current: id.current}
 }
 
-func (r *fixedWindow) SetCurrent(req SetCurrentRequest) error {
+func (r *fixedWindow) SetCurrent(ctx context.Context, req SetCurrentRequest) error {
+	ctx, span := tracing.Start(ctx, "fixedWindow.SetCurrent")
+	defer span.End()
 	key := buildKey(req.Identifier, req.Max, req.RefillInterval)
 
+	_, lockSpan := tracing.Start(ctx, "fixedWindow.SetCurrent.lock")
 	r.identifiersLock.Lock()
+	lockSpan.End()
 	defer r.identifiersLock.Unlock()
-
 	id, ok := r.identifiers[req.Identifier]
+	span.SetAttributes(attribute.Bool("identifierExisted", ok))
 	if !ok {
 		id = &identifierWindow{id: key, current: 0, reset: time.Now().Add(time.Duration(req.RefillInterval) * time.Millisecond)}
 		r.identifiers[req.Identifier] = id

apps/agent/pkg/ratelimit/interface.go

@@ -1,8 +1,10 @@
 package ratelimit
 
+import "context"
+
 type Ratelimiter interface {
-	Take(req RatelimitRequest) RatelimitResponse
-	SetCurrent(req SetCurrentRequest) error
+	Take(ctx context.Context, req RatelimitRequest) RatelimitResponse
+	SetCurrent(ctx context.Context, req SetCurrentRequest) error
 }
 
 type RatelimitRequest struct {

apps/agent/services/ratelimit/pushpull.go

@@ -6,7 +6,6 @@ import (
 
 	ratelimitv1 "github.com/unkeyed/unkey/apps/agent/gen/proto/ratelimit/v1"
 	"github.com/unkeyed/unkey/apps/agent/pkg/ratelimit"
-	"github.com/unkeyed/unkey/apps/agent/pkg/tracing"
 )
 
 func (s *service) PushPull(ctx context.Context, req *ratelimitv1.PushPullRequest) (*ratelimitv1.PushPullResponse, error) {
@@ -16,9 +15,7 @@ func (s *service) PushPull(ctx context.Context, req *ratelimitv1.PushPullRequest
 			Int64("latency", time.Since(start).Milliseconds()).
 			Msg("service.PushPull")
 	}()
-	ctx, span := tracing.Start(ctx, "PushPull")
-	defer span.End()
-	res := s.ratelimiter.Take(ratelimit.RatelimitRequest{
+	res := s.ratelimiter.Take(ctx, ratelimit.RatelimitRequest{
 		Identifier:     req.Identifier,
 		Max:            req.Limit,
 		RefillRate:     req.Limit,

apps/agent/services/ratelimit/ratelimit.go

@@ -16,9 +16,7 @@ func (s *service) Ratelimit(ctx context.Context, req *ratelimitv1.RatelimitReque
 			Int64("latency", time.Since(start).Milliseconds()).
 			Msg("service.Ratelimit")
 	}()
-	_, span := tracing.Start(ctx, "Ratelimit")
-	defer span.End()
-	res := s.ratelimiter.Take(ratelimit.RatelimitRequest{
+	res := s.ratelimiter.Take(ctx, ratelimit.RatelimitRequest{
 		Identifier:     req.Identifier,
 		Max:            req.Limit,
 		RefillRate:     req.Limit,
@@ -28,14 +26,15 @@ func (s *service) Ratelimit(ctx context.Context, req *ratelimitv1.RatelimitReque
 	s.logger.Info().Interface("req", req).Interface("res", res).Msg("ratelimit")
 
 	if s.pushPullC != nil {
+		_, span := tracing.Start(ctx, "emitting pushPull event")
 		e := pushPullEvent{
 			identifier: req.Identifier,
 			limit:      req.Limit,
 			duration:   req.Duration,
 			cost:       req.Cost,
 		}
-		s.logger.Info().Interface("event", e).Msg("queueing pushPull with origin")
 		s.pushPullC <- e
+		span.End()
 
 	}
 

apps/agent/services/ratelimit/sync.go

@@ -54,7 +54,7 @@ func (s *service) createWorker(id int) {
 		}
 		logger.Info().Str("peerId", peer.Id).Str("key", key).Interface("res", res).Msg("push pull came back")
 
-		err = s.ratelimiter.SetCurrent(ratelimit.SetCurrentRequest{
+		err = s.ratelimiter.SetCurrent(context.Background(), ratelimit.SetCurrentRequest{
 			Identifier:     e.identifier,
 			Max:            e.limit,
 			Current:        res.Msg.Current,