cmd/gitaly-lfs-smudge: Move creation of the context into `run()`

The context is currently created when smudging contents. We're about to
add another code path to the `run()` function though that will also need
a context.

Move creation of the context higher up into `run()` so that it can be
shared across both code paths.

3 files changed, 17 insertions, 11 deletions

diff --git a/cmd/gitaly-lfs-smudge/lfs_smudge.go b/cmd/gitaly-lfs-smudge/lfs_smudge.go
index d2676e14a..47689c5a3 100644
--- a/cmd/gitaly-lfs-smudge/lfs_smudge.go
+++ b/cmd/gitaly-lfs-smudge/lfs_smudge.go
@@ -11,16 +11,9 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/config/prometheus"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitlab"
 	"gitlab.com/gitlab-org/labkit/log"
-	"gitlab.com/gitlab-org/labkit/tracing"
 )
 
-func smudgeContents(cfg smudge.Config, to io.Writer, from io.Reader) (returnedErr error) {
-	// Since the environment is sanitized at the moment, we're only
-	// using this to extract the correlation ID. The finished() call
-	// to clean up the tracing will be a NOP here.
-	ctx, finished := tracing.ExtractFromEnv(context.Background())
-	defer finished()
-
+func smudgeContents(ctx context.Context, cfg smudge.Config, to io.Writer, from io.Reader) (returnedErr error) {
 	output, err := handleSmudge(ctx, cfg, from)
 	if err != nil {
 		return fmt.Errorf("smudging contents: %w", err)
diff --git a/cmd/gitaly-lfs-smudge/lfs_smudge_test.go b/cmd/gitaly-lfs-smudge/lfs_smudge_test.go
index 477a2131f..b63531124 100644
--- a/cmd/gitaly-lfs-smudge/lfs_smudge_test.go
+++ b/cmd/gitaly-lfs-smudge/lfs_smudge_test.go
@@ -10,6 +10,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/git/smudge"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitaly/config"
 	"gitlab.com/gitlab-org/gitaly/v15/internal/gitlab"
+	"gitlab.com/gitlab-org/gitaly/v15/internal/testhelper"
 )
 
 const (
@@ -61,6 +62,8 @@ func TestSuccessfulLfsSmudge(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
+			ctx := testhelper.Context(t)
+
 			var b bytes.Buffer
 			reader := strings.NewReader(tc.data)
 
@@ -76,7 +79,7 @@ func TestSuccessfulLfsSmudge(t *testing.T) {
 				},
 			}
 
-			require.NoError(t, smudgeContents(cfg, &b, reader))
+			require.NoError(t, smudgeContents(ctx, cfg, &b, reader))
 			require.Equal(t, testData, b.String())
 		})
 	}
@@ -169,13 +172,15 @@ func TestUnsuccessfulLfsSmudge(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.desc, func(t *testing.T) {
+			ctx := testhelper.Context(t)
+
 			gitlabCfg, cleanup := runTestServer(t, tc.options)
 			defer cleanup()
 
 			cfg := tc.setupCfg(t, gitlabCfg)
 
 			var b bytes.Buffer
-			err := smudgeContents(cfg, &b, strings.NewReader(tc.data))
+			err := smudgeContents(ctx, cfg, &b, strings.NewReader(tc.data))
 
 			if tc.expectedError {
 				require.Error(t, err)
diff --git a/cmd/gitaly-lfs-smudge/main.go b/cmd/gitaly-lfs-smudge/main.go
index 9db0fbd55..c9b9fcc4b 100644
--- a/cmd/gitaly-lfs-smudge/main.go
+++ b/cmd/gitaly-lfs-smudge/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"os"
@@ -10,6 +11,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/v15/internal/helper/env"
 	gitalylog "gitlab.com/gitlab-org/gitaly/v15/internal/log"
 	"gitlab.com/gitlab-org/labkit/log"
+	"gitlab.com/gitlab-org/labkit/tracing"
 )
 
 func requireStdin(msg string) {
@@ -59,12 +61,18 @@ func initLogging(environment []string) (io.Closer, error) {
 }
 
 func run(environment []string, out io.Writer, in io.Reader) error {
+	// Since the environment is sanitized at the moment, we're only
+	// using this to extract the correlation ID. The finished() call
+	// to clean up the tracing will be a NOP here.
+	ctx, finished := tracing.ExtractFromEnv(context.Background())
+	defer finished()
+
 	cfg, err := smudge.ConfigFromEnvironment(environment)
 	if err != nil {
 		return fmt.Errorf("loading configuration: %w", err)
 	}
 
-	if err := smudgeContents(cfg, out, in); err != nil {
+	if err := smudgeContents(ctx, cfg, out, in); err != nil {
 		return fmt.Errorf("running smudge filter: %w", err)
 	}