github.com/asynkron/protoactor-go@v0.0.0-20240308120642-ef91a6abee75/actor/throttler.go (about)

     1  package actor
     2  
     3  import (
     4  	"log/slog"
     5  	"sync/atomic"
     6  	"time"
     7  )
     8  
     9  type ShouldThrottle func() Valve
    10  
    11  type Valve int32
    12  
    13  const (
    14  	Open Valve = iota
    15  	Closing
    16  	Closed
    17  )
    18  
    19  // NewThrottle
    20  // This has no guarantees that the throttle opens exactly after the period, since it is reset asynchronously
    21  // Throughput has been prioritized over exact re-opening
    22  // throttledCallBack, This will be called with the number of events what was throttled after the period
    23  func NewThrottle(maxEventsInPeriod int32, period time.Duration, throttledCallBack func(int32)) ShouldThrottle {
    24  	currentEvents := int32(0)
    25  
    26  	startTimer := func(duration time.Duration, back func(int32)) {
    27  		go func() {
    28  			// crete ticker to mimic sleep, we do not want to put the goroutine to sleep
    29  			// as it will schedule it out of the P making a syscall, we just want it to
    30  			// halt for the given period of time
    31  			ticker := time.NewTicker(duration)
    32  			defer ticker.Stop()
    33  			<-ticker.C // wait for the ticker to tick once
    34  
    35  			timesCalled := atomic.SwapInt32(&currentEvents, 0)
    36  			if timesCalled > maxEventsInPeriod {
    37  				throttledCallBack(timesCalled - maxEventsInPeriod)
    38  			}
    39  		}()
    40  	}
    41  
    42  	return func() Valve {
    43  		tries := atomic.AddInt32(&currentEvents, 1)
    44  		if tries == 1 {
    45  			startTimer(period, throttledCallBack)
    46  		}
    47  
    48  		if tries == maxEventsInPeriod {
    49  			return Closing
    50  		} else if tries > maxEventsInPeriod {
    51  			return Closed
    52  		} else {
    53  			return Open
    54  		}
    55  	}
    56  }
    57  
    58  func NewThrottleWithLogger(logger *slog.Logger, maxEventsInPeriod int32, period time.Duration, throttledCallBack func(*slog.Logger, int32)) ShouldThrottle {
    59  	currentEvents := int32(0)
    60  
    61  	startTimer := func(duration time.Duration, back func(*slog.Logger, int32)) {
    62  		go func() {
    63  			// crete ticker to mimic sleep, we do not want to put the goroutine to sleep
    64  			// as it will schedule it out of the P making a syscall, we just want it to
    65  			// halt for the given period of time
    66  			ticker := time.NewTicker(duration)
    67  			defer ticker.Stop()
    68  			<-ticker.C // wait for the ticker to tick once
    69  
    70  			timesCalled := atomic.SwapInt32(&currentEvents, 0)
    71  			if timesCalled > maxEventsInPeriod {
    72  				throttledCallBack(logger, timesCalled-maxEventsInPeriod)
    73  			}
    74  		}()
    75  	}
    76  
    77  	return func() Valve {
    78  		tries := atomic.AddInt32(&currentEvents, 1)
    79  		if tries == 1 {
    80  			startTimer(period, throttledCallBack)
    81  		}
    82  
    83  		if tries == maxEventsInPeriod {
    84  			return Closing
    85  		} else if tries > maxEventsInPeriod {
    86  			return Closed
    87  		} else {
    88  			return Open
    89  		}
    90  	}
    91  }