[management] Add buffering for getAccount requests during login (#2449)

management/server/account.go

@@ -161,6 +161,8 @@ type DefaultAccountManager struct {
 	eventStore           activity.Store
 	geo                  *geolocation.Geolocation
 
+	cache *AccountCache
+
 	// singleAccountMode indicates whether the instance has a single account.
 	// If true, then every new user will end up under the same account.
 	// This value will be set to false if management service has more than one account.
@@ -967,6 +969,7 @@ func BuildManager(
 		userDeleteFromIDPEnabled: userDeleteFromIDPEnabled,
 		integratedPeerValidator:  integratedPeerValidator,
 		metrics:                  metrics,
+		cache:                    NewAccountCache(ctx, store),
 	}
 	allAccounts := store.GetAllAccounts(ctx)
 	// enable single account mode only if configured by user and number of existing accounts is not grater than 1

management/server/account_cache.go

@@ -0,0 +1,106 @@
+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 AccountCache struct {
+	store               Store
+	getAccountRequests  map[string][]*AccountRequest
+	mu                  sync.Mutex
+	getAccountRequestCh chan *AccountRequest
+	bufferInterval      time.Duration
+}
+
+func NewAccountCache(ctx context.Context, store Store) *AccountCache {
+	bufferIntervalStr := os.Getenv("NB_GET_ACCOUNT_BUFFER_INTERVAL")
+	bufferInterval, err := time.ParseDuration(bufferIntervalStr)
+	if err != nil && bufferIntervalStr != "" {
+		log.WithContext(ctx).Warnf("failed to parse account cache buffer interval: %s", err)
+		bufferInterval = 300 * time.Millisecond
+	}
+
+	log.WithContext(ctx).Infof("set account cache buffer interval to %s", bufferInterval)
+
+	ac := AccountCache{
+		store:               store,
+		getAccountRequests:  make(map[string][]*AccountRequest),
+		getAccountRequestCh: make(chan *AccountRequest),
+		bufferInterval:      bufferInterval,
+	}
+
+	go ac.processGetAccountRequests(ctx)
+
+	return &ac
+}
+func (ac *AccountCache) 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 *AccountCache) 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 *AccountCache) 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
+		}
+	}
+}