sync: fix "send on closed channel" panic in responsesWorker at teardown

sync/client.go

@@ -33,10 +33,19 @@ type DefaultClient struct {
 	nextMu     sync.Mutex
 	next       int
 	handlersMu sync.Mutex
-	handlers   map[string]chan *tgsync.Response
+	handlers   map[string]*pendingRequest
 	socket     *websocket.Conn
 }
 
+// pendingRequest is the per-request response handler. done is closed by the
+// request's cancellation goroutine to tell responsesWorker to stop delivering;
+// the channel ch is NEVER closed (responsesWorker is its only sender), so a
+// late response can never panic on send-to-closed.
+type pendingRequest struct {
+	ch   chan *tgsync.Response
+	done chan struct{}
+}
+
 // NewBoundClient returns a new sync DefaultClient that is bound to the provided
 // RunEnv. All operations will be automatically scoped to the keyspace of that
 // run.
@@ -99,7 +108,7 @@ func newClient(ctx context.Context, log *zap.SugaredLogger, extractor func(ctx c
 		cancel:    cancel,
 		log:       log,
 		extractor: extractor,
-		handlers:  map[string]chan *tgsync.Response{},
+		handlers:  map[string]*pendingRequest{},
 	}
 
 	c.sugarOperations = &sugarOperations{c}

sync/client_conn.go

@@ -32,15 +32,21 @@ func (c *DefaultClient) responsesWorker() {
 			c.log.Fatalw("error while reading socket", "error", err)
 		}
 
-		var ch chan *sync.Response
+		var h *pendingRequest
 		c.handlersMu.Lock()
-		ch = c.handlers[res.ID]
+		h = c.handlers[res.ID]
 		c.handlersMu.Unlock()
 
-		if ch == nil {
+		if h == nil {
 			c.log.Warnf("no handler available for response: %s", res.ID)
 		} else {
-			ch <- res
+			// Deliver, or drop if the request was torn down. h.ch is never closed,
+			// so this can never panic on send-to-closed; h.done guarantees we never
+			// block forever on a receiver that has gone away.
+			select {
+			case h.ch <- res:
+			case <-h.done:
+			}
 		}
 	}
 
@@ -57,9 +63,10 @@ func (c *DefaultClient) makeRequest(ctx context.Context, req *sync.Request) (cha
 	}
 
 	ch := make(chan *sync.Response)
+	done := make(chan struct{})
 
 	c.handlersMu.Lock()
-	c.handlers[req.ID] = ch
+	c.handlers[req.ID] = &pendingRequest{ch: ch, done: done}
 	c.handlersMu.Unlock()
 
 	err := c.writeSocket(req)
@@ -77,15 +84,29 @@ func (c *DefaultClient) makeRequest(ctx context.Context, req *sync.Request) (cha
 		}
 
 		c.handlersMu.Lock()
-		close(c.handlers[req.ID])
 		delete(c.handlers, req.ID)
 		c.handlersMu.Unlock()
+		close(done) // tell responsesWorker to stop; ch is intentionally NOT closed
 		c.wg.Done()
 	}()
 
 	return ch, nil
 }
 
+// awaitResponse waits for the single response to a one-shot request, or for the
+// request or client context to be cancelled. The handler channel is never closed,
+// so the context is the only unblock.
+func (c *DefaultClient) awaitResponse(ctx context.Context, ch chan *sync.Response) (*sync.Response, error) {
+	select {
+	case res := <-ch:
+		return res, nil
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-c.ctx.Done():
+		return nil, errors.New("client closed before getting response")
+	}
+}
+
 func (c *DefaultClient) readSocket() (*sync.Response, error) {
 	// After one hour without receiving information from the sync service,
 	// the test will inevitably fail. Note(hacdias): consider changing

sync/client_pubsub.go

@@ -22,9 +22,9 @@ func (c *DefaultClient) publish(ctx context.Context, topic string, payload inter
 		return -1, err
 	}
 
-	res, ok := <-ch
-	if !ok {
-		return -1, errors.New("channel closed before getting response")
+	res, err := c.awaitResponse(ctx, ch)
+	if err != nil {
+		return -1, err
 	}
 	if res.Error != "" {
 		return -1, errors.New(res.Error)

sync/client_state.go

@@ -59,9 +59,9 @@ func (c *DefaultClient) Barrier(ctx context.Context, state State, target int) (*
 	}
 
 	go func() {
-		res, ok := <-ch
-		if !ok {
-			b.C <- errors.New("channel closed before getting response")
+		res, err := c.awaitResponse(ctx, ch)
+		if err != nil {
+			b.C <- err
 		} else if res.Error == "" {
 			b.C <- nil
 		} else {
@@ -99,9 +99,9 @@ func (c *DefaultClient) SignalEntry(ctx context.Context, state State) (int64, er
 		return -1, err
 	}
 
-	res, ok := <-ch
-	if !ok {
-		return -1, errors.New("channel closed before getting response")
+	res, err := c.awaitResponse(ctx, ch)
+	if err != nil {
+		return -1, err
 	}
 	if res.Error != "" {
 		return -1, errors.New(res.Error)