fix: OnCompileMain needs to register new link.deps (#501)

The `OnCompile` hook cannot necessarily resolve the transitive
dependencies of all link-time dependencies, as those are normally used
to avoid creating dependency cycles.

The `OnCompileMain` hook however MUST resolve the transitive
dependencies of everything, as it needs to introduce `import`s of all
synthetic link-time dependencies so that they are correctly registered
(their `init` functions get scheduled, and they get presented to the
linker). This was previously not necessarily done, and could result in
link-time failures.

---

This involved re-working the way the `link.deps` file is added to the
`_pkg_.a` files so that both the `OnCompile` and `OnCompileMain`
functions contribute to the same closure and it gets written &
registered exactly once. As a result `Command.OnClose` was removed (no
longer used).

internal/cmd/toolexec.go

@@ -30,7 +30,7 @@ var Toolexec = &cli.Command{
 		if err != nil {
 			return err
 		}
-		defer proxyCmd.Close()
+		defer proxyCmd.Close(ctx.Context)
 
 		if proxyCmd.Type() == proxy.CommandTypeOther {
 			// Immediately run the command if it's of the Other type, as we do not do

internal/toolexec/aspect/linkdeps/linkdeps.go

@@ -205,8 +205,10 @@ func readArchiveData(archive string, entry string) (io.Reader, error) {
 	var list, data bytes.Buffer
 	cmd := exec.Command("go", "tool", "pack", "t", archive)
 	cmd.Stdout = &list
+	var stderr bytes.Buffer
+	cmd.Stderr = &stderr
 	if err := cmd.Run(); err != nil {
-		return nil, fmt.Errorf("running `go tool pack t %q`: %w", archive, err)
+		return nil, fmt.Errorf("running `go tool pack t %q`: %w\n%s", archive, err, stderr.String())
 	}
 	for {
 		line, err := list.ReadString('\n')

internal/toolexec/aspect/oncompile-main.go

@@ -55,17 +55,38 @@ func (Weaver) OnCompileMain(ctx context.Context, cmd *proxy.CompileCommand) erro
 	newDeps := linkDeps.Dependencies()
 
 	// Add package resolutions of link-time dependencies to the importcfg file:
-	for _, linkDepPath := range newDeps {
+	stack := append(make([]string, 0, len(newDeps)), newDeps...)
+	for len(stack) > 0 {
+		// Pop from the stack of things to process...
+		linkDepPath := stack[len(stack)-1]
+		stack = stack[:len(stack)-1]
+
 		deps, err := resolvePackageFiles(ctx, linkDepPath, cmd.WorkDir)
 		if err != nil {
 			return fmt.Errorf("resolving %q: %w", linkDepPath, err)
 		}
+
 		for p, a := range deps {
 			if _, found := reg.PackageFile[p]; found {
 				continue
 			}
 			log.Debug().Str("import-path", p).Str("archive", a).Msg("Recording resolved " + linkdeps.Filename + " dependency")
 			reg.PackageFile[p] = a
+
+			// The package may have its own link-time dependencies we need to resolve
+			tDeps, err := linkdeps.FromArchive(a)
+			if err != nil {
+				return fmt.Errorf("reading %s from %s[%s]: %w", linkdeps.Filename, p, a, err)
+			}
+			for _, tDep := range tDeps.Dependencies() {
+				if reg.PackageFile[tDep] != "" || slices.Contains(stack, tDep) {
+					// Already resolved, or already going to be resolved...
+					continue
+				}
+				stack = append(stack, tDep)     // Push it to the stack
+				newDeps = append(newDeps, tDep) // Record it as asynthetic import to add
+				cmd.LinkDeps.Add(tDep)          // Record it as a link-time dependency
+			}
 		}
 	}
 

internal/toolexec/aspect/oncompile.go

@@ -9,7 +9,6 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"os/exec"
 	"path/filepath"
 	"slices"
 	"strings"
@@ -74,24 +73,8 @@ func (w Weaver) OnCompile(ctx context.Context, cmd *proxy.CompileCommand) (err e
 		return fmt.Errorf("parsing %q: %w", cmd.Flags.ImportCfg, err)
 	}
 
-	linkDeps, err := linkdeps.FromImportConfig(&imports)
-	if err != nil {
-		return fmt.Errorf("reading %s closure from %s: %w", linkdeps.Filename, cmd.Flags.ImportCfg, err)
-	}
-
 	orchestrionDir := filepath.Join(filepath.Dir(cmd.Flags.Output), "orchestrion")
 
-	// Ensure we correctly register the [linkdeps.Filename] into the output
-	// archive upon returning, even if we made no changes. The contract is that
-	// an archive's [linkdeps.Filename] must represent all transitive link-time
-	// dependencies.
-	defer func() {
-		if err != nil {
-			return
-		}
-		err = writeLinkDeps(log, cmd, &linkDeps, orchestrionDir)
-	}()
-
 	goMod, err := goenv.GOMOD(".")
 	if err != nil {
 		return fmt.Errorf("go env GOMOD: %w", err)
@@ -182,13 +165,15 @@ func (w Weaver) OnCompile(ctx context.Context, cmd *proxy.CompileCommand) (err e
 		}
 
 		log.Debug().Stringer("kind", kind).Str("import-path", depImportPath).Msg("Recording synthetic " + linkdeps.Filename + " dependency")
-		linkDeps.Add(depImportPath)
+		cmd.LinkDeps.Add(depImportPath)
 
-		if kind != typed.ImportStatement {
+		if kind != typed.ImportStatement && cmd.Flags.Package != "main" {
 			// We cannot attempt to resolve link-time dependencies (relocation targets), as these are
 			// typically used to avoid creating dependency cycles. Corrollary to this, the `link.deps`
 			// file will not contain transitive closures for these packages, so we need to resolve these
-			// at link-time.
+			// at link-time. If the package being built is "main", then we can ignore this, as we are at
+			// the top-level of a dependency tree anyway, and if we cannot resolve a dependency, then we
+			// will not be able to link the final binary.
 			continue
 		}
 
@@ -206,7 +191,7 @@ func (w Weaver) OnCompile(ctx context.Context, cmd *proxy.CompileCommand) (err e
 			for _, tDep := range deps.Dependencies() {
 				if _, found := imports.PackageFile[tDep]; !found {
 					log.Debug().Str("import-path", dep).Str("transitive", tDep).Str("inherited-from", depImportPath).Msg("Copying transitive " + linkdeps.Filename + " dependency")
-					linkDeps.Add(tDep)
+					cmd.LinkDeps.Add(tDep)
 				}
 			}
 
@@ -245,34 +230,3 @@ func writeUpdatedImportConfig(log zerolog.Logger, reg importcfg.ImportConfig, fi
 
 	return nil
 }
-
-// writeLinkDeps writes the [linkdeps.Filename] file into the orchestrionDir,
-// and registers it to be packed into the output archive. Does nothing if the
-// provided [linkdeps.LinkDeps] is empty.
-func writeLinkDeps(log zerolog.Logger, cmd *proxy.CompileCommand, linkDeps *linkdeps.LinkDeps, orchestrionDir string) error {
-	if linkDeps.Empty() {
-		// Nothing to do...
-		return nil
-	}
-
-	// Write the link.deps file and add it to the output object once the compilation has completed.
-	if err := os.MkdirAll(orchestrionDir, 0o755); err != nil {
-		return fmt.Errorf("making directory %s: %w", orchestrionDir, err)
-	}
-
-	linkDepsFile := filepath.Join(orchestrionDir, linkdeps.Filename)
-	if err := linkDeps.WriteFile(linkDepsFile); err != nil {
-		return fmt.Errorf("writing %s file: %w", linkdeps.Filename, err)
-	}
-
-	cmd.OnClose(func() error {
-		log.Debug().Str("archive", cmd.Flags.Output).Array(linkdeps.Filename, linkDeps).Msg("Adding " + linkdeps.Filename + " file in archive")
-		child := exec.Command("go", "tool", "pack", "r", cmd.Flags.Output, linkDepsFile)
-		if err := child.Run(); err != nil {
-			return fmt.Errorf("running %q: %w", child.Args, err)
-		}
-		return nil
-	})
-
-	return nil
-}

internal/toolexec/proxy/compile.go

@@ -6,11 +6,18 @@
 package proxy
 
 import (
+	gocontext "context"
 	"errors"
+	"fmt"
+	"os"
+	"os/exec"
 	"path/filepath"
 	"strings"
 
 	"github.com/DataDog/orchestrion/internal/injector/aspect/context"
+	"github.com/DataDog/orchestrion/internal/toolexec/aspect/linkdeps"
+	"github.com/DataDog/orchestrion/internal/toolexec/importcfg"
+	"github.com/rs/zerolog"
 )
 
 //go:generate go run github.com/DataDog/orchestrion/internal/toolexec/proxy/generator -command=compile
@@ -30,6 +37,11 @@ type CompileCommand struct {
 	Files []string
 	// WorkDir is the $WORK directory managed by the go toolchain.
 	WorkDir string
+
+	// Link-time dependencies that must be honored to link a dependent of the
+	// built package. If not blank, this is written to disk, then appended to the
+	// archive output.
+	LinkDeps linkdeps.LinkDeps
 }
 
 func (*CompileCommand) Type() CommandType { return CommandTypeCompile }
@@ -105,6 +117,40 @@ func (cmd *CompileCommand) AddFiles(files []string) {
 	}
 }
 
+func (cmd *CompileCommand) Close(ctx gocontext.Context) (err error) {
+	defer func() { err = errors.Join(err, cmd.command.Close(ctx)) }()
+
+	if cmd.LinkDeps.Empty() {
+		return nil
+	}
+
+	if _, err := os.Stat(cmd.Flags.Output); errors.Is(err, os.ErrNotExist) {
+		// Already failing, not doing anything...
+		return nil
+	} else if err != nil {
+		return err
+	}
+
+	orchestrionDir := filepath.Join(cmd.Flags.Output, "..", "orchestrion")
+	if err := os.MkdirAll(orchestrionDir, 0o755); err != nil {
+		return fmt.Errorf("mkdir %q: %w", orchestrionDir, err)
+	}
+
+	linkDepsFile := filepath.Join(orchestrionDir, linkdeps.Filename)
+	if err := cmd.LinkDeps.WriteFile(linkDepsFile); err != nil {
+		return fmt.Errorf("writing %s file: %w", linkdeps.Filename, err)
+	}
+
+	log := zerolog.Ctx(ctx)
+	log.Debug().Str("archive", cmd.Flags.Output).Array(linkdeps.Filename, &cmd.LinkDeps).Msg("Adding " + linkdeps.Filename + " file in archive")
+
+	child := exec.Command("go", "tool", "pack", "r", cmd.Flags.Output, linkDepsFile)
+	if err := child.Run(); err != nil {
+		return fmt.Errorf("running %q: %w", child.Args, err)
+	}
+	return nil
+}
+
 func parseCompileCommand(args []string) (*CompileCommand, error) {
 	if len(args) == 0 {
 		return nil, errors.New("unexpected number of command arguments")
@@ -119,6 +165,16 @@ func parseCompileCommand(args []string) (*CompileCommand, error) {
 	if cmd.Flags.ImportCfg != "" {
 		// The WorkDir is the parent of the stage directory, which is where the importcfg file is located.
 		cmd.WorkDir = filepath.Dir(filepath.Dir(cmd.Flags.ImportCfg))
+
+		imports, err := importcfg.ParseFile(cmd.Flags.ImportCfg)
+		if err != nil {
+			return nil, fmt.Errorf("parsing %q: %w", cmd.Flags.ImportCfg, err)
+		}
+
+		cmd.LinkDeps, err = linkdeps.FromImportConfig(&imports)
+		if err != nil {
+			return nil, fmt.Errorf("reading %s closure from %s: %w", linkdeps.Filename, cmd.Flags.ImportCfg, err)
+		}
 	}
 
 	return &cmd, nil

internal/toolexec/proxy/proxy.go

@@ -22,7 +22,7 @@ type (
 	Command interface {
 		// Close invokes all registered OnClose callbacks and releases any resources associated with the
 		// command.
-		Close() error
+		Close(context.Context) error
 
 		// Args are all the command arguments, starting from the Go tool command
 		Args() []string
@@ -48,7 +48,6 @@ type (
 		args []string
 		// paramPos is the index in args of the *value* provided for the parameter stored in the key
 		paramPos map[string]int
-		onClose  []func() error
 	}
 )
 
@@ -84,18 +83,7 @@ func NewCommand(args []string) command {
 	return cmd
 }
 
-// OnClose registers a callback to be invoked when the command is closed, usually after it has run,
-// unless skipping was requested by the integration.
-func (cmd *command) OnClose(cb func() error) {
-	cmd.onClose = append(cmd.onClose, cb)
-}
-
-func (cmd *command) Close() error {
-	for _, cb := range cmd.onClose {
-		if err := cb(); err != nil {
-			return err
-		}
-	}
+func (*command) Close(context.Context) error {
 	return nil
 }