fix: auto-pin race condition on go.mod file

internal/filelock/filelock.go

@@ -22,8 +22,8 @@ import (
 // Upgrading a read-lock to a write lock, or vice-versa, is not guaranteed to
 // happen atomically (on Windows, it is guaranteed not to be atomic).
 type Mutex struct {
-	path   string
 	file   *os.File
+	path   string
 	locked lockState
 }
 

internal/pin/auto_test.go

@@ -19,6 +19,13 @@ import (
 )
 
 func TestAutoPin(t *testing.T) {
+	ctx := context.Background()
+	if d, ok := t.Deadline(); ok {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithDeadline(ctx, d)
+		defer cancel()
+	}
+
 	t.Run("simple", func(t *testing.T) {
 		tmp := scaffold(t, make(map[string]string))
 		chdir(t, tmp)
@@ -30,7 +37,7 @@ func TestAutoPin(t *testing.T) {
 		assert.FileExists(t, filepath.Join(tmp, config.FilenameOrchestrionToolGo))
 		assert.FileExists(t, filepath.Join(tmp, "go.sum"))
 
-		data, err := parseGoMod(filepath.Join(tmp, "go.mod"))
+		data, err := parseGoMod(ctx, filepath.Join(tmp, "go.mod"))
 		require.NoError(t, err)
 
 		rawTag, _ := version.TagInfo()
@@ -45,7 +52,7 @@ func TestAutoPin(t *testing.T) {
 
 		t.Setenv(envVarCheckedGoMod, "true")
 
-		AutoPinOrchestrion(context.Background(), io.Discard, io.Discard)
+		AutoPinOrchestrion(ctx, io.Discard, io.Discard)
 
 		assert.NoFileExists(t, filepath.Join(tmp, config.FilenameOrchestrionToolGo))
 		assert.NoFileExists(t, filepath.Join(tmp, "go.sum"))

internal/pin/gomod.go

@@ -7,6 +7,7 @@
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -47,9 +48,9 @@
 
 // parseGoMod processes the contents of the designated `go.mod` file using
 // `go mod edit -json` and returns the corresponding parsed [goMod].
-func parseGoMod(modfile string) (goMod, error) {
+func parseGoMod(ctx context.Context, modfile string) (goMod, error) {
 	var stdout bytes.Buffer
-	if err := runGoMod("edit", modfile, &stdout, "-json"); err != nil {
+	if err := runGoMod(ctx, "edit", modfile, &stdout, "-json"); err != nil {
 		return goMod{}, fmt.Errorf("running `go mod edit -json`: %w", err)
 	}
 
@@ -74,8 +75,8 @@
 
 // runGoGet executes the `go get <modSpecs...>` subcommand with the provided
 // module specifications on the designated `go.mod` file.
-func runGoGet(modfile string, modSpecs ...string) error {
-	cmd := exec.Command("go", "get", "-modfile", modfile)
+func runGoGet(ctx context.Context, modfile string, modSpecs ...string) error {
+	cmd := exec.CommandContext(ctx, "go", "get", "-modfile", modfile)
 	cmd.Args = append(cmd.Args, modSpecs...)
 	cmd.Env = append(os.Environ(), "GOTOOLCHAIN=local")
 	cmd.Stdin = os.Stdin
@@ -87,8 +88,8 @@
 // runGoMod executes the `go mod <command> <args...>` subcommand with the
 // provided arguments on the designated `go.mod` file, sending standard output
 // to the provided writer.
-func runGoMod(command string, modfile string, stdout io.Writer, args ...string) error {
-	cmd := exec.Command("go", "mod", command, "-modfile", modfile)
+func runGoMod(ctx context.Context, command string, modfile string, stdout io.Writer, args ...string) error {
+	cmd := exec.CommandContext(ctx, "go", "mod", command, "-modfile", modfile)
 	cmd.Args = append(cmd.Args, args...)
 	cmd.Env = append(os.Environ(), "GOTOOLCHAIN=local")
 	cmd.Stdin = os.Stdin
@@ -99,7 +100,7 @@
 
 // runGoModEdit makes the specified changes to the `go.mod` file, then runs `go mod tidy` if needed.
 // If there is a `vendor` directory, it also runs `go mod vendor` before returning.
-func runGoModEdit(modfile string, edits ...goModEdit) error {
+func runGoModEdit(ctx context.Context, modfile string, edits ...goModEdit) error {
 	if len(edits) == 0 {
 		// Nothing to do.
 		return nil
@@ -109,11 +110,11 @@
 	for i, edit := range edits {
 		editFlags[i] = edit.goModEditFlag()
 	}
-	if err := runGoMod("edit", modfile, os.Stdout, editFlags...); err != nil {
+	if err := runGoMod(ctx, "edit", modfile, os.Stdout, editFlags...); err != nil {
 		return fmt.Errorf("running `go mod edit %s`: %w", editFlags, err)
 	}
 
-	if err := runGoMod("tidy", modfile, os.Stdout); err != nil {
+	if err := runGoMod(ctx, "tidy", modfile, os.Stdout); err != nil {
 		return fmt.Errorf("running `go mod tidy`: %w", err)
 	}
 
@@ -127,7 +128,7 @@
 		return fmt.Errorf("checking for vendor directory %q: %w", vendorDir, err)
 	}
 
-	if err := runGoMod("vendor", modfile, os.Stdout); err != nil {
+	if err := runGoMod(ctx, "vendor", modfile, os.Stdout); err != nil {
 		return fmt.Errorf("running `go mod vendor`: %w", err)
 	}
 

internal/pin/pin.go

@@ -8,6 +8,8 @@
 import (
 	"bytes"
 	"context"
+	"crypto/sha512"
+	"encoding/base64"
 	"errors"
 	"fmt"
 	"go/parser"
@@ -18,6 +20,7 @@
 	"path/filepath"
 	"strings"
 
+	"github.com/DataDog/orchestrion/internal/filelock"
 	"github.com/DataDog/orchestrion/internal/goenv"
 	"github.com/DataDog/orchestrion/internal/injector/config"
 	"github.com/DataDog/orchestrion/internal/version"
@@ -71,6 +74,22 @@
 		return fmt.Errorf("getting GOMOD: %w", err)
 	}
 
+	// Acquire an advisory lock on the `go.mod` file, so that in `-toolexec` mode,
+	// multiple attempts to auto-pin don't try to modify the files at the same
+	// time. The `go mod tidy` command takes an advisory write-lock on `go.mod`,
+	// so we are using a separate file under [os.TempDir] to avoid deadlocking.
+	sha := sha512.Sum512([]byte(goMod))
+	flockname := filepath.Join(os.TempDir(), "orchestrion-pin_"+base64.URLEncoding.EncodeToString(sha[:])+"_go.mod.lock")
+	flock := filelock.MutexAt(flockname)
+	if err := flock.Lock(); err != nil {
+		return fmt.Errorf("failed to acquire lock on %q: %w", goMod, err)
+	}
+	defer func() {
+		if err := flock.Unlock(); err != nil {
+			log.Error().Err(err).Str("lock-file", goMod).Msg("Failed to release file lock")
+		}
+	}()
+
 	toolFile := filepath.Join(goMod, "..", config.FilenameOrchestrionToolGo)
 	dstFile, err := parseOrchestrionToolGo(toolFile)
 	if errors.Is(err, os.ErrNotExist) {
@@ -96,14 +115,14 @@
 
 	// Add the current version of orchestrion to the `go.mod` file.
 	var edits []goModEdit
-	curMod, err := parseGoMod(goMod)
+	curMod, err := parseGoMod(ctx, goMod)
 	if err != nil {
 		return fmt.Errorf("parsing %q: %w", goMod, err)
 	}
 
 	if ver, found := curMod.requires(datadogTracerV1); !found || semver.Compare(ver, "v1.72.0-rc.1") < 0 {
 		log.Info().Msg("Installing or upgrading " + datadogTracerV1)
-		if err := runGoGet(goMod, datadogTracerV1+"@>=v1.72.0-rc.1"); err != nil {
+		if err := runGoGet(ctx, goMod, datadogTracerV1+"@>=v1.72.0-rc.1"); err != nil {
 			return fmt.Errorf("go get "+datadogTracerV1+": %w", err)
 		}
 	}
@@ -119,7 +138,7 @@
 		edits = append(edits, goModRequire{Path: orchestrionImportPath, Version: version})
 	}
 
-	if err := runGoModEdit(goMod, edits...); err != nil {
+	if err := runGoModEdit(ctx, goMod, edits...); err != nil {
 		return fmt.Errorf("editing %q: %w", goMod, err)
 	}
 
@@ -130,13 +149,13 @@
 
 	if pruned {
 		// Run "go mod tidy" to ensure the `go.mod` file is up-to-date with detected dependencies.
-		if err := runGoMod("tidy", goMod, nil); err != nil {
+		if err := runGoMod(ctx, "tidy", goMod, nil); err != nil {
 			return fmt.Errorf("running `go mod tidy`: %w", err)
 		}
 	}
 
 	// Restore the previous toolchain directive if `go mod tidy` had the nerve to touch it...
-	if err := runGoModEdit(goMod, curMod.Toolchain); err != nil {
+	if err := runGoModEdit(ctx, goMod, curMod.Toolchain); err != nil {
 		return fmt.Errorf("restoring toolchain directive: %w", err)
 	}
 

internal/pin/pin_test.go

@@ -21,16 +21,23 @@ import (
 )
 
 func TestPin(t *testing.T) {
+	ctx := context.Background()
+	if d, ok := t.Deadline(); ok {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithDeadline(ctx, d)
+		defer cancel()
+	}
+
 	t.Run("simple", func(t *testing.T) {
 		tmp := scaffold(t, make(map[string]string))
 		chdir(t, tmp)
 
-		require.NoError(t, PinOrchestrion(context.Background(), Options{Writer: io.Discard, ErrWriter: io.Discard}))
+		require.NoError(t, PinOrchestrion(ctx, Options{Writer: io.Discard, ErrWriter: io.Discard}))
 
 		assert.FileExists(t, filepath.Join(tmp, config.FilenameOrchestrionToolGo))
 		assert.FileExists(t, filepath.Join(tmp, "go.sum"))
 
-		data, err := parseGoMod(filepath.Join(tmp, "go.mod"))
+		data, err := parseGoMod(ctx, filepath.Join(tmp, "go.mod"))
 		require.NoError(t, err)
 
 		rawTag, _ := version.TagInfo()
@@ -46,12 +53,12 @@ func TestPin(t *testing.T) {
 		tmp := scaffold(t, map[string]string{"github.com/DataDog/orchestrion": "v0.9.3"})
 		chdir(t, tmp)
 
-		require.NoError(t, PinOrchestrion(context.Background(), Options{Writer: io.Discard, ErrWriter: io.Discard}))
+		require.NoError(t, PinOrchestrion(ctx, Options{Writer: io.Discard, ErrWriter: io.Discard}))
 
 		assert.FileExists(t, filepath.Join(tmp, config.FilenameOrchestrionToolGo))
 		assert.FileExists(t, filepath.Join(tmp, "go.sum"))
 
-		data, err := parseGoMod(filepath.Join(tmp, "go.mod"))
+		data, err := parseGoMod(ctx, filepath.Join(tmp, "go.mod"))
 		require.NoError(t, err)
 
 		rawTag, _ := version.TagInfo()
@@ -62,7 +69,7 @@ func TestPin(t *testing.T) {
 		tmp := scaffold(t, make(map[string]string))
 		chdir(t, tmp)
 
-		require.NoError(t, PinOrchestrion(context.Background(), Options{Writer: io.Discard, ErrWriter: io.Discard, NoGenerate: true}))
+		require.NoError(t, PinOrchestrion(ctx, Options{Writer: io.Discard, ErrWriter: io.Discard, NoGenerate: true}))
 
 		content, err := os.ReadFile(filepath.Join(tmp, config.FilenameOrchestrionToolGo))
 		require.NoError(t, err)
@@ -74,9 +81,9 @@ func TestPin(t *testing.T) {
 		tmp := scaffold(t, map[string]string{"github.com/digitalocean/sample-golang": "v0.0.0-20240904143939-1e058723dcf4"})
 		chdir(t, tmp)
 
-		require.NoError(t, PinOrchestrion(context.Background(), Options{Writer: io.Discard, ErrWriter: io.Discard, NoGenerate: true}))
+		require.NoError(t, PinOrchestrion(ctx, Options{Writer: io.Discard, ErrWriter: io.Discard, NoGenerate: true}))
 
-		data, err := parseGoMod(filepath.Join(tmp, "go.mod"))
+		data, err := parseGoMod(ctx, filepath.Join(tmp, "go.mod"))
 		require.NoError(t, err)
 
 		assert.NotContains(t, data.Require, goModRequire{"github.com/digitalocean/sample-golang", "v0.0.0-20240904143939-1e058723dcf4"})
@@ -89,9 +96,9 @@ func TestPin(t *testing.T) {
 		})
 		chdir(t, tmp)
 
-		require.NoError(t, PinOrchestrion(context.Background(), Options{Writer: io.Discard, ErrWriter: io.Discard, NoGenerate: true}))
+		require.NoError(t, PinOrchestrion(ctx, Options{Writer: io.Discard, ErrWriter: io.Discard, NoGenerate: true}))
 
-		data, err := parseGoMod(filepath.Join(tmp, "go.mod"))
+		data, err := parseGoMod(ctx, filepath.Join(tmp, "go.mod"))
 		require.NoError(t, err)
 
 		assert.NotContains(t, data.Require, goModRequire{"github.com/digitalocean/sample-golang", "v0.0.0-20240904143939-1e058723dcf4"})
@@ -105,7 +112,7 @@ func TestPin(t *testing.T) {
 		toolDotGo := filepath.Join(tmp, config.FilenameOrchestrionToolGo)
 		require.NoError(t, os.WriteFile(toolDotGo, nil, 0644))
 
-		require.ErrorContains(t, PinOrchestrion(context.Background(), Options{Writer: io.Discard, ErrWriter: io.Discard}), "expected 'package', found 'EOF'")
+		require.ErrorContains(t, PinOrchestrion(ctx, Options{Writer: io.Discard, ErrWriter: io.Discard}), "expected 'package', found 'EOF'")
 	})
 }