rangefeed: explicitly initialize stream status

Since we do have a window after an error where the
unbufferedRegistration may send events, here we move to explicit
initialization of the stream status to avoid leaking any streamStatus
structs.

Epic: none
Release note: None

Author: stevendanna

pkg/kv/kvserver/rangefeed/buffered_sender.go

@@ -148,46 +148,48 @@ func (bs *BufferedSender) sendBuffered(
 	// registration. This error should be the next event that is sent to
 	// stream.
 	//
-	// NB: The zero-value of streamStatus is the valid state of a newly seen
-	// stream.
-	status := bs.queueMu.byStream[ev.StreamID]
-	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()
+	// 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. As a result, no non-error events are expected.
+			if ev.Error == nil {
+				panic("only error events expected after stream has exceeded capacity")
+			}
+			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))
 		}
-	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. As a result, no non-error events are expected.
-		if ev.Error == nil {
-			panic("only error events expected after stream has exceeded capacity")
-		}
-		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))
+		// We are admitting this event.
+		status.queueItems++
+		bs.queueMu.byStream[ev.StreamID] = status
 	}
 
-	// We are admitting this event.
-	status.queueItems++
-	bs.queueMu.byStream[ev.StreamID] = status
-
 	// TODO(wenyihu6): pass an actual context here
 	alloc.Use(context.Background())
 	bs.queueMu.buffer.pushBack(sharedMuxEvent{ev, alloc})
@@ -260,6 +262,14 @@ func (bs *BufferedSender) popFront() (e sharedMuxEvent, success bool) {
 	return event, ok
 }
 
+func (bs *BufferedSender) addStream(streamID int64) {
+	bs.queueMu.Lock()
+	defer bs.queueMu.Unlock()
+	if _, ok := bs.queueMu.byStream[streamID]; !ok {
+		bs.queueMu.byStream[streamID] = streamStatus{}
+	}
+}
+
 func (bs *BufferedSender) removeStream(streamID int64) {
 	bs.queueMu.Lock()
 	defer bs.queueMu.Unlock()

pkg/kv/kvserver/rangefeed/stream_manager.go

@@ -76,7 +76,9 @@ type sender interface {
 	// all streams in StreamManager.
 	run(ctx context.Context, stopper *stop.Stopper, onError func(int64)) error
 
-	// Remove stream is called when an individual stream is being removed.
+	// addStream is called when an individual stream is being added.
+	addStream(streamID int64)
+	// removeStream is called when an individual stream is being removed.
 	removeStream(streamID int64)
 
 	// cleanup is called when the sender is stopped. It is expected to clean up
@@ -156,6 +158,7 @@ func (sm *StreamManager) AddStream(streamID int64, d Disconnector) {
 	if _, ok := sm.streams.m[streamID]; ok {
 		log.KvDistribution.Fatalf(context.Background(), "stream %d already exists", streamID)
 	}
+	sm.sender.addStream(streamID)
 	sm.streams.m[streamID] = d
 	sm.metrics.ActiveMuxRangeFeed.Inc(1)
 	sm.metrics.NumMuxRangeFeed.Inc(1)

pkg/kv/kvserver/rangefeed/unbuffered_sender.go

@@ -132,6 +132,9 @@ func (ubs *UnbufferedSender) sendUnbuffered(event *kvpb.MuxRangeFeedEvent) error
 	return ubs.sender.Send(event)
 }
 
+// addStream implements sender.
+func (ubs *UnbufferedSender) addStream(int64) {}
+
 // removeStream implements sender.
 func (ubs *UnbufferedSender) removeStream(int64) {}