fix: keep the replay buffer when a request body reader is not read (#10805)

A resiliency attempt that hits its timeout is not waited for, so it can
call RawData() after the next attempt has read the body. Each RawData()
call reset the replay buffer, so if that reader was never read the body
was lost and every later retry sent an empty body.

Keep everything read so far in the replay buffer and give each reader a
snapshot of it followed by the rest of the stream.

Signed-off-by: Javier Aliaga <javier@diagrid.io>
Co-authored-by: Mike Nguyen <hey@mike.ee>
Co-authored-by: Dapr Bot <56698301+dapr-bot@users.noreply.github.com>

pkg/messaging/direct_messaging_replay_test.go

@@ -0,0 +1,133 @@
+/*
+Copyright 2026 The Dapr Authors
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package messaging
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+
+	invokev1 "github.com/dapr/dapr/pkg/messaging/v1"
+	internalv1pb "github.com/dapr/dapr/pkg/proto/internals/v1"
+)
+
+type staleFakeClient struct {
+	internalv1pb.ServiceInvocationClient
+	// gate, when non-nil, is waited on by the first CallLocalStream call.
+	gate  chan struct{}
+	mu    sync.Mutex
+	calls int
+	sent  map[int]*bytes.Buffer
+}
+
+func (c *staleFakeClient) CallLocalStream(ctx context.Context, _ ...grpc.CallOption) (internalv1pb.ServiceInvocation_CallLocalStreamClient, error) {
+	c.mu.Lock()
+	c.calls++
+	n := c.calls
+	buf := &bytes.Buffer{}
+	c.sent[n] = buf
+	c.mu.Unlock()
+	if n == 1 && c.gate != nil {
+		<-c.gate
+	}
+	return &staleFakeStream{buf: buf, mu: &c.mu}, nil
+}
+
+type staleFakeStream struct {
+	grpc.ClientStream
+	buf  *bytes.Buffer
+	mu   *sync.Mutex
+	recv int
+}
+
+func (s *staleFakeStream) Send(m *internalv1pb.InternalInvokeRequestStream) error {
+	return s.SendMsg(m)
+}
+
+func (s *staleFakeStream) SendMsg(m any) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if p := m.(*internalv1pb.InternalInvokeRequestStream).GetPayload(); p != nil {
+		s.buf.Write(p.GetData())
+	}
+	return nil
+}
+func (s *staleFakeStream) CloseSend() error { return nil }
+func (s *staleFakeStream) Recv() (*internalv1pb.InternalInvokeResponseStream, error) {
+	return nil, io.EOF
+}
+
+func (s *staleFakeStream) RecvMsg(m any) error {
+	s.recv++
+	if s.recv > 1 {
+		return io.EOF
+	}
+	m.(*internalv1pb.InternalInvokeResponseStream).Response = &internalv1pb.InternalInvokeResponse{
+		Status: &internalv1pb.Status{Code: 0},
+	}
+	return nil
+}
+
+// A resiliency attempt that timed out is not waited for. When it reaches RawData() after the next attempt sent the
+// body, and returns without reading because its context is canceled, the attempt after that must still send the body.
+func TestInvokeRemoteStreamStaleAttemptKeepsBody(t *testing.T) {
+	body := strings.Repeat("x", 61)
+	fc := &staleFakeClient{gate: make(chan struct{}), sent: map[int]*bytes.Buffer{}}
+	d := &directMessaging{}
+
+	req := invokev1.NewInvokeMethodRequest("method").
+		WithRawData(io.NopCloser(strings.NewReader(body))).
+		WithReplay(true)
+	defer req.Close()
+
+	noop := func(bool) {}
+
+	// Attempt 1: blocks inside CallLocalStream, before req.RawData().
+	ctx1, cancel1 := context.WithCancel(t.Context())
+	done1 := make(chan struct{})
+	go func() {
+		defer close(done1)
+		_, _, _ = d.invokeRemoteStream(ctx1, fc, req, "app", nil, noop)
+	}()
+	// Wait until attempt 1 is inside CallLocalStream.
+	require.Eventually(t, func() bool { fc.mu.Lock(); defer fc.mu.Unlock(); return fc.calls == 1 }, 5e9, 1e6)
+
+	// Timeout fires for attempt 1; attempt 2 runs to completion.
+	cancel1()
+	resp, _, err := d.invokeRemoteStream(t.Context(), fc, req, "app", nil, noop)
+	require.NoError(t, err)
+	_ = resp.Close()
+	fc.mu.Lock()
+	require.Equal(t, body, fc.sent[2].String(), "attempt 2 should send the body")
+	fc.mu.Unlock()
+
+	// Stale attempt 1 now proceeds.
+	close(fc.gate)
+	<-done1
+
+	// Attempt 3.
+	resp, _, err = d.invokeRemoteStream(t.Context(), fc, req, "app", nil, noop)
+	require.NoError(t, err)
+	_ = resp.Close()
+	fc.mu.Lock()
+	defer fc.mu.Unlock()
+	assert.Equal(t, body, fc.sent[3].String(), "attempt 3 should send the body")
+}

pkg/messaging/v1/replayable_request.go

@@ -18,19 +18,15 @@ import (
 	"io"
 	"sync"
 
-	"github.com/dapr/kit/byteslicepool"
 	streamutils "github.com/dapr/kit/streams"
 )
 
 // Minimum capacity for the slices is 2KB
 const minByteSliceCapacity = 2 << 10
 
-// Contain pools of *bytes.Buffer and []byte objects.
+// Contains a pool of *bytes.Buffer objects.
 // Used to reduce the number of allocations in replayableRequest for buffers and relieve pressure on the GC.
-var (
-	bufPool = sync.Pool{New: newBuffer}
-	bsPool  = byteslicepool.NewByteSlicePool(minByteSliceCapacity)
-)
+var bufPool = sync.Pool{New: newBuffer}
 
 func newBuffer() any {
 	return new(bytes.Buffer)
@@ -42,7 +38,6 @@ type replayableRequest struct {
 	replay           *bytes.Buffer
 	lock             sync.Mutex
 	currentTeeReader *streamutils.TeeReadCloser
-	currentData      []byte
 }
 
 // WithRawData sets message data.
@@ -91,28 +86,14 @@ func (rr *replayableRequest) RawData() (r io.Reader) {
 		// If there's no data, and there's never been, just return a reader with no data
 		r = bytes.NewReader(nil)
 	} else if rr.replay != nil {
-		// If there's replaying enabled, we need to create a new TeeReadCloser
-		// We need to copy the data read insofar from the reply buffer because the buffer becomes invalid after new data is written into the it, then reset the buffer
+		// Replay is enabled: everything read so far stays in the replay buffer (it is never reset),
+		// so a RawData() call that is never read from cannot lose data.
+		// New readers get a snapshot of the bytes read so far, followed by the rest of the stream,
+		// which is teed into the replay buffer.
 		l := rr.replay.Len()
-
-		// Get a new byte slice from the pool if we don't have one, and ensure it has enough capacity
-		if rr.currentData == nil {
-			rr.currentData = bsPool.Get(l)
-		}
-		rr.currentData = bsPool.Resize(rr.currentData, l)
-
-		// Copy the data from the replay buffer into the byte slice
-		copy(rr.currentData[0:l], rr.replay.Bytes())
-		rr.replay.Reset()
-
-		// Create a new TeeReadCloser that reads from the previously-read data and then the data not yet processed
-		// The TeeReadCloser also keeps all the data it reads into the replay buffer
-		mr := streamutils.NewMultiReaderCloser(
-			bytes.NewReader(rr.currentData[0:l]),
-			rr.data,
-		)
-		rr.currentTeeReader = streamutils.NewTeeReadCloser(mr, rr.replay)
-		r = rr.currentTeeReader
+		snapshot := rr.replay.Bytes()[:l:l]
+		rr.currentTeeReader = streamutils.NewTeeReadCloser(rr.data, rr.replay)
+		r = streamutils.NewMultiReaderCloser(bytes.NewReader(snapshot), rr.currentTeeReader)
 	} else {
 		// No replay enabled
 		r = rr.data
@@ -127,10 +108,6 @@ func (rr *replayableRequest) closeReplay() {
 		bufPool.Put(rr.replay)
 		rr.replay = nil
 	}
-	if rr.currentData != nil {
-		bsPool.Put(rr.currentData)
-		rr.currentData = nil
-	}
 }
 
 // Close the data stream and replay buffers.

pkg/messaging/v1/replayable_request_test.go

@@ -216,6 +216,27 @@ func TestReplayableRequest(t *testing.T) {
 			assert.Nil(t, rr.replay)
 		})
 	})
+
+	// A resiliency attempt that timed out keeps running and can call RawData() after the next attempt has read the body,
+	// then return without reading because its context is canceled.
+	t.Run("read in full, then RawData not read", func(t *testing.T) {
+		rr := newReplayable()
+		defer rr.Close()
+
+		t.Run("first read in full", func(t *testing.T) {
+			read, err := io.ReadAll(rr.RawData())
+			require.NoError(t, err)
+			assert.Equal(t, message, read)
+		})
+
+		_ = rr.RawData()
+
+		t.Run("third read in full", func(t *testing.T) {
+			read, err := io.ReadAll(rr.RawData())
+			require.NoError(t, err)
+			assert.Equal(t, message, read)
+		})
+	})
 }
 
 // readerCloser is a io.Reader that can be closed. Once the stream is closed, reading from it returns an error.