Data race on Request re-use

[mattias-p]

When I create a single Request and use it in a backoff.Retry callback, I get "DATA RACE" warnings. When I instead create the Request inside the backoff.Retry callback, everything works fine. The documentation states that Requests may be re-used. It looks to me like the current implementation is in violation of its documentation.

Software versions:

• github.com/0x4b53/amqp-rpc/v5 v5.0.1
• github.com/cenkalti/backoff/v5 v5.0.2 (not that I think it matters)

[bombsimon]

Thanks for the report!

Did a quick repro but didn't look much further for now.

func TestRequestReUse(t *testing.T) {
	url := "amqp://guest:guest@localhost:5672/"

	s := NewServer(url)
	s.Bind(DirectBinding("myqueue", func(_ context.Context, rw *ResponseWriter, d amqp.Delivery) {
		fmt.Fprintf(rw, "Got message: %s", d.Body)
	}))

	stop := startServerAndWait(s)
	defer stop()

	client := NewClient(url)
	defer client.Stop()

	request := NewRequest().
		WithRoutingKey("myqueue").
		WithResponse(true).
		WithBody("hello request")

	for range make([]struct{}, 2) {
		response, err := client.Send(request)
		require.NoError(t, err, "no errors sending request")
		assert.Equal(t, []byte("Got message: hello request"), response.Body, "correct body returned")
	}
}

[bombsimon]

I can fix this in a single thread by ensuring we're not accessing the context.

Diff

