limithandler: Log queue duration in limit.queue_ms log field

4 files changed, 109 insertions, 0 deletions

diff --git a/internal/gitaly/server/server.go b/internal/gitaly/server/server.go
index 94f11279b..c529a9038 100644
--- a/internal/gitaly/server/server.go
+++ b/internal/gitaly/server/server.go
@@ -20,6 +20,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/cache"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/commandstatshandler"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/featureflag"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/limithandler"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/metadatahandler"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/panichandler"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/middleware/sentryhandler"
@@ -111,6 +112,7 @@ func (s *GitalyServerFactory) New(secure bool, opts ...Option) (*grpc.Server, er
 			grpcstats.FieldsProducer,
 			featureflag.FieldsProducer,
 			structerr.FieldsProducer,
+			limithandler.FieldsProducer,
 		),
 	)
 
@@ -120,6 +122,7 @@ func (s *GitalyServerFactory) New(secure bool, opts ...Option) (*grpc.Server, er
 		metadatahandler.StreamInterceptor,
 		grpcprometheus.StreamServerInterceptor,
 		commandstatshandler.StreamInterceptor,
+		limithandler.StatsStreamInterceptor,
 		grpcmwlogrus.StreamServerInterceptor(s.logger,
 			grpcmwlogrus.WithTimestampFormat(gitalylog.LogTimestampFormat),
 			logMsgProducer,
@@ -136,6 +139,7 @@ func (s *GitalyServerFactory) New(secure bool, opts ...Option) (*grpc.Server, er
 		metadatahandler.UnaryInterceptor,
 		grpcprometheus.UnaryServerInterceptor,
 		commandstatshandler.UnaryInterceptor,
+		limithandler.StatsUnaryInterceptor,
 		grpcmwlogrus.UnaryServerInterceptor(s.logger,
 			grpcmwlogrus.WithTimestampFormat(gitalylog.LogTimestampFormat),
 			logMsgProducer,
diff --git a/internal/middleware/limithandler/monitor.go b/internal/middleware/limithandler/monitor.go
index 602ebcbb4..0c9ded13f 100644
--- a/internal/middleware/limithandler/monitor.go
+++ b/internal/middleware/limithandler/monitor.go
@@ -88,6 +88,10 @@ func (p *PromMonitor) Enter(ctx context.Context, acquireTime time.Duration) {
 	}
 
 	p.acquiringSecondsMetric.Observe(acquireTime.Seconds())
+
+	if stats := limitStatsFromContext(ctx); stats != nil {
+		stats.AddConcurrencyQueueMs(acquireTime.Milliseconds())
+	}
 }
 
 // Exit is called when a request has finished processing
diff --git a/internal/middleware/limithandler/stats.go b/internal/middleware/limithandler/stats.go
new file mode 100644
index 000000000..5f77766fe
--- /dev/null
+++ b/internal/middleware/limithandler/stats.go
@@ -0,0 +1,78 @@
+package limithandler
+
+import (
+	"context"
+	"sync/atomic"
+
+	grpcmw "github.com/grpc-ecosystem/go-grpc-middleware"
+	"github.com/sirupsen/logrus"
+	"google.golang.org/grpc"
+)
+
+type limitStatsKey struct{}
+
+// LimitStats contains info about the concurrency limiter.
+type LimitStats struct {
+	// concurrencyQueueMs milliseconds waiting in concurrency limit queue.
+	concurrencyQueueMs int64
+}
+
+// InitLimitStats initializes context with a per-RPC stats struct
+func InitLimitStats(ctx context.Context) context.Context {
+	return context.WithValue(ctx, limitStatsKey{}, &LimitStats{})
+}
+
+// AddConcurrencyQueueMs adds queue time.
+func (s *LimitStats) AddConcurrencyQueueMs(queueMs int64) {
+	atomic.AddInt64(&s.concurrencyQueueMs, queueMs)
+}
+
+// Fields returns logging info.
+func (s *LimitStats) Fields() logrus.Fields {
+	return logrus.Fields{
+		"limit.concurrency_queue_ms": s.concurrencyQueueMs,
+	}
+}
+
+// FieldsProducer extracts stats info from the context and returns it as a logging fields.
+func FieldsProducer(ctx context.Context, _ error) logrus.Fields {
+	stats := limitStatsFromContext(ctx)
+	if stats != nil {
+		return stats.Fields()
+	}
+	return nil
+}
+
+func limitStatsFromContext(ctx context.Context) *LimitStats {
+	v, ok := ctx.Value(limitStatsKey{}).(*LimitStats)
+	if !ok {
+		return nil
+	}
+	return v
+}
+
+// stats interceptors are separate from middleware and serve one main purpose:
+// initialize the stats object early, so that logrus can see it.
+// it must be placed before the limithandler middleware
+
+// StatsUnaryInterceptor returns a Unary Interceptor that initializes the context
+func StatsUnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+	ctx = InitLimitStats(ctx)
+
+	res, err := handler(ctx, req)
+
+	return res, err
+}
+
+// StatsStreamInterceptor returns a Stream Interceptor
+func StatsStreamInterceptor(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+	ctx := stream.Context()
+	ctx = InitLimitStats(ctx)
+
+	wrapped := grpcmw.WrapServerStream(stream)
+	wrapped.WrappedContext = ctx
+
+	err := handler(srv, wrapped)
+
+	return err
+}
diff --git a/internal/middleware/limithandler/stats_test.go b/internal/middleware/limithandler/stats_test.go
new file mode 100644
index 000000000..e25fd16fb
--- /dev/null
+++ b/internal/middleware/limithandler/stats_test.go
@@ -0,0 +1,23 @@
+//go:build !gitaly_test_sha256
+
+package limithandler
+
+import (
+	"testing"
+
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
+)
+
+func TestLimitStats(t *testing.T) {
+	t.Parallel()
+
+	ctx := testhelper.Context(t)
+	ctx = InitLimitStats(ctx)
+
+	stats := limitStatsFromContext(ctx)
+	stats.AddConcurrencyQueueMs(13)
+
+	assert.Equal(t, FieldsProducer(ctx, nil), logrus.Fields{"limit.concurrency_queue_ms": int64(13)})
+}