From 9e14214b1200675d10d6ff71349c12009e838bfe Mon Sep 17 00:00:00 2001
From: Copybara Service <copybara-servicebot@google.com>
Date: Thu, 16 Jul 2026 06:29:31 -0700
Subject: [PATCH] Fix race condition and test flakiness in polling communicator
 Flush and Reset.

In polling communicator, Reset forcefully canceled the context of active poll HTTP requests. The canceled poll deferred writing errors to pollComplete, which raced with Reset's non-blocking channel drain. Subsequent Flush calls received the canceled poll's error ("unable to contact any server").

Refactored polling communicator to use a clean broadcast channel pattern (pollDone) for Flush unblocking without forcibly canceling in-flight HTTP requests during Reset.

PiperOrigin-RevId: 948936881
---
 fleetspeak/src/client/https/polling.go | 59 ++++++++++----------------
 1 file changed, 22 insertions(+), 37 deletions(-)

diff --git a/fleetspeak/src/client/https/polling.go b/fleetspeak/src/client/https/polling.go
index 53a332a2..8505b157 100644
--- a/fleetspeak/src/client/https/polling.go
+++ b/fleetspeak/src/client/https/polling.go
@@ -59,10 +59,10 @@ type Communicator struct {
 
 	certBytes []byte
 
-	wakeUp       chan struct{}
-	pollComplete chan error
-	mu           sync.Mutex
-	pollCancel   context.CancelFunc
+	wakeUp      chan struct{}
+	mu          sync.Mutex
+	pollDone    chan struct{}
+	lastPollErr error
 }
 
 // Setup implements comms.Communicator.
@@ -114,7 +114,7 @@ func (c *Communicator) configure() error {
 	c.ctx, c.done = context.WithCancel(context.Background())
 	c.clientCertificateHeader = si.ClientCertificateHeader
 	c.wakeUp = make(chan struct{}, 1)
-	c.pollComplete = make(chan error, 1)
+	c.pollDone = make(chan struct{})
 	c.certBytes = certBytes
 	return nil
 }
@@ -137,17 +137,7 @@ func (c *Communicator) Stop() {
 // Reset implements comms.Communicator.
 func (c *Communicator) Reset() {
 	log.Infof("Reset called")
-	c.mu.Lock()
-	if c.pollCancel != nil {
-		c.pollCancel()
-	}
-	c.mu.Unlock()
 	c.hc.Transport.(*http.Transport).CloseIdleConnections()
-	// Drain pollComplete to ensure Flush waits for a new poll.
-	select {
-	case <-c.pollComplete:
-	default:
-	}
 	select {
 	case c.wakeUp <- struct{}{}:
 	default:
@@ -157,13 +147,17 @@ func (c *Communicator) Reset() {
 // Flush implements comms.Communicator.
 func (c *Communicator) Flush(ctx context.Context) error {
 	log.InfoContextf(ctx, "Flush called")
-	for {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case err := <-c.pollComplete:
-			return err
-		}
+	c.mu.Lock()
+	done := c.pollDone
+	c.mu.Unlock()
+
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-done:
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		return c.lastPollErr
 	}
 }
 
@@ -194,10 +188,11 @@ func (c *Communicator) processingLoop() {
 	poll := func() {
 		var err error
 		defer func() {
-			select {
-			case c.pollComplete <- err:
-			default:
-			}
+			c.mu.Lock()
+			c.lastPollErr = err
+			close(c.pollDone)
+			c.pollDone = make(chan struct{})
+			c.mu.Unlock()
 		}()
 		c.wd.Reset()
 		if c.cctx.CurrentID() != c.id {
@@ -396,17 +391,7 @@ func (c *Communicator) pollHost(host string, data []byte) (*fspb.ContactData, er
 	if sendErr != nil {
 		return nil, sendErr
 	}
-	var reqCtx context.Context
-	c.mu.Lock()
-	reqCtx, c.pollCancel = context.WithCancel(c.ctx)
-	c.mu.Unlock()
-	defer func() {
-		c.mu.Lock()
-		c.pollCancel()
-		c.pollCancel = nil
-		c.mu.Unlock()
-	}()
-	req = req.WithContext(reqCtx)
+	req = req.WithContext(c.ctx)
 	SetContentEncoding(req.Header, c.conf.GetCompression())
 	if c.clientCertificateHeader != "" {
 		bc := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: c.certBytes})