diff --git a/client.go b/client.go
index 6e71108..53a0e49 100644
--- a/client.go
+++ b/client.go
@@ -466,6 +466,8 @@ func (c *Client) runPublisher(ouputChan *amqp.Channel, wg *sync.WaitGroup) {
 			c.logger.Debug("publisher stopped after channel was closed")
 			return
 		case request := <-c.requests:
+			requestCtx := request.getContext()
+
 			// Set the ReplyTo if needed, or ensure it's empty if it's not.
 			if request.Reply {
 				request.Publishing.ReplyTo = c.replyToQueueName
@@ -473,7 +475,7 @@ func (c *Client) runPublisher(ouputChan *amqp.Channel, wg *sync.WaitGroup) {
 				request.Publishing.ReplyTo = ""
 			}
 
-			c.logger.DebugContext(request.Context, "publishing request",
+			c.logger.DebugContext(requestCtx, "publishing request",
 				slog.String("correlation_id", request.Publishing.CorrelationId),
 			)
 
@@ -501,7 +503,7 @@ func (c *Client) runPublisher(ouputChan *amqp.Channel, wg *sync.WaitGroup) {
 
 				c.retryRequest(request, err)
 
-				c.logger.ErrorContext(request.Context,
+				c.logger.ErrorContext(requestCtx,
 					"publisher stopped because of error",
 					slog.Any("error", err),
 					slogGroupFor("request", slogAttrsForRequest(request)),
@@ -523,9 +525,11 @@ func (c *Client) runPublisher(ouputChan *amqp.Channel, wg *sync.WaitGroup) {
 // has been retried too many times. Then the provided error will be sent to the
 // caller instead.
 func (c *Client) retryRequest(request *Request, err error) {
+	requestCtx := request.getContext()
+
 	if request.numRetries >= c.maxRetries {
 		// We have already retried too many times
-		c.logger.ErrorContext(request.Context,
+		c.logger.ErrorContext(requestCtx,
 			"could not publish, giving up",
 			slog.Any("error", err),
 			slogGroupFor("request", slogAttrsForRequest(request)),
@@ -540,16 +544,16 @@ func (c *Client) retryRequest(request *Request, err error) {
 	request.numRetries++
 
 	go func() {
-		c.logger.DebugContext(request.Context, "queuing request for retry",
+		c.logger.DebugContext(requestCtx, "queuing request for retry",
 			slog.Any("error", err),
 			slogGroupFor("request", slogAttrsForRequest(request)),
 		)
 
 		select {
 		case c.requests <- request:
-		case <-request.Context.Done():
-			c.logger.ErrorContext(request.Context, "canceled while waiting for retry",
-				slog.Any("error", errors.Join(context.Cause(request.Context), err)),
+		case <-requestCtx.Done():
+			c.logger.ErrorContext(requestCtx, "canceled while waiting for retry",
+				slog.Any("error", errors.Join(context.Cause(requestCtx), err)),
 				slogGroupFor("request", slogAttrsForRequest(request)),
 			)
 		}
@@ -590,7 +594,7 @@ func (c *Client) runConfirmsConsumer(confirms chan amqp.Confirmation, returns ch
 				continue
 			}
 
-			c.logger.DebugContext(request.Context, "publishing is returned by server",
+			c.logger.DebugContext(request.getContext(), "publishing is returned by server",
 				slog.String("correlation_id", ret.CorrelationId),
 			)
 
@@ -616,7 +620,7 @@ func (c *Client) runConfirmsConsumer(confirms chan amqp.Confirmation, returns ch
 				continue
 			}
 
-			c.logger.DebugContext(request.Context, "confirming request",
+			c.logger.DebugContext(request.getContext(), "confirming request",
 				slog.String("correlation_id", request.Publishing.CorrelationId),
 			)
 
@@ -650,13 +654,15 @@ func (c *Client) runConfirmsConsumer(confirms chan amqp.Confirmation, returns ch
 
 // respondToRequest will return the provided response to the caller.
 func (c *Client) respondToRequest(request *Request, d *amqp.Delivery, err error) {
+	requestCtx := request.getContext()
+
 	select {
 	case request.response <- response{delivery: d, err: err}:
 		return
-	case <-request.Context.Done():
-		c.logger.ErrorContext(request.Context,
+	case <-requestCtx.Done():
+		c.logger.ErrorContext(requestCtx,
 			"nobody is waiting for response",
-			slog.Any("error", errors.Join(context.Cause(request.Context), err)),
+			slog.Any("error", errors.Join(context.Cause(requestCtx), err)),
 			slogGroupFor("request", slogAttrsForRequest(request)),
 			slogGroupFor("delivery", slogAttrsForDelivery(d)),
 		)
@@ -719,7 +725,7 @@ func (c *Client) runRepliesConsumer(inChan *amqp.Channel, wg *sync.WaitGroup) (c
 				continue
 			}
 
-			c.logger.DebugContext(request.Context, "forwarding reply",
+			c.logger.DebugContext(request.getContext(), "forwarding reply",
 				slog.String("correlation_id", response.CorrelationId),
 			)
 
@@ -756,7 +762,12 @@ func (c *Client) send(r *Request) (*amqp.Delivery, error) {
 		r.Publishing.CorrelationId = uuid.New().String()
 	}
 
-	defer c.requestsMap.Delete(r)
+	requestCtx := r.getContext()
+
+	defer func() {
+		c.requestsMap.Delete(r)
+		r.WithContext(context.WithoutCancel(requestCtx))
+	}()
 
 	cancel := r.startTimeout(c.timeout)
 	defer cancel()
@@ -766,10 +777,10 @@ func (c *Client) send(r *Request) (*amqp.Delivery, error) {
 	select {
 	case c.requests <- r:
 		// successful send.
-	case <-r.Context.Done():
-		err := context.Cause(r.Context)
+	case <-requestCtx.Done():
+		err := context.Cause(requestCtx)
 
-		c.logger.DebugContext(r.Context,
+		c.logger.DebugContext(requestCtx,
 			"canceled while waiting for request queue",
 			slog.Any("error", err),
 			slog.String("correlation_id", r.Publishing.CorrelationId),
@@ -778,7 +789,7 @@ func (c *Client) send(r *Request) (*amqp.Delivery, error) {
 		return nil, fmt.Errorf("%w while waiting for request queue", err)
 	}
 
-	c.logger.DebugContext(r.Context,
+	c.logger.DebugContext(requestCtx,
 		"waiting for reply",
 		slog.String("correlation_id", r.Publishing.CorrelationId),
 	)
@@ -786,10 +797,10 @@ func (c *Client) send(r *Request) (*amqp.Delivery, error) {
 	select {
 	case res := <-r.response:
 		return res.delivery, res.err
-	case <-r.Context.Done():
-		err := context.Cause(r.Context)
+	case <-requestCtx.Done():
+		err := context.Cause(requestCtx)
 
-		c.logger.DebugContext(r.Context,
+		c.logger.DebugContext(requestCtx,
 			"canceled while waiting for response",
 			slog.Any("error", err),
 			slog.String("correlation_id", r.Publishing.CorrelationId),
diff --git a/request.go b/request.go
index 8b5ea8e..7968471 100644
--- a/request.go
+++ b/request.go
@@ -52,8 +52,8 @@ type Request struct {
 	numRetries int
 
 	deliveryTag uint64
-
-	returned *amqp.Return
+	returned    *amqp.Return
+	mu          sync.RWMutex
 }
 
 // NewRequest will generate a new request to be published. The default request
@@ -99,6 +99,9 @@ func (r *Request) WithCorrelationID(id string) *Request {
 
 // WithContext will set the context on the request.
 func (r *Request) WithContext(ctx context.Context) *Request {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
 	r.Context = ctx
 
 	return r
@@ -175,12 +178,22 @@ func (r *Request) AddMiddleware(m ClientMiddlewareFunc) *Request {
 	return r
 }
 
+func (r *Request) getContext() context.Context {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	return r.Context
+}
+
 // startTimeout will start the timeout counter. Is will also set the Expiration
 // field for the Publishing so that amqp won't hold on to the message in the
 // queue after the timeout has happened.
 func (r *Request) startTimeout(defaultTimeout time.Duration) context.CancelFunc {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
 	if r.Context == nil {
-		r.Context = context.Background()
+		r.WithContext(context.Background())
 	}
 
 	timeout := r.Timeout

But this still fails for multi threaded use. We read and write to the request all over the place so I'm curious if we should re-think this and put all read and writes of the request behind getters/setters/mutex to ensur safety. 🤔