Fix race conditions in polling client communicator Flush and Reset

Fixes concurrency and synchronization issues in the polling client communicator around Flush() and Reset() while preserving the original scheduling, throttling, and lifecycle
behaviors:

1. Replaces the single-consumer `pollComplete` channel with a broadcast channel closure mechanism (`pollDone chan struct{}`), ensuring multiple concurrent `Flush()` callers all wake up properly without channel read contention.
2. Replaces `pollDone` on `Reset()` and outbox enqueue to ensure `Flush()` waits for the appropriate poll cycle.
3. Differentiates `Reset()`-induced request cancellations from network failures (`wasReset`), preserving queued messages for immediate re-poll rather than NACKing them.
4. Preserves `lastActive` dynamic latency reduction, `oldestUnsent` buffer delay tracking, `FailureSuicideTimeSeconds` suicide timeout, and shutdown flush behavior.

PiperOrigin-RevId: 949430003

Author: Copybara Service

fleetspeak/src/client/https/polling.go

@@ -59,15 +59,19 @@ type Communicator struct {
 
 	certBytes []byte
 
-	wakeUp       chan struct{}
-	pollComplete chan error
-	mu           sync.Mutex
-	pollCancel   context.CancelFunc
+	// Synchronization for Reset/Flush
+	wakeUp      chan struct{}
+	mu          sync.Mutex
+	pollDone    chan struct{}
+	lastPollErr error
+	pollCancel  context.CancelFunc
 }
 
 // Setup implements comms.Communicator.
 func (c *Communicator) Setup(cl comms.Context) error {
 	c.cctx = cl
+	c.pollDone = make(chan struct{})
+	close(c.pollDone) // Start closed so Flush returns immediately if no messages.
 	return c.configure()
 }
 
@@ -114,7 +118,6 @@ 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.certBytes = certBytes
 	return nil
 }
@@ -141,13 +144,9 @@ func (c *Communicator) Reset() {
 	if c.pollCancel != nil {
 		c.pollCancel()
 	}
+	c.pollDone = make(chan struct{})
 	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:
@@ -158,10 +157,23 @@ func (c *Communicator) Reset() {
 func (c *Communicator) Flush(ctx context.Context) error {
 	log.InfoContextf(ctx, "Flush called")
 	for {
+		c.mu.Lock()
+		done := c.pollDone
+		c.mu.Unlock()
+
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-		case err := <-c.pollComplete:
+		case <-c.ctx.Done():
+			return c.ctx.Err()
+		case <-done:
+			c.mu.Lock()
+			err := c.lastPollErr
+			if err != nil && c.pollDone != done {
+				c.mu.Unlock()
+				continue
+			}
+			c.mu.Unlock()
 			return err
 		}
 	}
@@ -193,25 +205,48 @@ func (c *Communicator) processingLoop() {
 	// for the MinFailureDelay.
 	poll := func() {
 		var err error
+		c.mu.Lock()
+		select {
+		case <-c.pollDone:
+			c.pollDone = make(chan struct{})
+		default:
+		}
+		reqCtx, cancel := context.WithCancel(c.ctx)
+		c.pollCancel = cancel
+		myDone := c.pollDone
+		c.mu.Unlock()
+
 		defer func() {
-			select {
-			case c.pollComplete <- err:
-			default:
-			}
+			cancel()
+			c.mu.Lock()
+			c.pollCancel = nil
+			c.lastPollErr = err
+			close(myDone)
+			c.mu.Unlock()
 		}()
+
 		c.wd.Reset()
 		if c.cctx.CurrentID() != c.id {
 			c.configure()
 		}
 		var active bool
-		active, err = c.poll(toSend)
+		active, err = c.poll(reqCtx, toSend)
 		if err != nil {
+			c.mu.Lock()
+			wasReset := errors.Is(err, context.Canceled) && c.ctx.Err() == nil
+			c.mu.Unlock()
+
+			if wasReset {
+				return
+			}
+
 			log.Warningf("Failure during polling: %v", err)
 			for _, m := range toSend {
 				m.Nack()
 			}
 			toSend = nil
 			toSendSize = 0
+			oldestUnsent = time.Time{}
 
 			if (!lastPoll.IsZero()) && (time.Since(lastPoll) > time.Duration(c.conf.FailureSuicideTimeSeconds)*time.Second) {
 				// Die in the hopes that our replacement will be better configured, or otherwise have better luck.
@@ -310,6 +345,14 @@ func (c *Communicator) processingLoop() {
 			poll()
 		case m := <-c.cctx.Outbox():
 			t.Stop()
+			c.mu.Lock()
+			select {
+			case <-c.pollDone:
+				c.pollDone = make(chan struct{})
+			default:
+			}
+			c.mu.Unlock()
+
 			toSend = append(toSend, m)
 			toSendSize += 2 + proto.Size(m.M)
 			if toSendSize >= sendBytesThreshold ||
@@ -324,7 +367,7 @@ func (c *Communicator) processingLoop() {
 	}
 }
 
-func (c *Communicator) poll(toSend []comms.MessageInfo) (bool, error) {
+func (c *Communicator) poll(ctx context.Context, toSend []comms.MessageInfo) (bool, error) {
 	var sent bool // records whether an interesting (non-LOW) priority message was sent.
 	msgs := make([]*fspb.Message, 0, len(toSend))
 	for _, m := range toSend {
@@ -347,9 +390,12 @@ func (c *Communicator) poll(toSend []comms.MessageInfo) (bool, error) {
 	}
 
 	for i, host := range c.hosts {
-		cd, err := c.pollHost(host, data)
+		cd, err := c.pollHost(ctx, host, data)
 		if err != nil {
 			log.Warningf("Error polling %q for ContactData: %v", host, err)
+			if ctx.Err() != nil {
+				return false, ctx.Err()
+			}
 			continue
 		}
 		if i != 0 {
@@ -367,7 +413,7 @@ func (c *Communicator) poll(toSend []comms.MessageInfo) (bool, error) {
 	return false, errors.New("unable to contact any server")
 }
 
-func (c *Communicator) pollHost(host string, data []byte) (*fspb.ContactData, error) {
+func (c *Communicator) pollHost(ctx context.Context, host string, data []byte) (*fspb.ContactData, error) {
 	var sendErr, recvErr error
 	var sendSize, recvSize int
 	defer func() {
@@ -396,17 +442,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(ctx)
 	SetContentEncoding(req.Header, c.conf.GetCompression())
 	if c.clientCertificateHeader != "" {
 		bc := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: c.certBytes})