Rate limiter

3 files changed, 36 insertions, 10 deletions

diff --git a/internal/config/concurrency.go b/internal/config/concurrency.go
new file mode 100644
index 000000000..7211bd3e1
--- /dev/null
+++ b/internal/config/concurrency.go
@@ -0,0 +1,16 @@
+package config
+
+import (
+	"gitlab.com/gitlab-org/gitaly/internal/middleware/limithandler"
+)
+
+// ConfigureConcurrencyLimits configures the per-repo, per RPC rate limits
+func ConfigureConcurrencyLimits() {
+	maxConcurrencyPerRepoPerRPC := make(map[string]int)
+
+	for _, v := range Config.Concurrency {
+		maxConcurrencyPerRepoPerRPC[v.RPC] = v.MaxPerRepo
+	}
+
+	limithandler.SetMaxRepoConcurrency(maxConcurrencyPerRepoPerRPC)
+}
diff --git a/internal/config/config.go b/internal/config/config.go
index 61e0705fe..d1eac0ad5 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -19,16 +19,17 @@ var (
 )
 
 type config struct {
-	SocketPath           string      `toml:"socket_path" split_words:"true"`
-	ListenAddr           string      `toml:"listen_addr" split_words:"true"`
-	PrometheusListenAddr string      `toml:"prometheus_listen_addr" split_words:"true"`
-	Git                  Git         `toml:"git" envconfig:"git"`
-	Storages             []Storage   `toml:"storage" envconfig:"storage"`
-	Logging              Logging     `toml:"logging" envconfig:"logging"`
-	Prometheus           Prometheus  `toml:"prometheus"`
-	Auth                 Auth        `toml:"auth"`
-	Ruby                 Ruby        `toml:"gitaly-ruby"`
-	GitlabShell          GitlabShell `toml:"gitlab-shell"`
+	SocketPath           string        `toml:"socket_path" split_words:"true"`
+	ListenAddr           string        `toml:"listen_addr" split_words:"true"`
+	PrometheusListenAddr string        `toml:"prometheus_listen_addr" split_words:"true"`
+	Git                  Git           `toml:"git" envconfig:"git"`
+	Storages             []Storage     `toml:"storage" envconfig:"storage"`
+	Logging              Logging       `toml:"logging" envconfig:"logging"`
+	Prometheus           Prometheus    `toml:"prometheus"`
+	Auth                 Auth          `toml:"auth"`
+	Ruby                 Ruby          `toml:"gitaly-ruby"`
+	GitlabShell          GitlabShell   `toml:"gitlab-shell"`
+	Concurrency          []Concurrency `toml:"concurrency"`
 }
 
 // GitlabShell contains the settings required for executing `gitlab-shell`
@@ -58,6 +59,12 @@ type Prometheus struct {
 	GRPCLatencyBuckets []float64 `toml:"grpc_latency_buckets"`
 }
 
+// Concurrency allows endpoints to be limited to a maximum concurrency per repo
+type Concurrency struct {
+	RPC        string `toml:"rpc"`
+	MaxPerRepo int    `toml:"max_per_repo"`
+}
+
 // Load initializes the Config variable from file and the environment.
 //  Environment variables take precedence over the file.
 func Load(file io.Reader) error {
diff --git a/internal/config/prometheus.go b/internal/config/prometheus.go
index eb5dacc5b..2d8936f9a 100644
--- a/internal/config/prometheus.go
+++ b/internal/config/prometheus.go
@@ -4,6 +4,7 @@ import (
 	"github.com/grpc-ecosystem/go-grpc-prometheus"
 	"github.com/prometheus/client_golang/prometheus"
 	log "github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/gitaly/internal/middleware/limithandler"
 )
 
 // ConfigurePrometheus uses the global configuration to configure prometheus
@@ -17,4 +18,6 @@ func ConfigurePrometheus() {
 	grpc_prometheus.EnableHandlingTimeHistogram(func(histogramOpts *prometheus.HistogramOpts) {
 		histogramOpts.Buckets = Config.Prometheus.GRPCLatencyBuckets
 	})
+
+	limithandler.EnableAcquireTimeHistogram(Config.Prometheus.GRPCLatencyBuckets)
 }