housekeeping: Add context argument to optimization strategy

Add a context argument to the receiver functions of the optimization
strategy. This is required to make use of feature flags in the strategy
implementations.

3 files changed, 46 insertions, 35 deletions

diff --git a/internal/git/housekeeping/optimization_strategy.go b/internal/git/housekeeping/optimization_strategy.go
index 04235eabe..a2808758d 100644
--- a/internal/git/housekeeping/optimization_strategy.go
+++ b/internal/git/housekeeping/optimization_strategy.go
@@ -15,16 +15,16 @@ import (
 type OptimizationStrategy interface {
 	// ShouldRepackObjects determines whether the repository needs to be repacked and, if so,
 	// how it should be done.
-	ShouldRepackObjects() (bool, RepackObjectsConfig)
+	ShouldRepackObjects(context.Context) (bool, RepackObjectsConfig)
 	// ShouldPruneObjects determines whether the repository has stale objects that should be
 	// pruned.
-	ShouldPruneObjects() bool
+	ShouldPruneObjects(context.Context) bool
 	// ShouldRepackReferences determines whether the repository's references need to be
 	// repacked.
-	ShouldRepackReferences() bool
+	ShouldRepackReferences(context.Context) bool
 	// ShouldWriteCommitGraph determines whether we need to write the commit-graph and how it
 	// should be written.
-	ShouldWriteCommitGraph() (bool, WriteCommitGraphConfig)
+	ShouldWriteCommitGraph(context.Context) (bool, WriteCommitGraphConfig)
 }
 
 // HeuristicalOptimizationStrategy is an optimization strategy that is based on a set of
@@ -54,7 +54,7 @@ func NewHeuristicalOptimizationStrategy(ctx context.Context, repo *localrepo.Rep
 // ShouldRepackObjects checks whether the repository's objects need to be repacked. This uses a
 // set of heuristics that scales with the size of the object database: the larger the repository,
 // the less frequent does it get a full repack.
-func (s HeuristicalOptimizationStrategy) ShouldRepackObjects() (bool, RepackObjectsConfig) {
+func (s HeuristicalOptimizationStrategy) ShouldRepackObjects(context.Context) (bool, RepackObjectsConfig) {
 	// If there are neither packfiles nor loose objects in this repository then there is no need
 	// to repack anything.
 	if s.info.Packfiles.Count == 0 && s.info.LooseObjects.Count == 0 {
@@ -149,7 +149,7 @@ func (s HeuristicalOptimizationStrategy) ShouldRepackObjects() (bool, RepackObje
 
 // ShouldWriteCommitGraph determines whether we need to write the commit-graph and how it should be
 // written.
-func (s HeuristicalOptimizationStrategy) ShouldWriteCommitGraph() (bool, WriteCommitGraphConfig) {
+func (s HeuristicalOptimizationStrategy) ShouldWriteCommitGraph(ctx context.Context) (bool, WriteCommitGraphConfig) {
 	// If the repository doesn't have any references at all then there is no point in writing
 	// commit-graphs given that it would only contain reachable objects, of which there are
 	// none.
@@ -165,7 +165,7 @@ func (s HeuristicalOptimizationStrategy) ShouldWriteCommitGraph() (bool, WriteCo
 	//
 	// To fix this case we will replace the complete commit-chain when we have pruned objects
 	// from the repository.
-	if s.ShouldPruneObjects() {
+	if s.ShouldPruneObjects(ctx) {
 		return true, WriteCommitGraphConfig{
 			ReplaceChain: true,
 		}
@@ -183,7 +183,7 @@ func (s HeuristicalOptimizationStrategy) ShouldWriteCommitGraph() (bool, WriteCo
 
 	// When we repacked the repository then chances are high that we have accumulated quite some
 	// objects since the last time we wrote a commit-graph.
-	if needsRepacking, _ := s.ShouldRepackObjects(); needsRepacking {
+	if needsRepacking, _ := s.ShouldRepackObjects(ctx); needsRepacking {
 		return true, WriteCommitGraphConfig{}
 	}
 
@@ -193,7 +193,7 @@ func (s HeuristicalOptimizationStrategy) ShouldWriteCommitGraph() (bool, WriteCo
 // ShouldPruneObjects determines whether the repository has stale objects that should be pruned.
 // Object pools are never pruned to not lose data in them, but otherwise we prune when we've found
 // enough stale objects that might in fact get pruned.
-func (s HeuristicalOptimizationStrategy) ShouldPruneObjects() bool {
+func (s HeuristicalOptimizationStrategy) ShouldPruneObjects(context.Context) bool {
 	// Pool repositories must never prune any objects, or otherwise we may corrupt members of
 	// that pool if they still refer to that object.
 	if s.isObjectPool {
@@ -212,7 +212,7 @@ func (s HeuristicalOptimizationStrategy) ShouldPruneObjects() bool {
 // ShouldRepackReferences determines whether the repository's references need to be repacked based
 // on heuristics. The more references there are, the more loose referencos may exist until they are
 // packed again.
-func (s HeuristicalOptimizationStrategy) ShouldRepackReferences() bool {
+func (s HeuristicalOptimizationStrategy) ShouldRepackReferences(context.Context) bool {
 	// If there aren't any loose refs then there is nothing we need to do.
 	if s.info.References.LooseReferencesCount == 0 {
 		return false
@@ -274,7 +274,7 @@ func NewEagerOptimizationStrategy(ctx context.Context, repo *localrepo.Repo) (Ea
 // ShouldRepackObjects always instructs the caller to repack objects. The strategy will always be to
 // repack all objects into a single packfile. The bitmap will be written in case the repository does
 // not have any alterantes.
-func (s EagerOptimizationStrategy) ShouldRepackObjects() (bool, RepackObjectsConfig) {
+func (s EagerOptimizationStrategy) ShouldRepackObjects(context.Context) (bool, RepackObjectsConfig) {
 	return true, RepackObjectsConfig{
 		FullRepack:  true,
 		WriteBitmap: !s.hasAlternate,
@@ -283,7 +283,7 @@ func (s EagerOptimizationStrategy) ShouldRepackObjects() (bool, RepackObjectsCon
 
 // ShouldWriteCommitGraph always instructs the caller to write the commit-graph. The strategy will
 // always be to completely rewrite the commit-graph chain.
-func (s EagerOptimizationStrategy) ShouldWriteCommitGraph() (bool, WriteCommitGraphConfig) {
+func (s EagerOptimizationStrategy) ShouldWriteCommitGraph(context.Context) (bool, WriteCommitGraphConfig) {
 	return true, WriteCommitGraphConfig{
 		ReplaceChain: true,
 	}
@@ -291,11 +291,11 @@ func (s EagerOptimizationStrategy) ShouldWriteCommitGraph() (bool, WriteCommitGr
 
 // ShouldPruneObjects always instructs the caller to prune objects, unless the repository is an
 // object pool.
-func (s EagerOptimizationStrategy) ShouldPruneObjects() bool {
+func (s EagerOptimizationStrategy) ShouldPruneObjects(context.Context) bool {
 	return !s.isObjectPool
 }
 
 // ShouldRepackReferences always instructs the caller to repack references.
-func (s EagerOptimizationStrategy) ShouldRepackReferences() bool {
+func (s EagerOptimizationStrategy) ShouldRepackReferences(context.Context) bool {
 	return true
 }
diff --git a/internal/git/housekeeping/optimization_strategy_test.go b/internal/git/housekeeping/optimization_strategy_test.go
index f8497eadc..fa903aa44 100644
--- a/internal/git/housekeeping/optimization_strategy_test.go
+++ b/internal/git/housekeeping/optimization_strategy_test.go
@@ -1,6 +1,7 @@
 package housekeeping
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -317,6 +318,8 @@ func TestNewHeuristicalOptimizationStrategy_variousParameters(t *testing.T) {
 func TestHeuristicalOptimizationStrategy_ShouldRepackObjects(t *testing.T) {
 	t.Parallel()
 
+	ctx := testhelper.Context(t)
+
 	for _, tc := range []struct {
 		desc           string
 		strategy       HeuristicalOptimizationStrategy
@@ -374,7 +377,7 @@ func TestHeuristicalOptimizationStrategy_ShouldRepackObjects(t *testing.T) {
 		},
 	} {
 		t.Run(tc.desc, func(t *testing.T) {
-			repackNeeded, repackCfg := tc.strategy.ShouldRepackObjects()
+			repackNeeded, repackCfg := tc.strategy.ShouldRepackObjects(ctx)
 			require.Equal(t, tc.expectedNeeded, repackNeeded)
 			require.Equal(t, tc.expectedConfig, repackCfg)
 		})
@@ -458,14 +461,14 @@ func TestHeuristicalOptimizationStrategy_ShouldRepackObjects(t *testing.T) {
 						isObjectPool: tc.isPool,
 					}
 
-					repackNeeded, _ := strategy.ShouldRepackObjects()
+					repackNeeded, _ := strategy.ShouldRepackObjects(ctx)
 					require.False(t, repackNeeded)
 
 					// Now we add the last packfile that should bring us across
 					// the boundary of having to repack.
 					strategy.info.Packfiles.Count++
 
-					repackNeeded, repackCfg := strategy.ShouldRepackObjects()
+					repackNeeded, repackCfg := strategy.ShouldRepackObjects(ctx)
 					require.True(t, repackNeeded)
 					require.Equal(t, RepackObjectsConfig{
 						FullRepack:  true,
@@ -530,7 +533,7 @@ func TestHeuristicalOptimizationStrategy_ShouldRepackObjects(t *testing.T) {
 					isObjectPool: tc.isPool,
 				}
 
-				repackNeeded, repackCfg := strategy.ShouldRepackObjects()
+				repackNeeded, repackCfg := strategy.ShouldRepackObjects(ctx)
 				require.Equal(t, outerTC.expectedRepack, repackNeeded)
 				require.Equal(t, RepackObjectsConfig{
 					FullRepack:  false,
@@ -544,6 +547,8 @@ func TestHeuristicalOptimizationStrategy_ShouldRepackObjects(t *testing.T) {
 func TestHeuristicalOptimizationStrategy_ShouldPruneObjects(t *testing.T) {
 	t.Parallel()
 
+	ctx := testhelper.Context(t)
+
 	for _, tc := range []struct {
 		desc                       string
 		strategy                   HeuristicalOptimizationStrategy
@@ -590,13 +595,13 @@ func TestHeuristicalOptimizationStrategy_ShouldPruneObjects(t *testing.T) {
 	} {
 		t.Run(tc.desc, func(t *testing.T) {
 			t.Run("normal repository", func(t *testing.T) {
-				require.Equal(t, tc.expectedShouldPruneObjects, tc.strategy.ShouldPruneObjects())
+				require.Equal(t, tc.expectedShouldPruneObjects, tc.strategy.ShouldPruneObjects(ctx))
 			})
 
 			t.Run("object pool", func(t *testing.T) {
 				strategy := tc.strategy
 				strategy.isObjectPool = true
-				require.False(t, strategy.ShouldPruneObjects())
+				require.False(t, strategy.ShouldPruneObjects(ctx))
 			})
 		})
 	}
@@ -605,6 +610,8 @@ func TestHeuristicalOptimizationStrategy_ShouldPruneObjects(t *testing.T) {
 func TestHeuristicalOptimizationStrategy_ShouldRepackReferences(t *testing.T) {
 	t.Parallel()
 
+	ctx := testhelper.Context(t)
+
 	const kiloByte = 1024
 
 	for _, tc := range []struct {
@@ -650,11 +657,11 @@ func TestHeuristicalOptimizationStrategy_ShouldRepackReferences(t *testing.T) {
 				},
 			}
 
-			require.False(t, strategy.ShouldRepackReferences())
+			require.False(t, strategy.ShouldRepackReferences(ctx))
 
 			strategy.info.References.LooseReferencesCount++
 
-			require.True(t, strategy.ShouldRepackReferences())
+			require.True(t, strategy.ShouldRepackReferences(ctx))
 		})
 	}
 }
@@ -662,6 +669,8 @@ func TestHeuristicalOptimizationStrategy_ShouldRepackReferences(t *testing.T) {
 func TestHeuristicalOptimizationStrategy_NeedsWriteCommitGraph(t *testing.T) {
 	t.Parallel()
 
+	ctx := testhelper.Context(t)
+
 	for _, tc := range []struct {
 		desc           string
 		strategy       HeuristicalOptimizationStrategy
@@ -783,7 +792,7 @@ func TestHeuristicalOptimizationStrategy_NeedsWriteCommitGraph(t *testing.T) {
 		},
 	} {
 		t.Run(tc.desc, func(t *testing.T) {
-			needed, writeCommitGraphCfg := tc.strategy.ShouldWriteCommitGraph()
+			needed, writeCommitGraphCfg := tc.strategy.ShouldWriteCommitGraph(ctx)
 			require.Equal(t, tc.expectedNeeded, needed)
 			require.Equal(t, tc.expectedCfg, writeCommitGraphCfg)
 		})
@@ -851,6 +860,8 @@ func TestNewEagerOptimizationStrategy(t *testing.T) {
 func TestEagerOptimizationStrategy(t *testing.T) {
 	t.Parallel()
 
+	ctx := testhelper.Context(t)
+
 	for _, tc := range []struct {
 		desc                     string
 		strategy                 EagerOptimizationStrategy
@@ -885,21 +896,21 @@ func TestEagerOptimizationStrategy(t *testing.T) {
 		},
 	} {
 		t.Run(tc.desc, func(t *testing.T) {
-			shouldRepackObjects, repackObjectsCfg := tc.strategy.ShouldRepackObjects()
+			shouldRepackObjects, repackObjectsCfg := tc.strategy.ShouldRepackObjects(ctx)
 			require.True(t, shouldRepackObjects)
 			require.Equal(t, RepackObjectsConfig{
 				FullRepack:  true,
 				WriteBitmap: tc.expectWriteBitmap,
 			}, repackObjectsCfg)
 
-			shouldWriteCommitGraph, writeCommitGraphCfg := tc.strategy.ShouldWriteCommitGraph()
+			shouldWriteCommitGraph, writeCommitGraphCfg := tc.strategy.ShouldWriteCommitGraph(ctx)
 			require.True(t, shouldWriteCommitGraph)
 			require.Equal(t, WriteCommitGraphConfig{
 				ReplaceChain: true,
 			}, writeCommitGraphCfg)
 
-			require.Equal(t, tc.expectShouldPruneObjects, tc.strategy.ShouldPruneObjects())
-			require.True(t, tc.strategy.ShouldRepackReferences())
+			require.Equal(t, tc.expectShouldPruneObjects, tc.strategy.ShouldPruneObjects(ctx))
+			require.True(t, tc.strategy.ShouldRepackReferences(ctx))
 		})
 	}
 }
@@ -914,19 +925,19 @@ type mockOptimizationStrategy struct {
 	writeCommitGraphCfg    WriteCommitGraphConfig
 }
 
-func (m mockOptimizationStrategy) ShouldRepackObjects() (bool, RepackObjectsConfig) {
+func (m mockOptimizationStrategy) ShouldRepackObjects(context.Context) (bool, RepackObjectsConfig) {
 	return m.shouldRepackObjects, m.repackObjectsCfg
 }
 
-func (m mockOptimizationStrategy) ShouldPruneObjects() bool {
+func (m mockOptimizationStrategy) ShouldPruneObjects(context.Context) bool {
 	return m.shouldPruneObjects
 }
 
-func (m mockOptimizationStrategy) ShouldRepackReferences() bool {
+func (m mockOptimizationStrategy) ShouldRepackReferences(context.Context) bool {
 	return m.shouldRepackReferences
 }
 
-func (m mockOptimizationStrategy) ShouldWriteCommitGraph() (bool, WriteCommitGraphConfig) {
+func (m mockOptimizationStrategy) ShouldWriteCommitGraph(context.Context) (bool, WriteCommitGraphConfig) {
 	return m.shouldWriteCommitGraph, m.writeCommitGraphCfg
 }
 
diff --git a/internal/git/housekeeping/optimize_repository.go b/internal/git/housekeeping/optimize_repository.go
index 52fe8b156..b4827259c 100644
--- a/internal/git/housekeeping/optimize_repository.go
+++ b/internal/git/housekeeping/optimize_repository.go
@@ -159,7 +159,7 @@ func optimizeRepository(
 
 // repackIfNeeded repacks the repository according to the strategy.
 func repackIfNeeded(ctx context.Context, repo *localrepo.Repo, strategy OptimizationStrategy) (bool, RepackObjectsConfig, error) {
-	repackNeeded, cfg := strategy.ShouldRepackObjects()
+	repackNeeded, cfg := strategy.ShouldRepackObjects(ctx)
 	if !repackNeeded {
 		return false, RepackObjectsConfig{}, nil
 	}
@@ -173,7 +173,7 @@ func repackIfNeeded(ctx context.Context, repo *localrepo.Repo, strategy Optimiza
 
 // writeCommitGraphIfNeeded writes the commit-graph if required.
 func writeCommitGraphIfNeeded(ctx context.Context, repo *localrepo.Repo, strategy OptimizationStrategy) (bool, WriteCommitGraphConfig, error) {
-	needed, cfg := strategy.ShouldWriteCommitGraph()
+	needed, cfg := strategy.ShouldWriteCommitGraph(ctx)
 	if !needed {
 		return false, WriteCommitGraphConfig{}, nil
 	}
@@ -188,7 +188,7 @@ func writeCommitGraphIfNeeded(ctx context.Context, repo *localrepo.Repo, strateg
 // pruneIfNeeded removes objects from the repository which are either unreachable or which are
 // already part of a packfile. We use a grace period of two weeks.
 func pruneIfNeeded(ctx context.Context, repo *localrepo.Repo, strategy OptimizationStrategy) (bool, error) {
-	if !strategy.ShouldPruneObjects() {
+	if !strategy.ShouldPruneObjects(ctx) {
 		return false, nil
 	}
 
@@ -210,7 +210,7 @@ func pruneIfNeeded(ctx context.Context, repo *localrepo.Repo, strategy Optimizat
 }
 
 func packRefsIfNeeded(ctx context.Context, repo *localrepo.Repo, strategy OptimizationStrategy) (bool, error) {
-	if !strategy.ShouldRepackReferences() {
+	if !strategy.ShouldRepackReferences(ctx) {
 		return false, nil
 	}