diff --git a/management/internals/controllers/network_map/controller/controller.go b/management/internals/controllers/network_map/controller/controller.go index 5567af0b2..778f73c01 100644 --- a/management/internals/controllers/network_map/controller/controller.go +++ b/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 diff --git a/management/internals/shared/requestbuffer/buffer.go b/management/internals/shared/requestbuffer/buffer.go new file mode 100644 index 000000000..c3823776c --- /dev/null +++ b/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 +} diff --git a/management/internals/shared/requestbuffer/buffer_test.go b/management/internals/shared/requestbuffer/buffer_test.go new file mode 100644 index 000000000..9e356145e --- /dev/null +++ b/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) +} diff --git a/management/server/account_request_buffer.go b/management/server/account_request_buffer.go index e1672c2d0..792099431 100644 --- a/management/server/account_request_buffer.go +++ b/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 }