go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/rate/wait.go (about)

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