flaky tests: fix supervisor and relay manager data races, gate package on -race

agent/Makefile

@@ -96,19 +96,13 @@ test: mocks
 	go test -v ./...
 
 # Packages that currently pass the race detector. The rest of the tree does
-# not yet: server, server/handler, server/http, server/snykbroker and cmd all
-# have known races (supervisor process lifecycle, handler manager, the fx
-# server stack). Move packages up here as those are fixed.
-RACE_CLEAN_PKGS := ./common/... ./config/... ./scaffold/... ./server/api/... ./server/cron/... ./server/snykbroker/acceptfile/... ./test/httpfake/... ./util/...
-
-# The reflector lives in server/snykbroker, whose other tests are not yet
-# race-clean, so it is gated at test granularity instead. Derived from the
-# test files so new reflector tests are covered automatically.
-REFLECTOR_TESTS = $(shell grep -ho '^func Test[A-Za-z0-9_]*' server/snykbroker/reflector*_test.go | sed 's/^func //' | paste -sd'|' -)
+# not yet: server, server/handler, server/http and cmd have known races
+# (handler manager, the fx server stack). Move packages up here as those are
+# fixed.
+RACE_CLEAN_PKGS := ./common/... ./config/... ./scaffold/... ./server/api/... ./server/cron/... ./server/snykbroker/... ./test/httpfake/... ./util/...
 
 test-race: mocks
 	go test -race $(RACE_CLEAN_PKGS)
-	go test -race -run '^($(REFLECTOR_TESTS))$$' ./server/snykbroker/
 
 .PHONY: test-race
 

agent/server/snykbroker/relay_instance_manager.go

