github.com/safing/portbase@v0.19.5/utils/call_limiter.go (about)

     1  package utils
     2  
     3  import (
     4  	"sync"
     5  	"sync/atomic"
     6  	"time"
     7  )
     8  
     9  // CallLimiter bundles concurrent calls and optionally limits how fast a function is called.
    10  type CallLimiter struct {
    11  	pause time.Duration
    12  
    13  	inLock   sync.Mutex
    14  	lastExec time.Time
    15  
    16  	waiters atomic.Int32
    17  	outLock sync.Mutex
    18  }
    19  
    20  // NewCallLimiter returns a new call limiter.
    21  // Set minPause to zero to disable the minimum pause between calls.
    22  func NewCallLimiter(minPause time.Duration) *CallLimiter {
    23  	return &CallLimiter{
    24  		pause: minPause,
    25  	}
    26  }
    27  
    28  // Do executes the given function.
    29  // All concurrent calls to Do are bundled and return when f() finishes.
    30  // Waits until the minimum pause is over before executing f() again.
    31  func (l *CallLimiter) Do(f func()) {
    32  	// Wait for the previous waiters to exit.
    33  	l.inLock.Lock()
    34  
    35  	// Defer final unlock to safeguard from panics.
    36  	defer func() {
    37  		// Execution is finished - leave.
    38  		// If we are the last waiter, let the next batch in.
    39  		if l.waiters.Add(-1) == 0 {
    40  			l.inLock.Unlock()
    41  		}
    42  	}()
    43  
    44  	// Check if we are the first waiter.
    45  	if l.waiters.Add(1) == 1 {
    46  		// Take the lead on this execution run.
    47  		l.lead(f)
    48  	} else {
    49  		// We are not the first waiter, let others in.
    50  		l.inLock.Unlock()
    51  	}
    52  
    53  	// Wait for execution to complete.
    54  	l.outLock.Lock()
    55  	l.outLock.Unlock() //nolint:staticcheck
    56  
    57  	// Last statement is in defer above.
    58  }
    59  
    60  func (l *CallLimiter) lead(f func()) {
    61  	// Make all others wait while we execute the function.
    62  	l.outLock.Lock()
    63  
    64  	// Unlock in lock until execution is finished.
    65  	l.inLock.Unlock()
    66  
    67  	// Transition from out lock to in lock when done.
    68  	defer func() {
    69  		// Update last execution time.
    70  		l.lastExec = time.Now().UTC()
    71  		// Stop newcomers from waiting on previous execution.
    72  		l.inLock.Lock()
    73  		// Allow waiters to leave.
    74  		l.outLock.Unlock()
    75  	}()
    76  
    77  	// Wait for the minimum duration between executions.
    78  	if l.pause > 0 {
    79  		sinceLastExec := time.Since(l.lastExec)
    80  		if sinceLastExec < l.pause {
    81  			time.Sleep(l.pause - sinceLastExec)
    82  		}
    83  	}
    84  
    85  	// Execute.
    86  	f()
    87  }