Fix fork bomb

instrument/instrument.go

@@ -8,6 +8,7 @@ package instrument
 import (
 	"context"
 	"fmt"
+	"os"
 
 	"github.com/DataDog/orchestrion/instrument/event"
 
@@ -31,7 +32,7 @@ func Report(ctx context.Context, e event.Event, metadata ...any) context.Context
 		var ok bool
 		span, ok = tracer.SpanFromContext(ctx)
 		if !ok {
-			_, _ = fmt.Printf("Error: Received end/return event but have no corresponding span in the context.\n")
+			_, _ = fmt.Fprintf(os.Stderr, "Error: Received end/return event but have no corresponding span in the context.\n")
 			return ctx
 		}
 		span.Finish()

internal/jobserver/pkgs/resolve.go

@@ -137,7 +137,7 @@ func (s *service) resolve(req *ResolveRequest) (ResolveResponse, error) {
 
 		goFlags, err := goflags.Flags()
 		if err != nil {
-			return nil, fmt.Errorf("resolving go build flags: %w", err)
+			log.Warnf("Failed to obtain go build flags: %v\n", err)
 		}
 		goFlags.Trim(
 			"-a",        // Re-building everything here would be VERY expensive, as we'd re-build a lot of stuff multiple times

internal/jobserver/pkgs/resolve_test.go

@@ -7,10 +7,13 @@ package pkgs_test
 
 import (
 	"context"
+	"fmt"
 	"math/rand"
 	"os"
+	"os/exec"
 	"testing"
 
+	"github.com/DataDog/orchestrion/internal/goflags"
 	"github.com/DataDog/orchestrion/internal/jobserver"
 	"github.com/DataDog/orchestrion/internal/jobserver/client"
 	"github.com/DataDog/orchestrion/internal/jobserver/pkgs"
@@ -19,81 +22,111 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestCache(t *testing.T) {
-	server, err := jobserver.New(nil)
+func Test(t *testing.T) {
+	// Force the goflags so we don't get tainted by the `go test` flags!
+	wd, err := os.Getwd()
 	require.NoError(t, err)
-	defer server.Shutdown()
+	goflags.SetFlags(wd, []string{"test"})
 
-	conn, err := server.Connect()
-	require.NoError(t, err)
-	defer conn.Close()
-
-	env := os.Environ()
-
-	// First request is expected to always be a cache miss
-	resp, err := client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
-		context.Background(),
-		conn,
-		&pkgs.ResolveRequest{
-			Pattern: "net/http",
-			Env:     env,
-		},
-	)
-	require.NoError(t, err)
-	require.GreaterOrEqual(t, len(resp), 2)
-	require.EqualValues(t, 1, server.CacheStats.Count())
-	require.EqualValues(t, 0, server.CacheStats.Hits())
-
-	// Second request is equivalent, and should result in a cache hit. The order
-	// of entries in `env` is also shuffled, which should have no impact on the
-	// cache hitting or missing.
-	rand.Shuffle(len(env), func(i, j int) { env[i], env[j] = env[j], env[i] })
-	resp, err = client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
-		context.Background(),
-		conn,
-		&pkgs.ResolveRequest{
-			Pattern: "net/http",
-			Env:     env, // This was shuffled, so it's not the same as before
-		},
-	)
-	require.NoError(t, err)
-	require.GreaterOrEqual(t, len(resp), 2)
-	require.EqualValues(t, 2, server.CacheStats.Count())
-	require.EqualValues(t, 1, server.CacheStats.Hits())
-
-	// Third request is different, should result in a cache miss again
-	resp, err = client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
-		context.Background(),
-		conn,
-		&pkgs.ResolveRequest{
-			Pattern: "os", // Not the same package as before...
-			Env:     env,
-		},
-	)
-	require.NoError(t, err)
-	require.GreaterOrEqual(t, len(resp), 3)
-	require.EqualValues(t, 3, server.CacheStats.Count())
-	require.EqualValues(t, 1, server.CacheStats.Hits())
+	t.Run("Cache", func(t *testing.T) {
+
+		server, err := jobserver.New(nil)
+		require.NoError(t, err)
+		defer server.Shutdown()
+
+		conn, err := server.Connect()
+		require.NoError(t, err)
+		defer conn.Close()
+
+		env := os.Environ()
+
+		// First request is expected to always be a cache miss
+		resp, err := client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
+			context.Background(),
+			conn,
+			&pkgs.ResolveRequest{
+				Pattern: "net/http",
+				Env:     env,
+			},
+		)
+		require.NoError(t, err)
+		assert.GreaterOrEqual(t, len(resp), 2)
+		assert.EqualValues(t, 1, server.CacheStats.Count())
+		assert.EqualValues(t, 0, server.CacheStats.Hits())
+
+		// Second request is equivalent, and should result in a cache hit. The order
+		// of entries in `env` is also shuffled, which should have no impact on the
+		// cache hitting or missing.
+		rand.Shuffle(len(env), func(i, j int) { env[i], env[j] = env[j], env[i] })
+		resp, err = client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
+			context.Background(),
+			conn,
+			&pkgs.ResolveRequest{
+				Pattern: "net/http",
+				Env:     env, // This was shuffled, so it's not the same as before
+			},
+		)
+		require.NoError(t, err)
+		assert.GreaterOrEqual(t, len(resp), 2)
+		assert.EqualValues(t, 2, server.CacheStats.Count())
+		assert.EqualValues(t, 1, server.CacheStats.Hits())
+
+		// Third request is different, should result in a cache miss again
+		resp, err = client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
+			context.Background(),
+			conn,
+			&pkgs.ResolveRequest{
+				Pattern: "os", // Not the same package as before...
+				Env:     env,
+			},
+		)
+		require.NoError(t, err)
+		assert.GreaterOrEqual(t, len(resp), 3)
+		assert.EqualValues(t, 3, server.CacheStats.Count())
+		assert.EqualValues(t, 1, server.CacheStats.Hits())
+	})
+
+	t.Run("Error", func(t *testing.T) {
+		log.SetLevel(log.LevelTrace)
+		t.Cleanup(func() { log.SetLevel(log.LevelNone) })
+
+		server, err := jobserver.New(nil)
+		require.NoError(t, err)
+		defer server.Shutdown()
+
+		conn, err := server.Connect()
+		require.NoError(t, err)
+		defer conn.Close()
+
+		resp, err := client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
+			context.Background(),
+			conn,
+			&pkgs.ResolveRequest{BuildFlags: []string{"-definitely-not-a-valid-build=flag"}, Pattern: "definitely.not/a@valid\x01package"},
+		)
+		assert.Nil(t, resp)
+		assert.EqualValues(t, 0, server.CacheStats.Hits())
+		require.Error(t, err)
+	})
 }
 
