limithandler: Move per-repository concurrency keyer

The limithandler supports use of different keys to limit concurrency by
via a keyer function. The main function used by production code is
hosted in an unrelated package, mostly because it's where the limiter is
currently initialized.

Move it into the limithandler package to make it available at a central
place.

2 files changed, 18 insertions, 17 deletions

diff --git a/internal/gitaly/server/server.go b/internal/gitaly/server/server.go
index 9dcf07aa4..0ed05a358 100644
--- a/internal/gitaly/server/server.go
+++ b/internal/gitaly/server/server.go
@@ -1,7 +1,6 @@
 package server
 
 import (
-	"context"
 	"crypto/tls"
 	"fmt"
 	"time"
@@ -37,21 +36,6 @@ import (
 	"google.golang.org/grpc/keepalive"
 )
 
-func concurrencyKeyFn(ctx context.Context) string {
-	tags := grpcmwtags.Extract(ctx)
-	ctxValue := tags.Values()["grpc.request.repoPath"]
-	if ctxValue == nil {
-		return ""
-	}
-
-	s, ok := ctxValue.(string)
-	if ok {
-		return s
-	}
-
-	return ""
-}
-
 func init() {
 	for _, l := range gitalylog.Loggers {
 		urlSanitizer := logsanitizer.NewURLSanitizerHook()
@@ -80,7 +64,7 @@ func New(
 		grpcmwtags.WithFieldExtractorForInitialReq(fieldextractors.FieldExtractor),
 	}
 
-	lh := limithandler.New(concurrencyKeyFn)
+	lh := limithandler.New(limithandler.LimitConcurrencyByRepo)
 
 	transportCredentials := insecure.NewCredentials()
 	// If tls config is specified attempt to extract tls options and use it
diff --git a/internal/middleware/limithandler/middleware.go b/internal/middleware/limithandler/middleware.go
index df228ce6f..c3636a4de 100644
--- a/internal/middleware/limithandler/middleware.go
+++ b/internal/middleware/limithandler/middleware.go
@@ -4,6 +4,7 @@ import (
 	"context"
 
 	"github.com/grpc-ecosystem/go-grpc-middleware/logging/logrus/ctxlogrus"
+	grpcmwtags "github.com/grpc-ecosystem/go-grpc-middleware/tags"
 	"google.golang.org/grpc"
 )
 
@@ -12,6 +13,22 @@ var maxConcurrencyPerRepoPerRPC map[string]int
 // GetLockKey function defines the lock key of an RPC invocation based on its context
 type GetLockKey func(context.Context) string
 
+// LimitConcurrencyByRepo implements GetLockKey by using the repository path as lock.
+func LimitConcurrencyByRepo(ctx context.Context) string {
+	tags := grpcmwtags.Extract(ctx)
+	ctxValue := tags.Values()["grpc.request.repoPath"]
+	if ctxValue == nil {
+		return ""
+	}
+
+	s, ok := ctxValue.(string)
+	if ok {
+		return s
+	}
+
+	return ""
+}
+
 // LimiterMiddleware contains rate limiter state
 type LimiterMiddleware struct {
 	methodLimiters map[string]*ConcurrencyLimiter