fix(ci): serve drain race and danger test timeout

.github/workflows/test.yml

@@ -27,7 +27,7 @@ jobs:
         run: go build -o /tmp/odek-test github.com/BackendStack21/odek/cmd/odek && chmod +x /tmp/odek-test
 
       - name: test
-        run: ODEK_BINARY=/tmp/odek-test go test ./... -short -race -coverprofile=coverage.out -covermode=atomic -count=1
+        run: ODEK_BINARY=/tmp/odek-test go test ./... -short -race -timeout 20m -coverprofile=coverage.out -covermode=atomic -count=1
 
       - name: coverage
         run: go tool cover -func=coverage.out

cmd/odek/serve.go

@@ -211,9 +211,50 @@ func cancelPrompt(sessionID string) bool {
 // defers (sandbox cleanup, agent.Close) to run before the process exits.
 var wsConns sync.Map // map[*golangws.Conn]struct{}
 
-// wsHandlerWG counts live handleWS goroutines; serveOnListener waits on it
+// wsHandlers counts live handleWS goroutines; serveOnListener waits on it
 // after closing all connections to ensure cleanup completes.
-var wsHandlerWG sync.WaitGroup
+var wsHandlers liveCounter
+
+// liveCounter counts live goroutines and lets a drain wait until none are
+// left. Unlike sync.WaitGroup, add may run concurrently with wait: a
+// connection accepted while (or after) a drain waits is ordered by the
+// mutex, not left to the WaitGroup rule that an add from zero must happen
+// before Wait.
+type liveCounter struct {
+	mu   sync.Mutex
+	n    int
+	idle chan struct{} // closed when n returns to zero
+}
+
+func (c *liveCounter) add() {
+	c.mu.Lock()
+	if c.n == 0 {
+		c.idle = make(chan struct{})
+	}
+	c.n++
+	c.mu.Unlock()
+}
+
+func (c *liveCounter) done() {
+	c.mu.Lock()
+	c.n--
+	if c.n == 0 {
+		close(c.idle)
+	}
+	c.mu.Unlock()
+}
+
+// wait blocks until no goroutine is counted.
+func (c *liveCounter) wait() {
+	c.mu.Lock()
+	if c.n == 0 {
+		c.mu.Unlock()
+		return
+	}
+	idle := c.idle
+	c.mu.Unlock()
+	<-idle
+}
 
 // sessionLookupLimiter provides basic per-IP rate limiting for session detail
 // lookups, raising the cost of brute-force enumeration of session IDs.
@@ -867,7 +908,7 @@ func serveOnListener(listener net.Listener, mux *http.ServeMux) error {
 
 	// Phase 2: close active WebSocket connections so their handleWS goroutines
 	// unblock from Receive(), run their defers (agent.Close → docker rm -f),
-	// and decrement wsHandlerWG.
+	// and leave wsHandlers.
 	//
 	// http.Server.Shutdown does not forcibly close WebSocket connections because
 	// they are long-lived and never become "idle". Without this step, the
@@ -918,7 +959,7 @@ func serveOnListener(listener net.Listener, mux *http.ServeMux) error {
 func drainServeWork(timeout time.Duration) bool {
 	drained := make(chan struct{})
 	go func() {
-		wsHandlerWG.Wait()
+		wsHandlers.wait()
 		serveRunsWG.Wait()
 		close(drained)
 	}()
@@ -1532,11 +1573,11 @@ func handleWS(store *session.Store, resources *resource.Registry, resolved confi
 	// Register for graceful-shutdown tracking before anything else.
 	// serveOnListener closes all tracked connections on SIGINT/SIGTERM,
 	// which unblocks the Receive loop below and lets defers run.
-	wsHandlerWG.Add(1)
+	wsHandlers.add()
 	wsConns.Store(conn, struct{}{})
 	defer func() {
 		wsConns.Delete(conn)
-		wsHandlerWG.Done()
+		wsHandlers.done()
 	}()
 	defer conn.Close()
 	defer releaseConnWriter(conn) // drop the per-connection write state

cmd/odek/serve_drain_counter_test.go

@@ -0,0 +1,47 @@
+package main
+
+import (
+	"sync"
+	"testing"
+	"time"
+)
+
+// A WebSocket handler can start while a drain from an earlier shutdown is
+// still waiting (tests run several serve instances in one process). With a
+// sync.WaitGroup that add-from-zero concurrent with Wait is a data race the
+// race detector reports; liveCounter orders both under its mutex.
+func TestLiveCounter_AddConcurrentWithWait(t *testing.T) {
+	var c liveCounter
+	var wg sync.WaitGroup
+	for i := 0; i < 200; i++ {
+		wg.Add(2)
+		go func() { defer wg.Done(); c.wait() }()
+		go func() { defer wg.Done(); c.add(); c.done() }()
+	}
+	wg.Wait()
+	c.wait() // all handlers left: returns at once
+}
+
+func TestLiveCounter_WaitBlocksUntilLastDone(t *testing.T) {
+	var c liveCounter
+	c.add()
+	c.add()
+	waited := make(chan struct{})
+	go func() { c.wait(); close(waited) }()
+	c.done()
+	select {
+	case <-waited:
+		t.Fatal("wait returned with a handler still live")
+	case <-time.After(50 * time.Millisecond):
+	}
+	c.done()
+	select {
+	case <-waited:
+	case <-time.After(5 * time.Second):
+		t.Fatal("wait did not return after the last handler left")
+	}
+	// The counter is reusable after it drains.
+	c.add()
+	c.done()
+	c.wait()
+}

internal/danger/compound_commands_test.go

@@ -775,22 +775,28 @@ func TestCompound_WrappingKeepsTheVerdict(t *testing.T) {
 		"[[ a == a ]] && %s",
 		"{ ( %s ); }",
 	}
+	// Every dangerous command times every shape is slow under -race; the
+	// shapes are checked in parallel.
+	t.Parallel()
 	var cfg DangerousConfig
-	for _, danger := range monoDangerous {
-		base := Classify(danger)
-		for _, shape := range shapes {
-			cmd := fmt.Sprintf(shape, danger)
-			wrapped := Classify(cmd)
-			if monoRankDropped(base, wrapped) {
-				t.Errorf("Classify(%q) = %s ranks below Classify(%q) = %s", cmd, wrapped, danger, base)
+	for _, shape := range shapes {
+		t.Run(shape, func(t *testing.T) {
+			t.Parallel()
+			for _, danger := range monoDangerous {
+				base := Classify(danger)
+				cmd := fmt.Sprintf(shape, danger)
+				wrapped := Classify(cmd)
+				if monoRankDropped(base, wrapped) {
+					t.Errorf("Classify(%q) = %s ranks below Classify(%q) = %s", cmd, wrapped, danger, base)
+				}
+				if monoActionRank(cfg.ActionForCommand(cmd)) < monoActionRank(cfg.ActionForCommand(danger)) {
+					t.Errorf("ActionForCommand(%q) = %s is weaker than %q = %s", cmd, cfg.ActionForCommand(cmd), danger, cfg.ActionForCommand(danger))
+				}
+				if !monoEffectsSubset(Analyze(danger), Analyze(cmd)) && !strings.HasPrefix(shape, "!") {
+					t.Errorf("Analyze(%q).Effects = %v lost the effects %v", cmd, Analyze(cmd).Effects, Analyze(danger).Effects)
+				}
 			}
-			if monoActionRank(cfg.ActionForCommand(cmd)) < monoActionRank(cfg.ActionForCommand(danger)) {
-				t.Errorf("ActionForCommand(%q) = %s is weaker than %q = %s", cmd, cfg.ActionForCommand(cmd), danger, cfg.ActionForCommand(danger))
-			}
-			if !monoEffectsSubset(Analyze(danger), Analyze(cmd)) && !strings.HasPrefix(shape, "!") {
-				t.Errorf("Analyze(%q).Effects = %v lost the effects %v", cmd, Analyze(cmd).Effects, Analyze(danger).Effects)
-			}
-		}
+		})
 	}
 }
 

internal/danger/injection_prefilter_test.go

@@ -98,30 +98,41 @@ func testStringLiterals(t *testing.T) []string {
 	return out
 }
 
+// The corpus is every string literal in the package's tests, so it grows with
+// each regression case; it is checked in parallel shards (each with its own
+// mutation stream) to keep the package inside the CI test timeout under -race.
 func TestRED_Injection_PrefilterEquivalentToFullScan(t *testing.T) {
+	t.Parallel()
 	corpus := testStringLiterals(t)
 	for _, p := range injectionPatterns {
 		corpus = append(corpus, p.Label)
 	}
-	rng := rand.New(rand.NewSource(7))
-	check := func(s string) {
-		if got, want := ScanInjection(s), scanInjectionReference(s); !sameResults(got, want) {
-			t.Fatalf("verdict differs for %q:\n got %v\nwant %v", s, got, want)
-		}
-	}
-	for _, s := range corpus {
-		check(s)
-		check(strings.ToUpper(s))
-		// Random mutations: splice corpus entries, drop or duplicate a span.
-		for i := 0; i < 3; i++ {
-			o := corpus[rng.Intn(len(corpus))]
-			if len(s) > 0 {
-				a := rng.Intn(len(s) + 1)
-				b := a + rng.Intn(len(s)-a+1)
-				check(s[:a] + " " + o + " " + s[b:])
-				check(s[:b] + s[a:])
+	const shards = 4
+	for shard := 0; shard < shards; shard++ {
+		t.Run(strconv.Itoa(shard), func(t *testing.T) {
+			t.Parallel()
+			rng := rand.New(rand.NewSource(7 + int64(shard)))
+			check := func(s string) {
+				if got, want := ScanInjection(s), scanInjectionReference(s); !sameResults(got, want) {
+					t.Fatalf("verdict differs for %q:\n got %v\nwant %v", s, got, want)
+				}
 			}
-		}
+			for i := shard; i < len(corpus); i += shards {
+				s := corpus[i]
+				check(s)
+				check(strings.ToUpper(s))
+				// Random mutations: splice corpus entries, drop or duplicate a span.
+				for j := 0; j < 3; j++ {
+					o := corpus[rng.Intn(len(corpus))]
+					if len(s) > 0 {
+						a := rng.Intn(len(s) + 1)
+						b := a + rng.Intn(len(s)-a+1)
+						check(s[:a] + " " + o + " " + s[b:])
+						check(s[:b] + s[a:])
+					}
+				}
+			}
+		})
 	}
 }