[management] shift network map data store onto the read replica (#8302)

management/internals/shared/requestbuffer/buffer.go

@@ -26,6 +26,9 @@ type Buffer[T any] struct {
 
 	mu      sync.Mutex
 	waiting map[string][]chan result[T]
+	// inFlight holds the keys with a running fetch and whether their next flush
+	// is due once it ends. It is nil unless fetches are serialized.
+	inFlight map[string]bool
 }
 
 type result[T any] struct {
@@ -45,6 +48,15 @@ func New[T any](ctx context.Context, name string, interval time.Duration, fetch
 	}
 }
 
+// NewSerialized returns a Buffer like New that runs at most one fetch per key at
+// a time. Requests arriving during a fetch form the next batch, which is fetched
+// once both its window closed and the running fetch ended.
+func NewSerialized[T any](ctx context.Context, name string, interval time.Duration, fetch FetchFunc[T]) *Buffer[T] {
+	b := New(ctx, name, interval, fetch)
+	b.inFlight = make(map[string]bool)
+	return b
+}
+
 // Get returns the value for key, sharing one fetch with the other callers of
 // the current batch. The value is shared as is, so callers must treat it as
 // read-only unless the fetch hands out copies.
@@ -71,8 +83,16 @@ func (b *Buffer[T]) Get(ctx context.Context, key string) (T, error) {
 
 func (b *Buffer[T]) flush(key string) {
 	b.mu.Lock()
+	if _, running := b.inFlight[key]; running {
+		b.inFlight[key] = true
+		b.mu.Unlock()
+		return
+	}
 	waiting := b.waiting[key]
 	delete(b.waiting, key)
+	if b.inFlight != nil && len(waiting) > 0 {
+		b.inFlight[key] = false
+	}
 	b.mu.Unlock()
 
 	if len(waiting) == 0 {
@@ -86,6 +106,17 @@ func (b *Buffer[T]) flush(key string) {
 	for _, ch := range waiting {
 		ch <- result[T]{value: value, err: err}
 	}
+
+	if b.inFlight == nil {
+		return
+	}
+	b.mu.Lock()
+	due := b.inFlight[key]
+	delete(b.inFlight, key)
+	b.mu.Unlock()
+	if due {
+		go b.flush(key)
+	}
 }
 
 // Interval reads a buffer interval from envVar, falling back to def.

management/internals/shared/requestbuffer/buffer_test.go

@@ -6,6 +6,7 @@ import (
 	"sync"
 	"sync/atomic"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/stretchr/testify/assert"
@@ -104,3 +105,73 @@ func TestBufferHonorsCallerContext(t *testing.T) {
 	_, err := buffer.Get(ctx, "account")
 	assert.ErrorIs(t, err, context.DeadlineExceeded)
 }
+
+func TestSerializedBufferQueuesRequestsDuringFetch(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		var fetches, running, maxRunning atomic.Int32
+		buffer := NewSerialized(context.Background(), "test", 10*time.Millisecond,
+			func(ctx context.Context, key string) (int32, error) {
+				n := running.Add(1)
+				defer running.Add(-1)
+				if n > maxRunning.Load() {
+					maxRunning.Store(n)
+				}
+				time.Sleep(time.Second)
+				return fetches.Add(1), nil
+			})
+
+		var wg sync.WaitGroup
+		get := func(want int32) {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				value, err := buffer.Get(context.Background(), "account")
+				assert.NoError(t, err)
+				assert.Equal(t, want, value)
+			}()
+		}
+
+		get(1)
+		time.Sleep(100 * time.Millisecond)
+		synctest.Wait()
+		require.Equal(t, int32(1), running.Load(), "first fetch should be running")
+
+		for range 5 {
+			get(2)
+		}
+		wg.Wait()
+
+		assert.Equal(t, int32(2), fetches.Load(), "requests queued during the fetch should share one fetch")
+		assert.Equal(t, int32(1), maxRunning.Load(), "fetches of a key should not overlap")
+	})
+}
+
+func TestUnserializedBufferOverlapsFetches(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		var running, maxRunning atomic.Int32
+		buffer := New(context.Background(), "test", 10*time.Millisecond,
+			func(ctx context.Context, key string) (string, error) {
+				n := running.Add(1)
+				defer running.Add(-1)
+				if n > maxRunning.Load() {
+					maxRunning.Store(n)
+				}
+				time.Sleep(time.Second)
+				return key, nil
+			})
+
+		var wg sync.WaitGroup
+		wg.Go(func() {
+			_, err := buffer.Get(context.Background(), "account")
+			assert.NoError(t, err)
+		})
+		time.Sleep(100 * time.Millisecond)
+		wg.Go(func() {
+			_, err := buffer.Get(context.Background(), "account")
+			assert.NoError(t, err)
+		})
+		wg.Wait()
+
+		assert.Equal(t, int32(2), maxRunning.Load())
+	})
+}