github.com/blend/go-sdk@v1.20220411.3/ratelimiter/wait.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package ratelimiter
     9  
    10  import (
    11  	"context"
    12  	"time"
    13  )
    14  
    15  // Wait is a type that allows you to throttle actions
    16  // with sleeps based on a desired rate.
    17  type Wait struct {
    18  	NumberOfActions int64
    19  	Quantum         time.Duration
    20  }
    21  
    22  // Wait waits for a calculated throttling time based on the input options.
    23  func (w Wait) Wait(ctx context.Context, actions int64, quantum time.Duration) error {
    24  	return w.WaitTimer(ctx, actions, quantum, nil)
    25  }
    26  
    27  // WaitTimer waits with a given (re-used) timer reference.
    28  func (w Wait) WaitTimer(ctx context.Context, actions int64, quantum time.Duration, after *time.Timer) error {
    29  	waitFor := w.Calculate(actions, quantum)
    30  	if waitFor < 0 {
    31  		return nil
    32  	}
    33  	if after == nil {
    34  		after = time.NewTimer(waitFor)
    35  	} else {
    36  		after.Reset(waitFor)
    37  	}
    38  	defer after.Stop()
    39  	select {
    40  	case <-ctx.Done():
    41  		return context.Canceled
    42  	case <-after.C:
    43  		return nil
    44  	}
    45  }
    46  
    47  // Calculate takes the observed rate and the desired rate, and returns a quantum to sleep for
    48  // that adjusts the observed rate to match the desired rate.
    49  //
    50  // If the observed rate is _lower_ than the desired rate, the returned value will be negative
    51  // and you're free to ignore it.
    52  //
    53  // If the observed rate is _higher_ than the desired rate, a positive duration will be returned
    54  // which you can pass to a `time.Sleep(...)` or similar.
    55  //
    56  // The wait quantum is derrived from the following algebraic steps (where ? is what we're solving for):
    57  //
    58  //    pb/(pq+?) = rb/rq
    59  //    1/(pq+?) = rb/pb*rq
    60  //    pq+? = (pb*rq)/rb
    61  //    ? = ((pb*rq)/rb) - pq
    62  //
    63  func (w Wait) Calculate(actions int64, quantum time.Duration) time.Duration {
    64  	return time.Duration(((actions * int64(w.Quantum)) / w.NumberOfActions) - int64(quantum))
    65  }