@@ -10,6 +10,7 @@ import (
 	"os"
 	"path"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -44,10 +45,13 @@ type RelayInstanceManager interface {
 }
 
 type relayInstanceManager struct {
-	integrationInfo   common.IntegrationInfo
-	registration      Registration
-	config            config.AgentConfig
-	logger            *zap.Logger
+	integrationInfo common.IntegrationInfo
+	registration    Registration
+	config          config.AgentConfig
+	logger          *zap.Logger
+	// supervisor is written by Start's goroutine and taken by Close, from
+	// different goroutines; supervisorMu guards it. Use currentSupervisor.
+	supervisorMu      sync.Mutex
 	supervisor        *Supervisor
 	running           atomic.Bool
 	startCount        atomic.Int32
@@ -563,6 +567,7 @@ func (r *relayInstanceManager) Start() error {
 
 	gen := r.generation.Load()
 	done := make(chan struct{})
+	result := &startResult{}
 
 	// Auto-register: poll for token changes and request restart if changed.
 	if r.config.AutoRegisterFrequency != 0 {
@@ -646,7 +651,7 @@ func (r *relayInstanceManager) Start() error {
 		}
 
 		if errx != nil {
-			err = errx
+			result.set(errx)
 			requestRestartOnExit = true
 			return
 		}
@@ -697,7 +702,7 @@ func (r *relayInstanceManager) Start() error {
 			brokerEnv["LOG_LEVEL"] = "debug"
 		}
 
-		r.supervisor = NewSupervisor(
+		supervisor := NewSupervisor(
 			executable,
 			args,
 			brokerEnv,
@@ -706,17 +711,21 @@ func (r *relayInstanceManager) Start() error {
 		// Only panic on max retries during the very first start.
 		// On subsequent restarts, return the error so the recovery
 		// loop can retry.
-		r.supervisor.panicOnMaxRetries = r.startCount.Load() == 0
+		supervisor.panicOnMaxRetries = r.startCount.Load() == 0
 		r.startCount.Add(1)
 		requestRestartOnExit = true
-		supervisor := r.supervisor
-		err = supervisor.Start(5, 10*time.Second)
-		r.emitOperationCounter("broker_start", err == nil)
-		if err != nil {
-			r.logger.Warn("Supervisor has exited upon startup", zap.Error(err))
+		r.supervisorMu.Lock()
+		r.supervisor = supervisor
+		r.supervisorMu.Unlock()
+		startErr := supervisor.Start(5, 10*time.Second)
+		result.set(startErr)
+		r.emitOperationCounter("broker_start", startErr == nil)
+		if startErr != nil {
+			r.logger.Warn("Supervisor has exited upon startup", zap.Error(startErr))
 			return
 		}
-		err = supervisor.Wait()
+		err := supervisor.Wait()
+		result.set(err)
 		r.emitOperationCounter("broker_exit", err == nil)
 		if err == nil {
 			r.logger.Info("Supervisor has exited")
@@ -729,7 +738,33 @@ func (r *relayInstanceManager) Start() error {
 	case <-done:
 	case <-time.After(r.config.FailWaitTime):
 	}
-	return err
+	return result.get()
+}
+
+// startResult is the outcome Start's goroutine reports: it may still be
+// writing it when Start stops waiting and reads it.
+type startResult struct {
+	mu  sync.Mutex
+	err error
+}
+
+func (s *startResult) set(err error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.err = err
+}
+
+func (s *startResult) get() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.err
+}
+
+// currentSupervisor returns the running broker's supervisor, or nil.
+func (r *relayInstanceManager) currentSupervisor() *Supervisor {
+	r.supervisorMu.Lock()
+	defer r.supervisorMu.Unlock()
+	return r.supervisor
 }
 
 func (r *relayInstanceManager) setHttpProxyEnvVars(brokerEnv map[string]string) {
@@ -873,8 +908,10 @@ func (r *relayInstanceManager) applyClientValidationConfig(validationConfig *com
 func (r *relayInstanceManager) Close() error {
 
 	if r.running.CompareAndSwap(true, false) {
+		r.supervisorMu.Lock()
 		s := r.supervisor
 		r.supervisor = nil
+		r.supervisorMu.Unlock()
 		if s != nil {
 			return s.Close()
 		}

agent/server/snykbroker/relay_instance_manager_test.go

@@ -571,7 +571,8 @@ func TestBrokerSelfHealingRestart(t *testing.T) {
 
 	// Wait for the broker to be running
 	require.Eventually(t, func() bool {
-		return instance.supervisor != nil && instance.supervisor.Pid() > 0
+		s := instance.currentSupervisor()
+		return s != nil && s.Pid() > 0
 	}, 5*time.Second, 10*time.Millisecond, "Broker should be running")
 
 	initialStartCount := instance.startCount.Load()
@@ -580,7 +581,7 @@ func TestBrokerSelfHealingRestart(t *testing.T) {
 	// Kill the broker process to simulate unexpected death.
 	// The supervisor will exhaust retries and the restart consumer
 	// will process the broker_exit restart request.
-	pid := instance.supervisor.Pid()
+	pid := instance.currentSupervisor().Pid()
 	if pid > 0 {
 		process, err := os.FindProcess(pid)
 		require.NoError(t, err)

agent/server/snykbroker/supervisor.go

@@ -246,26 +246,45 @@ func (b *Supervisor) runCommand() error {
 	}()
 
 	var killed atomic.Bool
+	// procMu orders a stop against the start: cmd.Start writes cmd.Process
+	// on this goroutine while a stop may arrive on any other, so the stop
+	// reads the process only through proc, under the lock, and a stop that
+	// lands before the process exists is honoured as soon as it does.
+	var procMu sync.Mutex
+	var proc *os.Process
+	kill := func(p *os.Process) {
+		if err := p.Kill(); err != nil {
+			fmt.Printf("Error killing process: %v\n", err)
+		}
+	}
 
 	// our stopfunc allows anyone to close the process
 	// and wait for it to finish
 	b.setStopFunc(func() {
+		procMu.Lock()
 		killed.Store(true)
+		p := proc
+		procMu.Unlock()
 
 		// if process is running trigger a kill
-		if cmd.Process != nil {
-			err := cmd.Process.Kill()
-			if err != nil {
-				fmt.Printf("Error killing process: %v\n", err)
-			}
+		if p != nil {
+			kill(p)
 		}
 
 		// wait for it to actually finish
 		<-running
 	})
 
 	procStart := time.Now()
+	procMu.Lock()
 	err := cmd.Start()
+	if err == nil {
+		proc = cmd.Process
+		if killed.Load() {
+			kill(proc)
+		}
+	}
+	procMu.Unlock()
 
 	if err != nil {
 		b.lastRunNanos.Store(0)