Stop two grpctunnel tests racing their own teardown

agent/server/grpctunnel/response_header_test.go

@@ -33,7 +33,7 @@ func TestResponse_CarriesRelayInstanceHeader(t *testing.T) {
 
 	fc := newFrameCollector()
 	sc := newTestStreamCtx(fc.sendFn)
-	table := newCallTable()
+	table := newTestCallTable(t)
 
 	tc.startCall(sc, table, "c1", reqStart("GET", "/thing", 0))
 	tc.handleCallFrame(sc, table, &pb.CallFrame{CallId: "c1", Body: &pb.CallFrame_End{End: &pb.CallEnd{}}})
@@ -64,7 +64,7 @@ func TestResponse_OmitsRelayInstanceHeaderWhenUnset(t *testing.T) {
 
 	fc := newFrameCollector()
 	sc := newTestStreamCtx(fc.sendFn)
-	table := newCallTable()
+	table := newTestCallTable(t)
 
 	tc.startCall(sc, table, "c1", reqStart("GET", "/thing", 0))
 	tc.handleCallFrame(sc, table, &pb.CallFrame{CallId: "c1", Body: &pb.CallFrame_End{End: &pb.CallEnd{}}})

agent/server/grpctunnel/tunnel_client_test.go

@@ -368,6 +368,37 @@ func (fc *frameCollector) waitDone(t *testing.T, n int) {
 	}
 }
 
+// newTestCallTable returns a call table that holds the test open until every
+// call goroutine it tracks has actually finished.
+//
+// A call's last act is table.remove, in the outermost defer of the goroutine
+// startCall launches — which, being registered first, runs after runCall's own
+// deferred "Request completed" line. So an empty table is the only barrier
+// that says nothing will log any more. Waiting on the terminal frame is not
+// enough: that log defer runs after the frame is sent.
+//
+// Without this a test that returns while a call is still in flight panics the
+// whole package with "Log in goroutine after <test> has completed" — zaptest
+// logging into a finished test. It is intermittent, because it depends on
+// whether the goroutine gets scheduled before the test returns.
+func newTestCallTable(t *testing.T) *callTable {
+	t.Helper()
+	table := newCallTable()
+	t.Cleanup(func() {
+		// Cancel whatever is still running first, so a test that deliberately
+		// leaves a slow call in flight does not pay its full delay here.
+		// cancelAll empties the map itself and so cannot be the barrier —
+		// cancel in place and let each goroutine's own defer do the removing.
+		table.mu.Lock()
+		for _, c := range table.calls {
+			c.cancel()
+		}
+		table.mu.Unlock()
+		waitFor(t, 10*time.Second, func() bool { return !table.active() })
+	})
+	return table
+}
+
 func reqStart(method, path string, timeoutMs int32) *pb.CallStart {
 	return &pb.CallStart{
 		PseudoHeaders: map[string]string{":method": method, ":path": path},
@@ -483,8 +514,16 @@ func TestInitialConnectRetry_RefusedDial(t *testing.T) {
 	go func() { _ = grpcServer.Serve(lis2) }()
 	defer grpcServer.Stop()
 
+	// >= 1, not == 1: connectionsActive counts streams, and this client runs
+	// two idle streams per server, so the gauge settles at 2. It passes
+	// through 1 on the way, but waitFor polls every 10ms and the two slots
+	// connect within a few of each other — in CI they landed 6ms apart, so
+	// the poll saw 0 then 2 and the equality never held. What this test is
+	// about is that the client connects at all after a refused dial, which
+	// is what >= 1 says. TestClose_ShutsDownCleanly already spells it that
+	// way against the same gauge.
 	waitFor(t, 10*time.Second, func() bool {
-		return gaugeVecValue(t, tc.connectionsActive, "delayed-server") == 1
+		return gaugeVecValue(t, tc.connectionsActive, "delayed-server") >= 1
 	})
 }
 
@@ -544,7 +583,7 @@ func TestInflightCap_QueuesThenRuns(t *testing.T) {
 
 	fc := newFrameCollector()
 	sc := newTestStreamCtx(fc.sendFn)
-	table := newCallTable()
+	table := newTestCallTable(t)
 
 	tc.startCall(sc, table, "r1", reqStart("GET", "/", 0))
 	tc.handleCallFrame(sc, table, &pb.CallFrame{CallId: "r1", Body: &pb.CallFrame_End{End: &pb.CallEnd{}}})
@@ -582,7 +621,7 @@ func TestInflightCap_QueueTimeout(t *testing.T) {
 
 	fc := newFrameCollector()
 	sc := newTestStreamCtx(fc.sendFn)
-	table := newCallTable()
+	table := newTestCallTable(t)
 
 	tc.startCall(sc, table, "r1", reqStart("GET", "/", 0))
 	tc.handleCallFrame(sc, table, &pb.CallFrame{CallId: "r1", Body: &pb.CallFrame_End{End: &pb.CallEnd{}}})
@@ -623,7 +662,7 @@ func TestMaxRequestTimeout_AppliesWhenTimeoutMsZero(t *testing.T) {
 
 	fc := newFrameCollector()
 	sc := newTestStreamCtx(fc.sendFn)
-	table := newCallTable()
+	table := newTestCallTable(t)
 
 	start := time.Now()
 	tc.startCall(sc, table, "r1", reqStart("GET", "/", 0))
@@ -670,7 +709,7 @@ func TestCall_StreamedRequestBody(t *testing.T) {
 
 	fc := newFrameCollector()
 	sc := newTestStreamCtx(fc.sendFn)
-	table := newCallTable()
+	table := newTestCallTable(t)
 
 	tc.startCall(sc, table, "c1", reqStart("POST", "/upload", 0))
 	tc.handleCallFrame(sc, table, &pb.CallFrame{CallId: "c1", Body: &pb.CallFrame_Data{Data: &pb.CallData{Payload: []byte("part-1;")}}})
@@ -702,7 +741,7 @@ func TestCall_ServerCancelAbortsBackend(t *testing.T) {
 
 	fc := newFrameCollector()
 	sc := newTestStreamCtx(fc.sendFn)
-	table := newCallTable()
+	table := newTestCallTable(t)
 
 	tc.startCall(sc, table, "c1", reqStart("GET", "/slow", 0))
 	tc.handleCallFrame(sc, table, &pb.CallFrame{CallId: "c1", Body: &pb.CallFrame_End{End: &pb.CallEnd{}}})