From 7abe4d23194aea14cf2fda4856fe08bdc3d99bfc Mon Sep 17 00:00:00 2001 From: Copybara Service Date: Tue, 18 Aug 2026 08:53:03 -0700 Subject: [PATCH] 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: 966620893 --- fleetspeak/src/client/https/polling.go | 96 ++++++++++++++++++-------- 1 file changed, 66 insertions(+), 30 deletions(-) diff --git a/fleetspeak/src/client/https/polling.go b/fleetspeak/src/client/https/polling.go index 53a332a2..d194cd62 100644 --- a/fleetspeak/src/client/https/polling.go +++ b/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})