From 32ac83b15e92374c85383bf423dd9caba8f48d70 Mon Sep 17 00:00:00 2001 From: Andras Banki-Horvath Date: Thu, 19 Feb 2026 21:43:24 +0100 Subject: [PATCH] mailbox: fix AwaitRPC lost wakeup race MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TestClient_ConcurrentInFlightDoesNotDrop was flaking with "context deadline exceeded" due to a lost-wakeup race in AwaitRPC. AwaitRPC checked the pending-response map and registered a waiter in two separate steps under different lock acquisitions. If handleEnvelope cached a response between the pending check and waiter registration, it notified only the existing waiters — the new waiter was never signaled and blocked until its context timed out. Fix this by introducing popPendingOrAddWaiter, which atomically returns a pending response or registers a waiter, closing the missed-notification window. --- mailbox/client/client.go | 37 ++++++++++++++++--------------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/mailbox/client/client.go b/mailbox/client/client.go index 145c88af4..f9f9b9697 100644 --- a/mailbox/client/client.go +++ b/mailbox/client/client.go @@ -179,14 +179,15 @@ func (c *Client) AwaitRPC(ctx context.Context, correlationID string, resp proto.Message) error { for { - data, ok := c.popPending(correlationID) - if ok { - return (proto.UnmarshalOptions{ + data, ch, hasPending := c.popPendingOrAddWaiter(correlationID) + if hasPending { + unmarshal := proto.UnmarshalOptions{ DiscardUnknown: true, - }).Unmarshal(data, resp) + } + + return unmarshal.Unmarshal(data, resp) } - ch := c.addWaiter(correlationID) select { case <-ch: case <-ctx.Done(): @@ -341,31 +342,25 @@ func (c *Client) handleEnvelope(env *mailboxpb.Envelope) { delete(c.waiters, correlationID) } -// popPending returns and removes a cached response for correlationID. -func (c *Client) popPending(correlationID string) ([]byte, bool) { +// popPendingOrAddWaiter atomically checks for a pending response and, if one +// is not present, registers a waiter channel for a future response. +func (c *Client) popPendingOrAddWaiter(correlationID string) ( + []byte, chan struct{}, bool) { + c.mu.Lock() defer c.mu.Unlock() data, ok := c.pending[correlationID] - if !ok { - return nil, false - } - - delete(c.pending, correlationID) + if ok { + delete(c.pending, correlationID) - return data, true -} + return data, nil, true + } -// addWaiter registers a waiter for correlationID and returns its channel. -func (c *Client) addWaiter(correlationID string) chan struct{} { ch := make(chan struct{}) - - c.mu.Lock() - defer c.mu.Unlock() - c.waiters[correlationID] = append(c.waiters[correlationID], ch) - return ch + return nil, ch, false } // removeWaiter removes a previously registered waiter channel.