From 802368dea689bd28f48ce37c049210b6b6f4c324 Mon Sep 17 00:00:00 2001
From: Ahmad Sherif <ahmad@gitlab.com>
Date: Fri, 19 Apr 2024 23:05:16 +0200
Subject: [PATCH] Fix Gitaly hanging after cancelled InfoRefs request

When we have an info ref cache miss and the request context is
cancelled when `missFn` [1] is being executed, we could end up in a
situation where the command is blocked trying to write/copy the last bytes of
stdout to a pipe[2] that's not being read from (because this goroutine[3]
has exited at this point after the context was cancelled).

Discarding the remaining bytes [4] doesn't work because all reads from
the `TeeReader` are matched with writes to the writer it writes to, in
this case it is a gRPC stream that's closed because the context was
cancelled, so it will err out.

Closing the reading-end of the pipe side fixes this problem and unblocks
the command's `Wait()` call.

Fixes https://gitlab.com/gitlab-org/gitaly/-/issues/5953

[1]: https://gitlab.com/gitlab-org/gitaly/-/blob/faac94442a49a42427236f28d5190c6d721c172b/internal/gitaly/service/smarthttp/cache.go#L112
[2]: https://gitlab.com/gitlab-org/gitaly/-/blob/faac94442a49a42427236f28d5190c6d721c172b/internal/gitaly/service/smarthttp/cache.go#L92
[3]: https://gitlab.com/gitlab-org/gitaly/-/blob/faac94442a49a42427236f28d5190c6d721c172b/internal/gitaly/service/smarthttp/cache.go#L95-110
[4]: https://gitlab.com/gitlab-org/gitaly/-/blob/faac94442a49a42427236f28d5190c6d721c172b/internal/gitaly/service/smarthttp/cache.go#L104
---
 internal/gitaly/service/smarthttp/cache.go    |  1 +
 .../gitaly/service/smarthttp/inforefs_test.go | 55 +++++++++++++++++++
 2 files changed, 56 insertions(+)

diff --git a/internal/gitaly/service/smarthttp/cache.go b/internal/gitaly/service/smarthttp/cache.go
index 515029c0fc..b0fda42f5f 100644
--- a/internal/gitaly/service/smarthttp/cache.go
+++ b/internal/gitaly/service/smarthttp/cache.go
@@ -106,6 +106,7 @@ func (c infoRefCache) tryCache(ctx context.Context, in *gitalypb.InfoRefsRequest
 					c.logger.WithError(err).
 						ErrorContext(ctx, "unable to discard remaining InfoRefsUploadPack cache stream")
 				}
+				_ = pr.CloseWithError(err) // always returns nil
 			}
 		}()
 
diff --git a/internal/gitaly/service/smarthttp/inforefs_test.go b/internal/gitaly/service/smarthttp/inforefs_test.go
index 8d864e6452..20a42ab074 100644
--- a/internal/gitaly/service/smarthttp/inforefs_test.go
+++ b/internal/gitaly/service/smarthttp/inforefs_test.go
@@ -10,6 +10,7 @@ import (
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v16/internal/cache"
@@ -306,6 +307,60 @@ func testInfoRefsUploadPackGitProtocol(t *testing.T, ctx context.Context) {
 	require.Contains(t, envData, fmt.Sprintf("GIT_PROTOCOL=%s\n", git.ProtocolV2))
 }
 
+func TestInfoRefsUploadPack_cacheNotStuckUponContextCancellation(t *testing.T) {
+	t.Parallel()
+
+	testhelper.NewFeatureSets(
+		featureflag.BundleURI,
+	).Run(t, testInfoRefsUploadPackCacheNotStuckUponContextCancellation)
+}
+
+// We mainly test for goroutine leakage here, so if the server is hanging on finishing the request,
+// we will be leaking a goroutine and the test framework will panic.
+func testInfoRefsUploadPackCacheNotStuckUponContextCancellation(t *testing.T, ctx context.Context) {
+	t.Parallel()
+
+	cfg := testcfg.Build(t)
+
+	locator := config.NewLocator(cfg)
+	logger := testhelper.NewLogger(t)
+	cache := cache.New(cfg, locator, logger)
+
+	streamer := mockStreamer{
+		Streamer: cache,
+	}
+	mockInfoRefCache := newInfoRefCache(logger, &streamer)
+
+	// This is to simulate a long-running `git upload-pack` command so that we can have time to cancel the request abruptly.
+	gitCmdFactory := gittest.NewInterceptingCommandFactory(t, ctx, cfg, func(execEnv git.ExecutionEnvironment) string {
+		return fmt.Sprintf(`#!/usr/bin/env bash
+						if [[ "$@" =~ "upload-pack" ]]; then
+							while true; do echo "foo"; done
+							exit 1
+						fi
+						exec %q "$@"`, execEnv.BinaryPath)
+	})
+	server := startSmartHTTPServerWithOptions(t, cfg,
+		[]ServerOpt{withInfoRefCache(mockInfoRefCache)},
+		[]testserver.GitalyServerOpt{testserver.WithGitCommandFactory(gitCmdFactory)},
+	)
+	cfg.SocketPath = server.Address()
+
+	repo, repoPath := gittest.CreateRepository(t, ctx, cfg)
+
+	commitID := gittest.WriteCommit(t, cfg, repoPath, gittest.WithBranch("main"), gittest.WithParents())
+	_ = gittest.WriteTag(t, cfg, repoPath, "v1.0.0", commitID.Revision(), gittest.WriteTagConfig{
+		Message: "annotated tag",
+	})
+
+	requestCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
+	defer cancel()
+
+	rpcRequest := &gitalypb.InfoRefsRequest{Repository: repo}
+	_, err := makeInfoRefsUploadPackRequest(t, requestCtx, cfg.SocketPath, cfg.Auth.Token, rpcRequest)
+	require.Error(t, err)
+}
+
 func makeInfoRefsUploadPackRequest(t *testing.T, ctx context.Context, serverSocketPath, token string, rpcRequest *gitalypb.InfoRefsRequest) ([]byte, error) {
 	t.Helper()
 
-- 
GitLab