Merge #156458

156458: rangefeed: drop messages for overflowed streams r=wenyihu6,arulajmani a=stevendanna

This PR fixes some bugs discovered when trying to increase the test coverage of the buffered sender unit tests:

- A stream would actually never move into the "overflowing" state. Rather it would stay in the "active" state and keep returning a buffer capacity error until the registration sent the error back to us, in which case it would go straight to overflowed.

- This was papering over another bug in which the overflowing state could produce an invalid stream of events.

Here, I add test coverage for the expected state transitions and fix these bugs by:

  - Dropping buffered messages for streams not in the stream map. This works because the stream will be in the stream map from the point of the RegisteringStream call and is only removed when the stream finally sends the error. Thus any message for a stream not in map represents either (1) a message sent to a StreamID before RegisteringStream has been called or (2) a message sent to a StreamID after an Error. (1) would be a programming error and (2) would be messages that are dropped server-side anyway.

  - Re-organizing the code such that it is harder to miss saving the stream status back to the status map.

Note any potential bug here only existed on master and thus no release note is needed.

Epic: CRDB-55217

Release note: None

Co-authored-by: Steven Danna <[email protected]>

Author: craig[bot]

pkg/kv/kvserver/rangefeed/BUILD.bazel

@@ -59,6 +59,7 @@ go_library(
         "//pkg/util/uuid",
         "@com_github_cockroachdb_crlib//crtime",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_redact//:redact",
     ],
 )
 

pkg/kv/kvserver/rangefeed/buffered_sender.go

@@ -8,17 +8,18 @@ package rangefeed
 import (
 	"context"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
-	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/redact"
 )
 
 //            ┌─────────────────┐
@@ -104,17 +105,45 @@ const (
 	// streamOverflowing is the state we are in when the stream has reached its
 	// limit and is waiting to deliver an error.
 	streamOverflowing streamState = iota
-	// streamOverflowed means the stream has overflowed and the error has been
-	// placed in the queue.
-	streamOverflowed streamState = iota
+	// streamErrored means an error has been enqueued for this stream and further
+	// buffered sends will be ignored. Streams in this state will not be found in
+	// the status map.
+	streamErrored streamState = iota
 )
 
+func (s streamState) String() string {
+	switch s {
+	case streamActive:
+		return "active"
+	case streamOverflowing:
+		return "overflowing"
+	case streamErrored:
+		return "errored"
+	default:
+		return "unknown"
+	}
+}
+
+func (s streamState) SafeFormat(w redact.SafePrinter, _ rune) {
+	w.Printf("%s", redact.SafeString(s.String()))
+}
+
 type streamStatus struct {
 	// queueItems is the number of items for a given stream in the event queue.
 	queueItems int64
 	state      streamState
 }
 