-func TestError(t *testing.T) {
-	log.SetLevel(log.LevelTrace)
-	t.Cleanup(func() { log.SetLevel(log.LevelNone) })
+func init() {
+	if len(os.Args) <= 2 || os.Args[1] != "toolexec" {
+		return
+	}
 
-	server, err := jobserver.New(nil)
-	require.NoError(t, err)
-	defer server.Shutdown()
+	// We're invoked with `toolexec` so pretend we're a toolexec proxy...
+	cmd := exec.Command(os.Args[2], os.Args[3:]...)
+	cmd.Stdin = os.Stdin
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
 
-	conn, err := server.Connect()
-	require.NoError(t, err)
-	defer conn.Close()
-
-	resp, err := client.Request[*pkgs.ResolveRequest, pkgs.ResolveResponse](
-		context.Background(),
-		conn,
-		&pkgs.ResolveRequest{BuildFlags: []string{"-definitely-not-a-valid-build=flag"}, Pattern: "definitely.not/a@valid\x00package"},
-	)
-	assert.Nil(t, resp)
-	assert.EqualValues(t, 0, server.CacheStats.Hits())
-	require.Error(t, err)
+	if err := cmd.Run(); err != nil {
+		if err, ok := err.(*exec.ExitError); ok {
+			os.Exit(err.ExitCode())
+			return
+		}
+		_, _ = fmt.Fprintln(os.Stderr, err)
+		os.Exit(1)
+	}
+	os.Exit(0)
 }