netbird/management/server/account_request_buffer.go

package server

import (
	"context"
	"os"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
)

// 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 *Account
	Err     error
}

type AccountRequestBuffer struct {
	store               Store
	getAccountRequests  map[string][]*AccountRequest
	mu                  sync.Mutex
	getAccountRequestCh chan *AccountRequest
	bufferInterval      time.Duration
}

func NewAccountRequestBuffer(ctx context.Context, 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
	}

	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) (*Account, error) {
	req := &AccountRequest{
		AccountID:  accountID,
		ResultChan: make(chan *AccountResult, 1),
	}

	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 {
		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
		}
	}
}