From 5b5e00ce786334918883869eb24af64ccde4cf9c Mon Sep 17 00:00:00 2001
From: Igor Drozdov <idrozdov@gitlab.com>
Date: Fri, 6 May 2022 12:43:32 +0400
Subject: Put reuse Gitaly connections behind feature flag

The connections get canceled due to some shared state
Let's try disabling reusing Gitaly connections and see
whether anything changes
---
 internal/gitaly/gitaly.go                   |  6 +++++-
 internal/gitaly/gitaly_test.go              | 10 +++++++---
 internal/gitlabnet/accessverifier/client.go |  9 +++++----
 internal/handler/exec.go                    |  2 +-
 internal/handler/exec_test.go               |  7 ++++---
 5 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/internal/gitaly/gitaly.go b/internal/gitaly/gitaly.go
index 9f73661..f11c72a 100644
--- a/internal/gitaly/gitaly.go
+++ b/internal/gitaly/gitaly.go
@@ -42,7 +42,11 @@ func (c *Client) InitSidechannelRegistry(ctx context.Context) {
 	c.SidechannelRegistry = gitalyclient.NewSidechannelRegistry(log.ContextLogger(ctx))
 }
 
-func (c *Client) GetConnection(ctx context.Context, cmd Command) (*grpc.ClientConn, error) {
+func (c *Client) GetConnection(ctx context.Context, reuseConnections bool, cmd Command) (*grpc.ClientConn, error) {
+	if !reuseConnections {
+		return c.newConnection(ctx, cmd)
+	}
+
 	c.cache.RLock()
 	conn := c.cache.connections[cmd]
 	c.cache.RUnlock()
diff --git a/internal/gitaly/gitaly_test.go b/internal/gitaly/gitaly_test.go
index abfb764..9c8e04b 100644
--- a/internal/gitaly/gitaly_test.go
+++ b/internal/gitaly/gitaly_test.go
@@ -37,17 +37,21 @@ func TestCachedConnections(t *testing.T) {
 
 	cmd := Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9999"}
 
-	conn, err := c.GetConnection(context.Background(), cmd)
+	conn, err := c.GetConnection(context.Background(), false, cmd)
+	require.NoError(t, err)
+	require.Len(t, c.cache.connections, 0)
+
+	conn, err = c.GetConnection(context.Background(), true, cmd)
 	require.NoError(t, err)
 	require.Len(t, c.cache.connections, 1)
 
-	newConn, err := c.GetConnection(context.Background(), cmd)
+	newConn, err := c.GetConnection(context.Background(), true, cmd)
 	require.NoError(t, err)
 	require.Len(t, c.cache.connections, 1)
 	require.Equal(t, conn, newConn)
 
 	cmd = Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9998"}
-	_, err = c.GetConnection(context.Background(), cmd)
+	_, err = c.GetConnection(context.Background(), true, cmd)
 	require.NoError(t, err)
 	require.Len(t, c.cache.connections, 2)
 }
diff --git a/internal/gitlabnet/accessverifier/client.go b/internal/gitlabnet/accessverifier/client.go
index bce32cf..0a8a1d9 100644
--- a/internal/gitlabnet/accessverifier/client.go
+++ b/internal/gitlabnet/accessverifier/client.go
@@ -32,10 +32,11 @@ type Request struct {
 }
 
 type Gitaly struct {
-	Repo     pb.Repository     `json:"repository"`
-	Address  string            `json:"address"`
-	Token    string            `json:"token"`
-	Features map[string]string `json:"features"`
+	Repo             pb.Repository     `json:"repository"`
+	Address          string            `json:"address"`
+	Token            string            `json:"token"`
+	Features         map[string]string `json:"features"`
+	ReuseConnections bool              `json:"reuse_gitaly_connections"`
 }
 
 type CustomPayloadData struct {
diff --git a/internal/handler/exec.go b/internal/handler/exec.go
index 0870d2b..1f6c41c 100644
--- a/internal/handler/exec.go
+++ b/internal/handler/exec.go
@@ -119,5 +119,5 @@ func withOutgoingMetadata(ctx context.Context, features map[string]string) conte
 }
 
 func (gc *GitalyCommand) getConn(ctx context.Context) (*grpc.ClientConn, error) {
-	return gc.Config.GitalyClient.GetConnection(ctx, gc.Command)
+	return gc.Config.GitalyClient.GetConnection(ctx, gc.Response.Gitaly.ReuseConnections, gc.Command)
 }
diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
index 791fe79..c9c5855 100644
--- a/internal/handler/exec_test.go
+++ b/internal/handler/exec_test.go
@@ -50,8 +50,9 @@ func TestCachingOfGitalyConnections(t *testing.T) {
 	response := &accessverifier.Response{
 		Username: "user",
 		Gitaly: accessverifier.Gitaly{
-			Address: "tcp://localhost:9999",
-			Token:   "token",
+			Address:          "tcp://localhost:9999",
+			Token:            "token",
+			ReuseConnections: true,
 		},
 	}
 
@@ -69,7 +70,7 @@ func TestCachingOfGitalyConnections(t *testing.T) {
 }
 
 func TestMissingGitalyAddress(t *testing.T) {
-	cmd := GitalyCommand{Config: newConfig()}
+	cmd := GitalyCommand{Config: newConfig(), Response: &accessverifier.Response{}}
 
 	err := cmd.RunGitalyCommand(context.Background(), makeHandler(t, nil))
 	require.EqualError(t, err, "no gitaly_address given")
-- 
cgit v1.2.1