feat: support non-root build paths (#209)

Author: txabman42

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/urfave/cli/v3 v3.6.1
 	golang.org/x/mod v0.30.0
 	golang.org/x/sync v0.18.0
+	golang.org/x/tools v0.38.0
 	gopkg.in/yaml.v3 v3.0.1
 	gotest.tools/v3 v3.5.2
 )
@@ -21,6 +22,5 @@ require (
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/rogpeppe/go-internal v1.13.1 // indirect
-	golang.org/x/tools v0.38.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
 )

pkg/instrumentation/grpc/go.mod

@@ -4,4 +4,4 @@ go 1.23.0
 
 replace github.com/open-telemetry/opentelemetry-go-compile-instrumentation/pkg => ../..
 
-require github.com/open-telemetry/opentelemetry-go-compile-instrumentation/pkg 947df235a956
+require github.com/open-telemetry/opentelemetry-go-compile-instrumentation/pkg 947df235a956

tool/internal/setup/add.go

@@ -5,6 +5,9 @@ package setup
 
 import (
 	"fmt"
+	"maps"
+	"path/filepath"
+	"slices"
 
 	"github.com/dave/dst"
 
@@ -27,9 +30,10 @@ func genImportDecl(matched []*rule.InstFuncRule) []dst.Decl {
 	for _, m := range matched {
 		requiredImports[m.Path] = ast.IdentIgnore
 	}
-	importDecls := make([]dst.Decl, 0)
-	for k, v := range requiredImports {
-		importDecls = append(importDecls, ast.ImportDecl(v, k))
+	importDecls := make([]dst.Decl, 0, len(requiredImports))
+	// Sort the keys to ensure deterministic order
+	for _, k := range slices.Sorted(maps.Keys(requiredImports)) {
+		importDecls = append(importDecls, ast.ImportDecl(requiredImports[k], k))
 	}
 	return importDecls
 }
@@ -96,7 +100,9 @@ func buildOtelRuntimeAst(decls []dst.Decl) *dst.File {
 	}
 }
 
-func (sp *SetupPhase) addDeps(matched []*rule.InstRuleSet) error {
+// addDeps generates and writes otel.runtime.go with required imports and variable
+// declarations for OpenTelemetry instrumentation based on matched rules.
+func (sp *SetupPhase) addDeps(matched []*rule.InstRuleSet, packagePath string) error {
 	rules := make([]*rule.InstFuncRule, 0)
 	for _, m := range matched {
 		funcRules := m.GetFuncRules()
@@ -113,10 +119,12 @@ func (sp *SetupPhase) addDeps(matched []*rule.InstRuleSet) error {
 	// Build the ast
 	root := buildOtelRuntimeAst(append(importDecls, varDecls...))
 	// Write the ast to file
-	err := ast.WriteFile(OtelRuntimeFile, root)
+	otelRuntimeFilePath := filepath.Join(packagePath, OtelRuntimeFile)
+	err := ast.WriteFile(otelRuntimeFilePath, root)
 	if err != nil {
 		return err
 	}
-	sp.keepForDebug(OtelRuntimeFile)
+	sp.keepForDebug(otelRuntimeFilePath)
+	sp.Info("Created otel.runtime.go", "path", otelRuntimeFilePath)
 	return nil
 }

tool/internal/setup/add_test.go

@@ -0,0 +1,136 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+//go:build !windows
+
+// Package setup tests verify that the addDeps function generates
+// the expected otel.runtime.go file by comparing against golden files.
+//
+// To update golden files after intentional changes:
+//
+//	go test -update ./tool/internal/setup/...
+
+package setup
+
+import (
+	"io"
+	"log/slog"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/open-telemetry/opentelemetry-go-compile-instrumentation/tool/internal/rule"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gotest.tools/v3/golden"
+)
+
+func TestAddDeps(t *testing.T) {
+	tests := []struct {
+		name       string
+		matched    []*rule.InstRuleSet
+		goldenFile string // Empty means no file should be generated
+	}{
+		{
+			name:       "empty_matched_rules",
+			matched:    []*rule.InstRuleSet{},
+			goldenFile: "",
+		},
+		{
+			name: "single_func_rule",
+			matched: []*rule.InstRuleSet{
+				newTestRuleSet(
+					"github.com/example/pkg",
+					newTestFuncRule("github.com/example/pkg", "github.com/example/pkg"),
+				),
+			},
+			goldenFile: "single_func_rule.otel.runtime.go.golden",
+		},
+		{
+			name: "no_func_rules",
+			matched: []*rule.InstRuleSet{
+				newTestRuleSet("github.com/example/pkg"),
+			},
+			goldenFile: "",
+		},
+		{
+			name: "multiple_rule_sets",
+			matched: []*rule.InstRuleSet{
+				newTestRuleSet(
+					"github.com/example/pkg1",
+					newTestFuncRule("github.com/example/pkg1", "github.com/example/pkg1"),
+				),
+				newTestRuleSet(
+					"github.com/example/pkg2",
+					newTestFuncRule("github.com/example/pkg2", "github.com/example/pkg2"),
+				),
+			},
+			goldenFile: "multiple_rule_sets.otel.runtime.go.golden",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			tmpDir := t.TempDir()
+			sp := newTestSetupPhase()
+
+			err := sp.addDeps(tt.matched, tmpDir)
+			require.NoError(t, err)
+
+			runtimeFilePath := filepath.Join(tmpDir, OtelRuntimeFile)
+
+			if tt.goldenFile == "" {
+				assert.NoFileExists(t, runtimeFilePath)
+				return
+			}
+
+			assert.FileExists(t, runtimeFilePath)
+			actual, err := os.ReadFile(runtimeFilePath)
+			require.NoError(t, err)
+
+			golden.Assert(t, string(actual), tt.goldenFile)
+		})
+	}
+}
+
+func TestAddDeps_FileWriteError(t *testing.T) {
+	matched := []*rule.InstRuleSet{
+		newTestRuleSet(
+			"github.com/example/pkg",
+			newTestFuncRule("github.com/example/pkg", "github.com/example/pkg"),
+		),
+	}
+
+	// Use a non-existent parent directory to cause write error
+	invalidPath := filepath.Join(t.TempDir(), "nonexistent", "subdir")
+	sp := newTestSetupPhase()
+
+	err := sp.addDeps(matched, invalidPath)
+	assert.Error(t, err)
+}
+
+// Helper functions for constructing test data
+
+func newTestSetupPhase() *SetupPhase {
+	return &SetupPhase{
+		logger: slog.New(slog.NewTextHandler(io.Discard, nil)),
+	}
+}
+
+func newTestFuncRule(path, target string) *rule.InstFuncRule {
+	return &rule.InstFuncRule{
+		InstBaseRule: rule.InstBaseRule{
+			Target: target,
+		},
+		Path: path,
+	}
+}
+
+func newTestRuleSet(modulePath string, funcRules ...*rule.InstFuncRule) *rule.InstRuleSet {
+	rs := rule.NewInstRuleSet(modulePath)
+	fakeFilePath := filepath.Join(os.TempDir(), "file.go")
+	for _, fr := range funcRules {
+		rs.AddFuncRule(fakeFilePath, fr)
+	}
+	return rs
+}

tool/internal/setup/setup.go

@@ -9,9 +9,11 @@ import (
 	"log/slog"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/open-telemetry/opentelemetry-go-compile-instrumentation/tool/ex"
 	"github.com/open-telemetry/opentelemetry-go-compile-instrumentation/tool/util"
+	"golang.org/x/tools/go/packages"
 )
 
 type SetupPhase struct {
@@ -39,6 +41,90 @@ func isSetup() bool {
 	return false
 }
 
+// flagsWithPathValues contains flags that accept a value from "go build" command.
+//
+//nolint:gochecknoglobals // private lookup table
+var flagsWithPathValues = map[string]bool{
+	"-C":             true,
+	"-o":             true,
+	"-p":             true,
+	"-covermode":     true,
+	"-coverpkg":      true,
+	"-asmflags":      true,
+	"-buildmode":     true,
+	"-buildvcs":      true,
+	"-compiler":      true,
+	"-gccgoflags":    true,
+	"-gcflags":       true,
+	"-installsuffix": true,
+	"-ldflags":       true,
+	"-mod":           true,
+	"-modfile":       true,
+	"-overlay":       true,
+	"-pgo":           true,
+	"-pkgdir":        true,
+	"-tags":          true,
+	"-toolexec":      true,
+}
+
+// GetBuildPackages loads all packages from the go build command arguments.
+// Returns a list of loaded packages. If no package patterns are found in args,
+// defaults to loading the current directory package.
+// The args parameter should be the go build command arguments (e.g., ["build", "-a", "./cmd"]).
+// Returns an error if package loading fails or if invalid patterns are provided.
+// For example:
+//   - args ["build", "-a", "./cmd"] returns packages for "./cmd"
+//   - args ["build", "-a", "cmd"] returns packages for the "cmd" package in the module
+//   - args ["build", "-a", ".", "./cmd"] returns packages for both "." and "./cmd"
+//   - args ["build"] returns packages for "."
+func getBuildPackages(ctx context.Context, args []string) ([]*packages.Package, error) {
+	logger := util.LoggerFromContext(ctx)
+
+	buildPkgs := make([]*packages.Package, 0)
+	cfg := &packages.Config{
+		Mode: packages.NeedName | packages.NeedFiles | packages.NeedModule,
+	}
+	found := false
+	for i := len(args) - 1; i >= 0; i-- {
+		arg := args[i]
+
+		// If preceded by a flag that takes a path value, this is a flag value
+		// We want to avoid scenarios like "go build -o ./tmp ./app" where tmp also contains Go files,
+		// as it would be treated as a package.
+		if i > 0 && flagsWithPathValues[args[i-1]] {
+			break
+		}
+
+		// If we hit a flag, stop. Packages come after all flags
+		// go build [-o output] [build flags] [packages]
+		if strings.HasPrefix(arg, "-") || arg == "go" || arg == "build" || arg == "install" {
+			break
+		}
+
+		pkgs, err := packages.Load(cfg, arg)
+		if err != nil {
+			return nil, ex.Wrapf(err, "failed to load packages for pattern %s", arg)
+		}
+		for _, pkg := range pkgs {
+			if pkg.Errors != nil || pkg.Module == nil {
+				logger.DebugContext(ctx, "skipping package", "pattern", arg, "errors", pkg.Errors, "module", pkg.Module)
+				continue
+			}
+			buildPkgs = append(buildPkgs, pkg)
+			found = true
+		}
+	}
+
+	if !found {
+		var err error
+		buildPkgs, err = packages.Load(cfg, ".")
+		if err != nil {
+			return nil, ex.Wrapf(err, "failed to load packages for pattern .")
+		}
+	}
+	return buildPkgs, nil
+}
+
 // Setup prepares the environment for further instrumentation.
 func Setup(ctx context.Context, args []string) error {
 	logger := util.LoggerFromContext(ctx)
@@ -51,6 +137,14 @@ func Setup(ctx context.Context, args []string) error {
 	sp := &SetupPhase{
 		logger: logger,
 	}
+
+	// Introduce additional hook code by generating otel.runtime.go
+	// Use GetPackage to determine the build target directory
+	pkgs, err := getBuildPackages(ctx, args)
+	if err != nil {
+		return err
+	}
+
 	// Find all dependencies of the project being build
 	deps, err := sp.findDeps(ctx, args)
 	if err != nil {
@@ -61,10 +155,11 @@ func Setup(ctx context.Context, args []string) error {
 	if err != nil {
 		return err
 	}
-	// Introduce additional hook code by generating otel.instrumentation.go
-	err = sp.addDeps(matched)
-	if err != nil {
-		return err
+	for _, pkg := range pkgs {
+		// Introduce additional hook code by generating otel.runtime.go
+		if err = sp.addDeps(matched, pkg.Dir); err != nil {
+			return err
+		}
 	}
 	// Extract the embedded instrumentation modules into local directory
 	err = sp.extract()
@@ -77,11 +172,7 @@ func Setup(ctx context.Context, args []string) error {
 		return err
 	}
 	// Write the matched hook to matched.txt for further instrument phase
-	err = sp.store(matched)
-	if err != nil {
-		return err
-	}
-	return nil
+	return sp.store(matched)
 }
 
 // BuildWithToolexec builds the project with the toolexec mode
@@ -128,16 +219,21 @@ func GoBuild(ctx context.Context, args []string) error {
 		logger.DebugContext(ctx, "failed to back up files", "error", err)
 	}
 	defer func() {
-		err = os.RemoveAll(OtelRuntimeFile)
+		var pkgs []*packages.Package
+		pkgs, err = getBuildPackages(ctx, args)
 		if err != nil {
-			logger.DebugContext(ctx, "failed to remove otel runtime file", "error", err)
+			logger.DebugContext(ctx, "failed to get build packages", "error", err)
 		}
-		err = os.RemoveAll(unzippedPkgDir)
-		if err != nil {
+		for _, pkg := range pkgs {
+			if err = os.RemoveAll(filepath.Join(pkg.Dir, OtelRuntimeFile)); err != nil {
+				logger.DebugContext(ctx, "failed to remove generated file from package",
+					"file", filepath.Join(pkg.Dir, OtelRuntimeFile), "error", err)
+			}
+		}
+		if err = os.RemoveAll(unzippedPkgDir); err != nil {
 			logger.DebugContext(ctx, "failed to remove unzipped pkg", "error", err)
 		}
-		err = util.RestoreFile(backupFiles)
-		if err != nil {
+		if err = util.RestoreFile(backupFiles); err != nil {
 			logger.DebugContext(ctx, "failed to restore files", "error", err)
 		}
 	}()