rename request buffer and update default interval (#2459)

management/server/account_request_buffer.go

@@ -0,0 +1,108 @@
+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
+		}
+	}
+}