github.com/sandwich-go/boost@v1.3.29/ratelimiter/limiter_atomic.go (about)

     1  package ratelimiter
     2  
     3  import (
     4  	"sync/atomic"
     5  	"time"
     6  	_ "unsafe"
     7  
     8  	"github.com/sandwich-go/boost/z"
     9  )
    10  
    11  type atomicInt64Limiter struct {
    12  	//lint:ignore U1000 Padding is unused but it is crucial to maintain performance
    13  	// of this rate limiter in case of collocation with other frequently accessed memory.
    14  	prepadding [64]byte // cache line size = 64; created to avoid false sharing.
    15  	state      int64    // unix nanoseconds of the next permissions issue.
    16  	//lint:ignore U1000 like prepadding.
    17  	postpadding [56]byte // cache line size - state size = 64 - 8; created to avoid false sharing.
    18  
    19  	perRequest time.Duration
    20  	maxSlack   time.Duration
    21  }
    22  
    23  // newAtomicBased returns a new atomic based limiter.
    24  func newAtomicInt64Based(rate int, opts ...Option) *atomicInt64Limiter {
    25  	// TODO consider moving config building to the implementation
    26  	// independent code.
    27  	config := NewOptions(opts...)
    28  	perRequest := config.per / time.Duration(rate)
    29  	l := &atomicInt64Limiter{
    30  		perRequest: perRequest,
    31  		maxSlack:   time.Duration(config.Slack) * perRequest,
    32  	}
    33  	atomic.StoreInt64(&l.state, 0)
    34  	return l
    35  }
    36  
    37  // Take blocks to ensure that the time spent between multiple
    38  // Take calls is on average time.Second/rate.
    39  func (t *atomicInt64Limiter) Take() time.Time {
    40  	var (
    41  		newTimeOfNextPermissionIssue int64
    42  		now                          int64
    43  	)
    44  	for {
    45  		now = int64(z.MonoOffset())
    46  		timeOfNextPermissionIssue := atomic.LoadInt64(&t.state)
    47  
    48  		switch {
    49  		case timeOfNextPermissionIssue == 0 || (t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest)):
    50  			// if this is our first call or t.maxSlack == 0 we need to shrink issue time to now
    51  			newTimeOfNextPermissionIssue = now
    52  		case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack)+int64(t.perRequest):
    53  			// a lot of nanoseconds passed since the last Take call
    54  			// we will limit max accumulated time to maxSlack
    55  			newTimeOfNextPermissionIssue = now - int64(t.maxSlack)
    56  		default:
    57  			// calculate the time at which our permission was issued
    58  			newTimeOfNextPermissionIssue = timeOfNextPermissionIssue + int64(t.perRequest)
    59  		}
    60  
    61  		if atomic.CompareAndSwapInt64(&t.state, timeOfNextPermissionIssue, newTimeOfNextPermissionIssue) {
    62  			break
    63  		}
    64  	}
    65  
    66  	sleepDuration := time.Duration(newTimeOfNextPermissionIssue - now)
    67  	if sleepDuration > 0 {
    68  		time.Sleep(sleepDuration)
    69  		return z.NowWithOffset(time.Duration(newTimeOfNextPermissionIssue))
    70  	}
    71  	// return now if we don't sleep as atomicLimiter does
    72  	return z.NowWithOffset(time.Duration(now))
    73  }