Skip to content

Commit 6470590

Browse files
wenyihu6wenyihu6
committed
kvserver/rangefeed: dynamically adjust capacity per buffered sender
Previously, buffered sender capacity was configured via a cluster setting. This commit removes that setting and makes capacity scale dynamically with the number of active registrations. The minimum (and default) is 4096 * 20, and capacity never drops below that. Above the minimum, it grows as 4096 * <active registrations>. Events already buffered will not dropped when capacity decreases, and capacity is not adjusted during shutdown (0).
1 parent e3f35e5 commit 6470590

File tree

8 files changed

+239
-44
lines changed

8 files changed

+239
-44
lines changed

‎pkg/kv/kvserver/rangefeed/buffered_sender.go‎

Lines changed: 36 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,6 @@ import (
1010
"time"
1111

1212
"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
13-
"github.com/cockroachdb/cockroach/pkg/util/log"
1413
"github.com/cockroachdb/cockroach/pkg/util/retry"
1514
"github.com/cockroachdb/cockroach/pkg/util/stop"
1615
"github.com/cockroachdb/cockroach/pkg/util/syncutil"
@@ -57,8 +56,23 @@ type BufferedSender struct {
5756
// queueMu protects the buffer queue.
5857
queueMu struct {
5958
syncutil.Mutex
60-
stopped bool
61-
buffer *eventQueue
59+
stopped bool
60+
buffer *eventQueue
61+
// capacity is the maximum number of events that can be buffered. It
62+
// dynamically scales based on the number of active registrations.
63+
//
64+
// The capacity is calculated as:
65+
// - Minimum: minBufferedSenderQueueCapacity (20 registrations)
66+
// - Active streams: perUnbufferedRegCapacity * number of active
67+
// registrations
68+
//
69+
// This scaling is based on the intuition that each stream should have
70+
// equivalent buffer space to what unbuffered registrations receive (4096
71+
// events per registration). Since all registrations share the same buffered
72+
// sender, registrations may affect each other.
73+
//
74+
// Note that when capacity shrinks, events already buffered will not be
75+
// dropped. Capacity is not adjusted during shutdown.
6276
capacity int64
6377
overflowed bool
6478
}
@@ -74,8 +88,14 @@ type BufferedSender struct {
7488
metrics *BufferedSenderMetrics
7589
}
7690

91+
// TODO(wenyihu6): This value is set to the same value as
92+
// https://github.com/cockroachdb/cockroach/blob/5536f0828f50bb21ec1577c77d388c4303d124a4/pkg/kv/kvserver/replica_rangefeed.go#L121.
93+
// Should I move this const to kvserverbase and import this value?
94+
const perUnbufferedRegCapacity = int64(4096)
95+
const minBufferedSenderQueueCapacity = int64(4096 * 20)
96+
7797
func NewBufferedSender(
78-
sender ServerStreamSender, bsMetrics *BufferedSenderMetrics, maxQueueSize int64,
98+
sender ServerStreamSender, bsMetrics *BufferedSenderMetrics,
7999
) *BufferedSender {
80100
bs := &BufferedSender{
81101
sender: sender,
@@ -84,7 +104,7 @@ func NewBufferedSender(
84104
bs.queueMu.buffer = newEventQueue()
85105
bs.notifyDataC = make(chan struct{}, 1)
86106
bs.queueMu.buffer = newEventQueue()
87-
bs.queueMu.capacity = maxQueueSize
107+
bs.queueMu.capacity = minBufferedSenderQueueCapacity
88108
return bs
89109
}
90110

@@ -101,8 +121,6 @@ func (bs *BufferedSender) sendBuffered(
101121
return errors.New("stream sender is stopped")
102122
}
103123
if bs.queueMu.overflowed {
104-
// Is this too spammy
105-
log.Dev.Error(context.Background(), "buffer capacity exceeded")
106124
return newRetryErrBufferCapacityExceeded()
107125
}
108126
if bs.queueMu.buffer.len() >= bs.queueMu.capacity {
@@ -190,7 +208,17 @@ func (bs *BufferedSender) cleanup(ctx context.Context) {
190208
bs.metrics.BufferedSenderQueueSize.Dec(remaining)
191209
}
192210

193-
// Used for testing only.
211+
// onStreamConnectOrDisconnect is called when a stream is added or removed. This
212+
// is currently used to dynamically adjust its queue capacity based on number of
213+
// active registrations. Note that this is not called during shutdown, so strict
214+
// dependency on this contract is discouraged. And note that we do not drop
215+
// events if the capacity is shrunk.
216+
func (bs *BufferedSender) onStreamConnectOrDisconnect(activeStreamCount int64) {
217+
bs.queueMu.Lock()
218+
defer bs.queueMu.Unlock()
219+
bs.queueMu.capacity = max(minBufferedSenderQueueCapacity, activeStreamCount*perUnbufferedRegCapacity)
220+
}
221+
194222
func (bs *BufferedSender) len() int {
195223
bs.queueMu.Lock()
196224
defer bs.queueMu.Unlock()

‎pkg/kv/kvserver/rangefeed/buffered_sender_test.go‎

Lines changed: 153 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313

1414
"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
1515
"github.com/cockroachdb/cockroach/pkg/roachpb"
16+
"github.com/cockroachdb/cockroach/pkg/testutils"
1617
"github.com/cockroachdb/cockroach/pkg/util/hlc"
1718
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
1819
"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -34,7 +35,7 @@ func TestBufferedSenderDisconnectStream(t *testing.T) {
3435
defer stopper.Stop(ctx)
3536
testServerStream := newTestServerStream()
3637
smMetrics := NewStreamManagerMetrics()
37-
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
38+
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
3839
sm := NewStreamManager(bs, smMetrics)
3940
require.NoError(t, sm.Start(ctx, stopper))
4041
defer sm.Stop(ctx)
@@ -87,7 +88,7 @@ func TestBufferedSenderChaosWithStop(t *testing.T) {
8788
testServerStream := newTestServerStream()
8889

8990
smMetrics := NewStreamManagerMetrics()
90-
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
91+
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
9192
sm := NewStreamManager(bs, smMetrics)
9293
require.NoError(t, sm.Start(ctx, stopper))
9394

@@ -164,3 +165,153 @@ func TestBufferedSenderChaosWithStop(t *testing.T) {
164165
require.Equal(t, 0, bs.len())
165166
})
166167
}
168+
169+
// TestBufferedSenderOnOverflow tests that BufferedSender handles overflow
170+
// properly. It does not test the shutdown flow with stream manager.
171+
func TestBufferedSenderOnOverflow(t *testing.T) {
172+
defer leaktest.AfterTest(t)()
173+
defer log.Scope(t).Close(t)
174+
ctx := context.Background()
175+
176+
stopper := stop.NewStopper()
177+
defer stopper.Stop(ctx)
178+
testServerStream := newTestServerStream()
179+
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
180+
require.Equal(t, minBufferedSenderQueueCapacity, bs.queueMu.capacity)
181+
182+
val1 := roachpb.Value{RawBytes: []byte("val"), Timestamp: hlc.Timestamp{WallTime: 1}}
183+
ev1 := new(kvpb.RangeFeedEvent)
184+
ev1.MustSetValue(&kvpb.RangeFeedValue{Key: keyA, Value: val1})
185+
muxEv := &kvpb.MuxRangeFeedEvent{RangeFeedEvent: *ev1, RangeID: 0, StreamID: 1}
186+
187+
for i := int64(0); i < minBufferedSenderQueueCapacity; i++ {
188+
require.NoError(t, bs.sendBuffered(muxEv, nil))
189+
}
190+
require.Equal(t, minBufferedSenderQueueCapacity, int64(bs.len()))
191+
e, success, overflowed, remains := bs.popFront()
192+
require.Equal(t, sharedMuxEvent{
193+
ev: muxEv,
194+
alloc: nil,
195+
}, e)
196+
require.True(t, success)
197+
require.False(t, overflowed)
198+
require.Equal(t, minBufferedSenderQueueCapacity-1, remains)
199+
require.Equal(t, minBufferedSenderQueueCapacity-1, int64(bs.len()))
200+
require.NoError(t, bs.sendBuffered(muxEv, nil))
201+
require.Equal(t, minBufferedSenderQueueCapacity, int64(bs.len()))
202+
203+
// Overflow now.
204+
require.Equal(t, bs.sendBuffered(muxEv, nil).Error(),
205+
newRetryErrBufferCapacityExceeded().Error())
206+
}
207+
208+
// TestBufferedSenderOnStreamShutdown tests that BufferedSender and
209+
// StreamManager handle overflow and shutdown properly.
210+
211+
func TestBufferedSenderOnStreamShutdown(t *testing.T) {
212+
defer leaktest.AfterTest(t)()
213+
defer log.Scope(t).Close(t)
214+
ctx := context.Background()
215+
216+
stopper := stop.NewStopper()
217+
defer stopper.Stop(ctx)
218+
testServerStream := newTestServerStream()
219+
smMetrics := NewStreamManagerMetrics()
220+
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
221+
require.Equal(t, minBufferedSenderQueueCapacity, bs.queueMu.capacity)
222+
sm := NewStreamManager(bs, smMetrics)
223+
require.NoError(t, sm.Start(ctx, stopper))
224+
defer sm.Stop(ctx)
225+
226+
p, h, pStopper := newTestProcessor(t, withRangefeedTestType(scheduledProcessorWithBufferedSender))
227+
defer pStopper.Stop(ctx)
228+
229+
val1 := roachpb.Value{RawBytes: []byte("val"), Timestamp: hlc.Timestamp{WallTime: 1}}
230+
ev1 := new(kvpb.RangeFeedEvent)
231+
ev1.MustSetValue(&kvpb.RangeFeedValue{Key: keyA, Value: val1})
232+
muxEv := &kvpb.MuxRangeFeedEvent{RangeFeedEvent: *ev1, RangeID: 0, StreamID: 1}
233+
234+
// Add 21 streams and overflow the buffer.
235+
t.Run("add 21 streams", func(t *testing.T) {
236+
numStreams := int64(21)
237+
expectedCapacity := perUnbufferedRegCapacity * numStreams
238+
// Block the stream to help the queue to overflow.
239+
unblock := testServerStream.BlockSend()
240+
for id := int64(0); id < numStreams; id++ {
241+
registered, d, _ := p.Register(ctx, h.span, hlc.Timestamp{}, nil, /* catchUpIter */
242+
false /* withDiff */, false /* withFiltering */, false /* withOmitRemote */, noBulkDelivery,
243+
sm.NewStream(id, 1 /*rangeID*/))
244+
require.True(t, registered)
245+
sm.AddStream(id, d)
246+
}
247+
require.Equal(t, expectedCapacity, bs.queueMu.capacity)
248+
for int64(bs.len()) != expectedCapacity {
249+
require.NoError(t, sm.sender.sendBuffered(muxEv, nil))
250+
}
251+
require.Equal(t, expectedCapacity, int64(bs.len()))
252+
require.Equal(t, bs.sendBuffered(muxEv, nil).Error(),
253+
newRetryErrBufferCapacityExceeded().Error())
254+
require.True(t, bs.queueMu.overflowed)
255+
unblock()
256+
})
257+
258+
t.Run("overflow clean up", func(t *testing.T) {
259+
// All events buffered should still be sent to the stream.
260+
testutils.SucceedsSoon(t, func() error {
261+
if bs.len() == 0 {
262+
return nil
263+
}
264+
return errors.Newf("expected 0 registrations, found %d", bs.len())
265+
})
266+
// Overflow cleanup.
267+
err := <-sm.Error()
268+
require.Equal(t, newRetryErrBufferCapacityExceeded().Error(), err.Error())
269+
// Note that we expect the stream manager to shut down here, but no actual
270+
// error events would be sent during the shutdown.
271+
require.Equal(t, bs.sendBuffered(muxEv, nil).Error(), newRetryErrBufferCapacityExceeded().Error())
272+
})
273+
}
274+
275+
// TestBufferedSenderOnStreamDisconnect tests that BufferedSender dynamically
276+
// adjusts its capacity when streams are connected or disconnected.
277+
func TestBufferedSenderOnStreamDisconnect(t *testing.T) {
278+
defer leaktest.AfterTest(t)()
279+
defer log.Scope(t).Close(t)
280+
ctx := context.Background()
281+
282+
stopper := stop.NewStopper()
283+
defer stopper.Stop(ctx)
284+
testServerStream := newTestServerStream()
285+
smMetrics := NewStreamManagerMetrics()
286+
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
287+
require.Equal(t, minBufferedSenderQueueCapacity, bs.queueMu.capacity)
288+
289+
sm := NewStreamManager(bs, smMetrics)
290+
require.NoError(t, sm.Start(ctx, stopper))
291+
292+
p, h, pStopper := newTestProcessor(t, withRangefeedTestType(scheduledProcessorWithBufferedSender))
293+
defer pStopper.Stop(ctx)
294+
defer sm.Stop(ctx)
295+
296+
numStreams := int64(21)
297+
expectedCapacity := perUnbufferedRegCapacity * numStreams
298+
for id := int64(0); id < numStreams; id++ {
299+
registered, d, _ := p.Register(ctx, h.span, hlc.Timestamp{}, nil, /* catchUpIter */
300+
false /* withDiff */, false /* withFiltering */, false /* withOmitRemote */, noBulkDelivery,
301+
sm.NewStream(id, 1 /*rangeID*/))
302+
require.True(t, registered)
303+
sm.AddStream(id, d)
304+
}
305+
require.Equal(t, expectedCapacity, bs.queueMu.capacity)
306+
sm.DisconnectStream(int64(0), newErrBufferCapacityExceeded())
307+
testServerStream.waitForEvent(t, makeMuxRangefeedErrorEvent(0, 1, newErrBufferCapacityExceeded()))
308+
testutils.SucceedsSoon(t, func() error {
309+
bs.queueMu.Lock()
310+
defer bs.queueMu.Unlock()
311+
if bs.queueMu.capacity == expectedCapacity-perUnbufferedRegCapacity {
312+
return nil
313+
}
314+
return errors.Newf("expected %d cap to be %d", bs.queueMu.capacity,
315+
expectedCapacity-perUnbufferedRegCapacity)
316+
})
317+
}

‎pkg/kv/kvserver/rangefeed/stream_manager.go‎

Lines changed: 41 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,12 @@ type sender interface {
7878
// cleanup is called when the sender is stopped. It is expected to clean up
7979
// any resources used by the sender.
8080
cleanup(ctx context.Context)
81+
// onStreamConnectOrDisconnect is called when a stream is added or removed.
82+
// This is currently used for buffered sender to dynamically adjust its queue
83+
// capacity based on number of active registrations. Note that this is not
84+
// called during shutdown, so strict dependency on this contract is
85+
// discouraged.
86+
onStreamConnectOrDisconnect(activeStreamCount int64)
8187
}
8288

8389
func NewStreamManager(sender sender, metrics *StreamManagerMetrics) *StreamManager {
@@ -108,12 +114,18 @@ func (sm *StreamManager) NewStream(streamID int64, rangeID roachpb.RangeID) (sin
108114
// streamID to avoid metrics inaccuracy when the error is sent before the stream
109115
// is added to the StreamManager.
110116
func (sm *StreamManager) OnError(streamID int64) {
111-
sm.streams.Lock()
112-
defer sm.streams.Unlock()
113-
if _, ok := sm.streams.m[streamID]; ok {
114-
delete(sm.streams.m, streamID)
115-
sm.metrics.ActiveMuxRangeFeed.Dec(1)
116-
}
117+
func() {
118+
sm.streams.Lock()
119+
defer sm.streams.Unlock()
120+
if _, ok := sm.streams.m[streamID]; ok {
121+
delete(sm.streams.m, streamID)
122+
sm.metrics.ActiveMuxRangeFeed.Dec(1)
123+
}
124+
}()
125+
126+
// Call onStreamConnectOrDisconnect regardless of whether ActiveMuxRangeFeed
127+
// has changed for simplicity.
128+
sm.sender.onStreamConnectOrDisconnect(sm.metrics.ActiveMuxRangeFeed.Value())
117129
}
118130

119131
// DisconnectStream disconnects the stream with the given streamID.
@@ -138,19 +150,27 @@ func (sm *StreamManager) AddStream(streamID int64, d Disconnector) {
138150
// At this point, the stream had been registered with the processor and
139151
// started receiving events. We need to lock here to avoid race conditions
140152
// with a disconnect error passing through before the stream is added.
141-
sm.streams.Lock()
142-
defer sm.streams.Unlock()
143-
if d.IsDisconnected() {
144-
// If the stream is already disconnected, we don't add it to streams. The
145-
// registration will have already sent an error to the client.
146-
return
147-
}
148-
if _, ok := sm.streams.m[streamID]; ok {
149-
log.KvDistribution.Fatalf(context.Background(), "stream %d already exists", streamID)
150-
}
151-
sm.streams.m[streamID] = d
152-
sm.metrics.ActiveMuxRangeFeed.Inc(1)
153-
sm.metrics.NumMuxRangeFeed.Inc(1)
153+
func() {
154+
sm.streams.Lock()
155+
defer sm.streams.Unlock()
156+
if d.IsDisconnected() {
157+
// If the stream is already disconnected, we don't add it to streams. The
158+
// registration will have already sent an error to the client.
159+
return
160+
}
161+
if _, ok := sm.streams.m[streamID]; ok {
162+
log.KvDistribution.Fatalf(context.Background(), "stream %d already exists", streamID)
163+
}
164+
sm.streams.m[streamID] = d
165+
sm.metrics.ActiveMuxRangeFeed.Inc(1)
166+
sm.metrics.NumMuxRangeFeed.Inc(1)
167+
}()
168+
169+
// TODO(during review) Is it okay to trust the metrics gauge value here? We
170+
// can maintain our own counter here as well
171+
// Call onStreamConnectOrDisconnect regardless of whether ActiveMuxRangeFeed
172+
// has changed for simplicity.
173+
sm.sender.onStreamConnectOrDisconnect(sm.metrics.ActiveMuxRangeFeed.Value())
154174
}
155175

156176
// Start launches sender.run in the background if no error is returned.
@@ -186,6 +206,8 @@ func (sm *StreamManager) Start(ctx context.Context, stopper *stop.Stopper) error
186206
func (sm *StreamManager) Stop(ctx context.Context) {
187207
sm.taskCancel()
188208
sm.wg.Wait()
209+
// Since this is called during shutdown, sm.sender.onStreamConnectOrDisconnect
210+
// is not being called explicitly to update the queue capacity.
189211
sm.sender.cleanup(ctx)
190212
sm.streams.Lock()
191213
defer sm.streams.Unlock()

‎pkg/kv/kvserver/rangefeed/stream_manager_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ func TestStreamManagerErrorHandling(t *testing.T) {
180180
case scheduledProcessorWithUnbufferedSender:
181181
s = NewUnbufferedSender(testServerStream)
182182
case scheduledProcessorWithBufferedSender:
183-
s = NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
183+
s = NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
184184
default:
185185
t.Fatalf("unknown rangefeed test type %v", rt)
186186
}

‎pkg/kv/kvserver/rangefeed/unbuffered_registration_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ func TestUnbufferedRegWithStreamManager(t *testing.T) {
3434
defer stopper.Stop(ctx)
3535
testServerStream := newTestServerStream()
3636
smMetrics := NewStreamManagerMetrics()
37-
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
37+
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
3838
sm := NewStreamManager(bs, smMetrics)
3939
require.NoError(t, sm.Start(ctx, stopper))
4040

@@ -106,7 +106,7 @@ func TestUnbufferedRegCorrectnessOnDisconnect(t *testing.T) {
106106
defer stopper.Stop(ctx)
107107
testServerStream := newTestServerStream()
108108
smMetrics := NewStreamManagerMetrics()
109-
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics(), 1000)
109+
bs := NewBufferedSender(testServerStream, NewBufferedSenderMetrics())
110110
sm := NewStreamManager(bs, smMetrics)
111111
require.NoError(t, sm.Start(ctx, stopper))
112112
defer sm.Stop(ctx)

‎pkg/kv/kvserver/rangefeed/unbuffered_sender.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -186,4 +186,7 @@ func (ubs *UnbufferedSender) detachMuxErrors() []*kvpb.MuxRangeFeedEvent {
186186
return toSend
187187
}
188188

189-
func (ubs *UnbufferedSender) cleanup(context.Context) {}
189+
// The following methods are no-op implementations to satisfy the sender
190+
// interface.
191+
func (ubs *UnbufferedSender) cleanup(context.Context) {}
192+
func (ubs *UnbufferedSender) onStreamConnectOrDisconnect(_ int64) { return }

0 commit comments

Comments
 (0)