add request buffer for nmap data store

management/internals/controllers/network_map/controller/controller.go

@@ -21,6 +21,7 @@ import (
 	networkmapdb "github.com/netbirdio/netbird/management/internals/network_map_db"
 	"github.com/netbirdio/netbird/management/internals/server/config"
 	"github.com/netbirdio/netbird/management/internals/shared/grpc"
+	"github.com/netbirdio/netbird/management/internals/shared/requestbuffer"
 	"github.com/netbirdio/netbird/management/server/account"
 	"github.com/netbirdio/netbird/management/server/integrations/integrated_validator"
 	"github.com/netbirdio/netbird/management/server/integrations/port_forwarding"
@@ -39,6 +40,8 @@ import (
 	"github.com/netbirdio/netbird/version"
 )
 
+const defaultNetworkMapDataBufferInterval = 100 * time.Millisecond
+
 type Controller struct {
 	repo    Repository
 	metrics *metrics
@@ -65,7 +68,8 @@ type Controller struct {
 
 	perAccountServerSupportedSyncMessageVersions map[string]sharedgrpc.SyncMessageVersion
 
-	nmdataStore *networkmapdb.NetworkMapDBStoreImpl
+	nmdataStore  *networkmapdb.NetworkMapDBStoreImpl
+	nmdataBuffer *requestbuffer.Buffer[*networkmap.NetworkMapData]
 }
 
 type bufferUpdate struct {
@@ -89,7 +93,7 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
 		log.Fatal(fmt.Errorf("error creating metrics: %w", err))
 	}
 
-	return &Controller{
+	c := &Controller{
 		repo:                    newRepository(store),
 		metrics:                 nMetrics,
 		accountManagerMetrics:   metrics.AccountManagerMetrics(),
@@ -106,6 +110,14 @@ func NewController(ctx context.Context, store store.Store, metrics telemetry.App
 		perAccountServerSupportedSyncMessageVersions: sharedgrpc.SyncMessageVersionsFromMap(config.PerAccountHighestSupportedSyncMessageVersion),
 		nmdataStore:                                  nmdataStore,
 	}
+
+	if nmdataStore != nil {
+		interval := requestbuffer.Interval(ctx, "NB_NETWORK_MAP_DATA_BUFFER_INTERVAL", defaultNetworkMapDataBufferInterval)
+		log.WithContext(ctx).Infof("set network map data request buffer interval to %s", interval)
+		c.nmdataBuffer = requestbuffer.New(ctx, "network map data request buffer", interval, c.fetchNetworkMapData)
+	}
+
+	return c
 }
 
 func (c *Controller) OnPeerConnected(ctx context.Context, accountID string, peerID string) (chan *network_map.UpdateMessage, error) {
@@ -392,8 +404,6 @@ func (c *Controller) sendUpdatesFromData(ctx context.Context, accountID string,
 		return fmt.Errorf("failed to get flow enabled status: %v", err)
 	}
 
-	nmData.PrecomputePostureValidation()
-
 	dnsCache := &cache.DNSConfigCache{}
 	dnsDomain := c.getDNSDomainFromData(nmData.AccountSettings)
 	peersCustomZone := networkmap.PeersCustomZone(ctx, accountID, dnsDomain, nmData.Peers, IPv6AllowedPeersFromData(nmData))
@@ -476,21 +486,37 @@ func (c *Controller) sendUpdatesFromData(ctx context.Context, accountID string,
 }
 
 func (c *Controller) getNetworkMapData(ctx context.Context, accountID string) *networkmap.NetworkMapData {
-	if c.nmdataStore == nil {
+	if c.nmdataBuffer == nil {
 		return nil
 	}
 
-	nmData, err := c.nmdataStore.GetNetworkMapData(ctx, accountID)
+	nmData, err := c.nmdataBuffer.Get(ctx, accountID)
 	if err != nil {
 		log.WithContext(ctx).Errorf("failed to get network map data for account %s, falling back to account-based computation: %v", accountID, err)
 		return nil
 	}
 
-	nmData.Services = c.proxyServicesFromRepo(ctx, accountID)
-
 	return nmData
 }
 
+// fetchNetworkMapData reads the twin once per buffer window. Its result is
+// shared by every waiter of that window, so the mutating steps run here, before
+// it is handed out: the twin the callers see is read-only. Injected proxy
+// policies carry no posture checks, so precomputing after the injection yields
+// the same validation as precomputing before it.
+func (c *Controller) fetchNetworkMapData(ctx context.Context, accountID string) (*networkmap.NetworkMapData, error) {
+	nmData, err := c.nmdataStore.GetNetworkMapData(ctx, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	nmData.Services = c.proxyServicesFromRepo(ctx, accountID)
+	nmData.InjectProxyPolicies()
+	nmData.PrecomputePostureValidation()
+
+	return nmData, nil
+}
+
 func (c *Controller) getDNSDomainFromData(settings *nmdata.AccountSettingsInfo) string {
 	if settings == nil || settings.DNSDomain == "" {
 		return c.dnsDomain

management/internals/shared/requestbuffer/buffer.go

@@ -0,0 +1,102 @@
+// Package requestbuffer coalesces concurrent reads of the same expensive
+// resource into a single fetch.
+package requestbuffer
+
+import (
+	"context"
+	"os"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// FetchFunc reads the resource identified by key.
+type FetchFunc[T any] func(ctx context.Context, key string) (T, error)
+
+// Buffer batches requests per key: the first request opens a window, every
+// request arriving within it joins the batch, and a single fetch serves them
+// all. The fetch starts only after the window closed, so a caller never
+// observes data read before its own request.
+type Buffer[T any] struct {
+	ctx      context.Context
+	name     string
+	fetch    FetchFunc[T]
+	interval time.Duration
+
+	mu      sync.Mutex
+	waiting map[string][]chan result[T]
+}
+
+type result[T any] struct {
+	value T
+	err   error
+}
+
+// New returns a Buffer serving batched requests through fetch. ctx bounds the
+// fetches, not the callers, and must outlive them.
+func New[T any](ctx context.Context, name string, interval time.Duration, fetch FetchFunc[T]) *Buffer[T] {
+	return &Buffer[T]{
+		ctx:      ctx,
+		name:     name,
+		fetch:    fetch,
+		interval: interval,
+		waiting:  make(map[string][]chan result[T]),
+	}
+}
+
+// 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.
+func (b *Buffer[T]) Get(ctx context.Context, key string) (T, error) {
+	ch := make(chan result[T], 1)
+
+	b.mu.Lock()
+	b.waiting[key] = append(b.waiting[key], ch)
+	first := len(b.waiting[key]) == 1
+	b.mu.Unlock()
+
+	if first {
+		time.AfterFunc(b.interval, func() { b.flush(key) })
+	}
+
+	select {
+	case res := <-ch:
+		return res.value, res.err
+	case <-ctx.Done():
+		var zero T
+		return zero, ctx.Err()
+	}
+}
+
+func (b *Buffer[T]) flush(key string) {
+	b.mu.Lock()
+	waiting := b.waiting[key]
+	delete(b.waiting, key)
+	b.mu.Unlock()
+
+	if len(waiting) == 0 {
+		return
+	}
+
+	start := time.Now()
+	value, err := b.fetch(b.ctx, key)
+	log.WithContext(b.ctx).Tracef("%s: fetched %s for %d waiters in %s", b.name, key, len(waiting), time.Since(start))
+
+	for _, ch := range waiting {
+		ch <- result[T]{value: value, err: err}
+	}
+}
+
+// Interval reads a buffer interval from envVar, falling back to def.
+func Interval(ctx context.Context, envVar string, def time.Duration) time.Duration {
+	value := os.Getenv(envVar)
+	interval, err := time.ParseDuration(value)
+	if err != nil {
+		if value != "" {
+			log.WithContext(ctx).Warnf("failed to parse %s: %s", envVar, err)
+		}
+		return def
+	}
+	return interval
+}

management/internals/shared/requestbuffer/buffer_test.go

@@ -0,0 +1,106 @@
+package requestbuffer
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBufferCoalescesConcurrentRequests(t *testing.T) {
+	var fetches atomic.Int32
+	buffer := New(context.Background(), "test", 50*time.Millisecond,
+		func(ctx context.Context, key string) (string, error) {
+			fetches.Add(1)
+			return key, nil
+		})
+
+	var wg sync.WaitGroup
+	for range 10 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			value, err := buffer.Get(context.Background(), "account")
+			assert.NoError(t, err)
+			assert.Equal(t, "account", value)
+		}()
+	}
+	wg.Wait()
+
+	assert.Equal(t, int32(1), fetches.Load())
+}
+
+func TestBufferSeparatesKeys(t *testing.T) {
+	keys := make(chan string, 2)
+	buffer := New(context.Background(), "test", 10*time.Millisecond,
+		func(ctx context.Context, key string) (string, error) {
+			keys <- key
+			return key, nil
+		})
+
+	var wg sync.WaitGroup
+	for _, key := range []string{"a", "b"} {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			_, err := buffer.Get(context.Background(), key)
+			assert.NoError(t, err)
+		}()
+	}
+	wg.Wait()
+	close(keys)
+
+	var fetched []string
+	for key := range keys {
+		fetched = append(fetched, key)
+	}
+	assert.ElementsMatch(t, []string{"a", "b"}, fetched)
+}
+
+func TestBufferFetchesAfterRequest(t *testing.T) {
+	var version atomic.Int32
+	buffer := New(context.Background(), "test", 10*time.Millisecond,
+		func(ctx context.Context, key string) (int32, error) {
+			return version.Load(), nil
+		})
+
+	first, err := buffer.Get(context.Background(), "account")
+	require.NoError(t, err)
+	assert.Equal(t, int32(0), first)
+
+	version.Store(1)
+
+	second, err := buffer.Get(context.Background(), "account")
+	require.NoError(t, err)
+	assert.Equal(t, int32(1), second)
+}
+
+func TestBufferPropagatesError(t *testing.T) {
+	fetchErr := errors.New("fetch failed")
+	buffer := New(context.Background(), "test", 10*time.Millisecond,
+		func(ctx context.Context, key string) (*int, error) {
+			return nil, fetchErr
+		})
+
+	value, err := buffer.Get(context.Background(), "account")
+	assert.ErrorIs(t, err, fetchErr)
+	assert.Nil(t, value)
+}
+
+func TestBufferHonorsCallerContext(t *testing.T) {
+	buffer := New(context.Background(), "test", time.Minute,
+		func(ctx context.Context, key string) (string, error) {
+			return key, nil
+		})
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
+	defer cancel()
+
+	_, err := buffer.Get(ctx, "account")
+	assert.ErrorIs(t, err, context.DeadlineExceeded)
+}

management/server/account_request_buffer.go

@@ -2,117 +2,38 @@ package server
 
 import (
 	"context"
-	"os"
-	"sync"
 	"time"
 
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/management/internals/shared/requestbuffer"
 	"github.com/netbirdio/netbird/management/server/store"
 	"github.com/netbirdio/netbird/management/server/types"
 )
 
-// AccountRequest holds the result channel to return the requested account.
-type AccountRequest struct {
-	AccountID  string
-	ResultChan chan *AccountResult
-}
-
-// AccountResult holds the account data or an error.
-type AccountResult struct {
-	Account *types.Account
-	Err     error
-}
+const defaultAccountBufferInterval = 100 * time.Millisecond
 
 type AccountRequestBuffer struct {
-	store               store.Store
-	getAccountRequests  map[string][]*AccountRequest
-	mu                  sync.Mutex
-	getAccountRequestCh chan *AccountRequest
-	bufferInterval      time.Duration
+	buffer *requestbuffer.Buffer[*types.Account]
 }
 
 func NewAccountRequestBuffer(ctx context.Context, store store.Store) *AccountRequestBuffer {
-	bufferIntervalStr := os.Getenv("NB_GET_ACCOUNT_BUFFER_INTERVAL")
-	bufferInterval, err := time.ParseDuration(bufferIntervalStr)
-	if err != nil {
-		if bufferIntervalStr != "" {
-			log.WithContext(ctx).Warnf("failed to parse account request buffer interval: %s", err)
-		}
-		bufferInterval = 100 * time.Millisecond
+	interval := requestbuffer.Interval(ctx, "NB_GET_ACCOUNT_BUFFER_INTERVAL", defaultAccountBufferInterval)
+	log.WithContext(ctx).Infof("set account request buffer interval to %s", interval)
+
+	return &AccountRequestBuffer{
+		buffer: requestbuffer.New(ctx, "account request buffer", interval, store.GetAccount),
 	}
-
-	log.WithContext(ctx).Infof("set account request buffer interval to %s", bufferInterval)
-
-	ac := AccountRequestBuffer{
-		store:               store,
-		getAccountRequests:  make(map[string][]*AccountRequest),
-		getAccountRequestCh: make(chan *AccountRequest),
-		bufferInterval:      bufferInterval,
-	}
-
-	go ac.processGetAccountRequests(ctx)
-
-	return &ac
 }
+
 func (ac *AccountRequestBuffer) GetAccountWithBackpressure(ctx context.Context, accountID string) (*types.Account, error) {
-	req := &AccountRequest{
-		AccountID:  accountID,
-		ResultChan: make(chan *AccountResult, 1),
+	account, err := ac.buffer.Get(ctx, accountID)
+	if err != nil || account == nil {
+		return account, err
 	}
 
-	log.WithContext(ctx).Tracef("requesting account %s with backpressure", accountID)
-	startTime := time.Now()
-	ac.getAccountRequestCh <- req
-
-	result := <-req.ResultChan
-	log.WithContext(ctx).Tracef("got account with backpressure after %s", time.Since(startTime))
-	return result.Account, result.Err
-}
-
-func (ac *AccountRequestBuffer) processGetAccountBatch(ctx context.Context, accountID string) {
-	ac.mu.Lock()
-	requests := ac.getAccountRequests[accountID]
-	delete(ac.getAccountRequests, accountID)
-	ac.mu.Unlock()
-
-	if len(requests) == 0 {
-		return
-	}
-
-	startTime := time.Now()
-	account, err := ac.store.GetAccount(ctx, accountID)
-	log.WithContext(ctx).Tracef("getting account %s in batch took %s", accountID, time.Since(startTime))
-	result := &AccountResult{Account: account, Err: err}
-
-	for _, req := range requests {
-		if account != nil {
-			// Shallow copy the account so each goroutine gets its own struct value.
-			// This prevents data races when callers mutate fields like Policies.
-			accountCopy := *account
-			req.ResultChan <- &AccountResult{Account: &accountCopy, Err: err}
-		} else {
-			req.ResultChan <- result
-		}
-		close(req.ResultChan)
-	}
-}
-
-func (ac *AccountRequestBuffer) processGetAccountRequests(ctx context.Context) {
-	for {
-		select {
-		case req := <-ac.getAccountRequestCh:
-			ac.mu.Lock()
-			ac.getAccountRequests[req.AccountID] = append(ac.getAccountRequests[req.AccountID], req)
-			if len(ac.getAccountRequests[req.AccountID]) == 1 {
-				go func(ctx context.Context, accountID string) {
-					time.Sleep(ac.bufferInterval)
-					ac.processGetAccountBatch(ctx, accountID)
-				}(ctx, req.AccountID)
-			}
-			ac.mu.Unlock()
-		case <-ctx.Done():
-			return
-		}
-	}
+	// Shallow copy the account so each caller gets its own struct value.
+	// This prevents data races when callers mutate fields like Policies.
+	accountCopy := *account
+	return &accountCopy, nil
 }