Merge branch 'ps-fix-optimize-job-cancellation' into 'master'

maintenance: Override deadline for RPC call

Closes #3963

See merge request gitlab-org/gitaly!4206

2 files changed, 47 insertions, 0 deletions

diff --git a/internal/gitaly/maintenance/optimize.go b/internal/gitaly/maintenance/optimize.go
index f4da15c84..d7231fb24 100644
--- a/internal/gitaly/maintenance/optimize.go
+++ b/internal/gitaly/maintenance/optimize.go
@@ -56,6 +56,14 @@ func optimizeRepoAtPath(ctx context.Context, l logrus.FieldLogger, s config.Stor
 		Repository: repo,
 	}
 
+	// In order to prevent RPC cancellation because of the parent context cancellation
+	// (job execution duration passed) we suppress cancellation of the parent and add a
+	// new time limit to make sure it won't block forever.
+	// It also helps us to be protected over repositories that are taking too long to complete
+	// a request, so no any other repository can be optimized.
+	ctx, cancel := context.WithTimeout(helper.SuppressCancellation(ctx), 5*time.Minute)
+	defer cancel()
+
 	start := time.Now()
 	if _, err := o.OptimizeRepository(ctx, optimizeReq); err != nil {
 		l.WithFields(map[string]interface{}{
diff --git a/internal/gitaly/maintenance/optimize_test.go b/internal/gitaly/maintenance/optimize_test.go
index 2ca182f13..8c03229d4 100644
--- a/internal/gitaly/maintenance/optimize_test.go
+++ b/internal/gitaly/maintenance/optimize_test.go
@@ -5,6 +5,7 @@ import (
 	"math/rand"
 	"path/filepath"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -121,3 +122,41 @@ func TestOptimizeReposRandomly(t *testing.T) {
 		})
 	}
 }
+
+type mockOptimizerCancel struct {
+	t         *testing.T
+	startedAt time.Time
+}
+
+func (m mockOptimizerCancel) OptimizeRepository(ctx context.Context, _ *gitalypb.OptimizeRepositoryRequest, _ ...grpc.CallOption) (*gitalypb.OptimizeRepositoryResponse, error) {
+	timeline, ok := ctx.Deadline()
+	if assert.True(m.t, ok) {
+		assert.True(m.t, timeline.After(m.startedAt), m.startedAt)
+		future := m.startedAt.Add(10 * time.Minute)
+		assert.True(m.t, timeline.Before(future), future)
+	}
+	return &gitalypb.OptimizeRepositoryResponse{}, nil
+}
+
+func TestOptimizeReposRandomly_cancellationOverride(t *testing.T) {
+	cfgBuilder := testcfg.NewGitalyCfgBuilder()
+	cfg := cfgBuilder.Build(t)
+
+	gittest.InitRepo(t, cfg, cfg.Storages[0])
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	// The timeout should be overwritten by the default 5 min timeout.
+	//nolint:forbidigo // We're explicitly testing deadline override.
+	ctx, cancel = context.WithTimeout(ctx, 72*time.Hour)
+	defer cancel()
+
+	ticker := helper.NewManualTicker()
+	ticker.Tick()
+
+	mo := &mockOptimizerCancel{t: t, startedAt: time.Now()}
+	walker := OptimizeReposRandomly(cfg.Storages, mo, ticker, rand.New(rand.NewSource(1)))
+
+	require.NoError(t, walker(ctx, testhelper.NewDiscardingLogEntry(t), []string{cfg.Storages[0].Name}))
+}