+func (s streamStatus) String() string {
+	return fmt.Sprintf("%s [queue_len:%d]", s.state, s.queueItems)
+}
+
+func (s streamStatus) SafeFormat(w redact.SafePrinter, _ rune) {
+	w.Printf("%s", redact.SafeString(s.String()))
+}
+
+var errNoSuchStream = errors.New("stream already encountered an error or has not be added to buffered sender")
+
 func NewBufferedSender(
 	sender ServerStreamSender, settings *cluster.Settings, bsMetrics *BufferedSenderMetrics,
 ) *BufferedSender {
@@ -130,8 +159,12 @@ func NewBufferedSender(
 }
 
 // sendBuffered buffers the event before sending it to the underlying gRPC
-// stream. It does not block. It errors in the case of a stopped sender of if
-// the registration has exceeded its capacity.
+// stream. It does not block.
+//
+// It errors in the case of a stopped sender of if the registration has exceeded
+// its capacity. sendBuffered with rangefeed events for streams that have
+// already enqueued an error event or have not been added via addStream will
+// return an error.
 func (bs *BufferedSender) sendBuffered(
 	ev *kvpb.MuxRangeFeedEvent, alloc *SharedBudgetAllocation,
 ) error {
@@ -145,55 +178,24 @@ func (bs *BufferedSender) sendBuffered(
 	// request. If the stream has hit its limit, we return an error to the
 	// registration. This error should be the next event that is sent to
 	// stream.
-	//
-	// NB: We don't error if the stream status is not found as this may be an
-	// event for an already closed stream. Such events are possible while the
-	// registration publishes the catch up scan buffer.
 	status, ok := bs.queueMu.byStream[ev.StreamID]
-	if ok {
-		switch status.state {
-		case streamActive:
-			if bs.queueMu.perStreamCapacity > 0 && status.queueItems == bs.queueMu.perStreamCapacity {
-				if ev.Error != nil {
-					// If _this_ event is an error, no use sending another error. This stream
-					// is going down. Admit this error and mark the stream as overflowed.
-					status.state = streamOverflowed
-				} else {
-					// This stream is at capacity, return an error to the registration that it
-					// should send back to us after cleaning up.
-					status.state = streamOverflowing
-					return newRetryErrBufferCapacityExceeded()
-				}
-			}
-		case streamOverflowing:
-			// The unbufferedRegistration is the only component that sends non-error
-			// events to our stream. In response to the error we return when moving to
-			// stateOverflowing, it should immediately send us an error and mark itself
-			// as disconnected.
-			//
-			// The only unfortunate exception is if we get disconnected while flushing
-			// the catch-up scan buffer. In this case we admit the event and stay in
-			// state overflowing until we actually receive the error.
-			//
-			// TODO(ssd): Given the above exception, we should perhaps just move
-			// directly to streamOverflowed. But, I think instead we want to remove
-			// that exception if possible.
-			if ev.Error != nil {
-				status.state = streamOverflowed
-			}
-		case streamOverflowed:
-			// If we are overflowed, we don't expect any further events because the
-			// registration should have disconnected in response to the error.
-			//
-			// TODO(ssd): Consider adding an assertion here.
-			return nil
-		default:
-			panic(fmt.Sprintf("unhandled stream state: %v", status.state))
+	if !ok {
+		return errNoSuchStream
+	}
+	nextState, err := bs.nextPerStreamStateLocked(status, ev)
+	if nextState == streamErrored {
+		// We will be admitting this event but no events after this.
+		assertTrue(err == nil, "expected error event to be admitted")
+		delete(bs.queueMu.byStream, ev.StreamID)
+	} else {
+		if err == nil {
+			status.queueItems++
 		}
-		// We are admitting this event.
-		status.queueItems++
+		status.state = nextState
 		bs.queueMu.byStream[ev.StreamID] = status
-
+	}
+	if err != nil {
+		return err
 	}
 
 	// TODO(wenyihu6): pass an actual context here
@@ -207,8 +209,51 @@ func (bs *BufferedSender) sendBuffered(
 	return nil
 }
 
+// nextPerStreamStateLocked returns the next state that should be stored for the
+// stream related to the given rangefeed event. If an error is returned, the
+// event should not be admitted and the given error should be returned to the
+// client.
+func (bs *BufferedSender) nextPerStreamStateLocked(
+	status streamStatus, ev *kvpb.MuxRangeFeedEvent,
+) (streamState, error) {
+	// An error should always put us into stateErrored, so let's do that first.
+	if ev.Error != nil {
+		if status.state == streamErrored {
+			assumedUnreachable("unexpected buffered event on stream in state streamErrored")
+		}
+		return streamErrored, nil
+	}
+
+	switch status.state {
+	case streamActive:
+		if bs.queueMu.perStreamCapacity > 0 && status.queueItems == bs.queueMu.perStreamCapacity {
+			// This stream is at capacity, return an error to the registration that it
+			// should send back to us after cleaning up.
+			return streamOverflowing, newRetryErrBufferCapacityExceeded()
+		}
+		// Happy path.
+		return streamActive, nil
+	case streamOverflowing:
+		// The only place we do concurrent buffered sends is during catch-up scan
+		// publishing which may be concurrent with a disconnect. The catch-up scan
+		// will stop publishing if it receives an error and try to send an error
+		// back. A disconnect only sends an error. This path exclusively handles
+		// non-errors.
+		assumedUnreachable("unexpected buffered event on stream in state streamOverflowing")
+		return streamOverflowing, newRetryErrBufferCapacityExceeded()
+	case streamErrored:
+		// This is unexpected because streamErrored streams are removed from the
+		// status map and thus should be handled in sendBuffered before this
+		// function is called.
+		assumedUnreachable("unexpected buffered event on stream in state streamErrored")
+		return streamErrored, errNoSuchStream
+	default:
+		panic(fmt.Sprintf("unhandled stream state: %v", status.state))
+	}
+}
+
 // sendUnbuffered sends the event directly to the underlying
-// ServerStreamSender.  It bypasses the buffer and thus may block.
+// ServerStreamSender. It bypasses the buffer and thus may block.
 func (bs *BufferedSender) sendUnbuffered(ev *kvpb.MuxRangeFeedEvent) error {
 	return bs.sender.Send(ev)
 }
@@ -273,24 +318,10 @@ func (bs *BufferedSender) addStream(streamID int64) {
 	if _, ok := bs.queueMu.byStream[streamID]; !ok {
 		bs.queueMu.byStream[streamID] = streamStatus{}
 	} else {
-		if buildutil.CrdbTestBuild {
-			panic(fmt.Sprintf("stream %d already exists in buffered sender", streamID))
-		}
+		assumedUnreachable(fmt.Sprintf("stream %d already exists in buffered sender", streamID))
 	}
 }
 
-// removeStream removes the per-stream state tracking from the sender.
-//
-// TODO(ssd): There may be items still in the queue when removeStream is called.
-// We'd like to solve this by removing this as a possibility. But this is OK
-// since we will eventually process the events and the client knows to ignore
-// them.
-func (bs *BufferedSender) removeStream(streamID int64) {
-	bs.queueMu.Lock()
-	defer bs.queueMu.Unlock()
-	delete(bs.queueMu.byStream, streamID)
-}
-
 // cleanup is called when the sender is stopped. It is expected to free up
 // buffer queue and no new events should be buffered after this.
 func (bs *BufferedSender) cleanup(ctx context.Context) {
@@ -309,6 +340,18 @@ func (bs *BufferedSender) len() int {
 	return int(bs.queueMu.buffer.len())
 }
 
+func (bs *BufferedSender) TestingBufferSummary() string {
+	bs.queueMu.Lock()
+	defer bs.queueMu.Unlock()
+
+	summary := &strings.Builder{}
+	fmt.Fprintf(summary, "buffered sender: queue_len=%d streams=%d", bs.queueMu.buffer.len(), len(bs.queueMu.byStream))
+	for id, stream := range bs.queueMu.byStream {
+		fmt.Fprintf(summary, "\n    %d: %s", id, stream)
+	}
+	return summary.String()
+}
+
 // Used for testing only.
 func (bs *BufferedSender) waitForEmptyBuffer(ctx context.Context) error {
 	opts := retry.Options{