58 lines
1.2 KiB
Go
58 lines
1.2 KiB
Go
package throttling
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
// Client is an interface that calls something with a payload.
|
|
type Client interface {
|
|
Call(*Payload)
|
|
}
|
|
|
|
// Payload is some payload a Client would send in a call.
|
|
type Payload struct {
|
|
ClientResult string
|
|
W http.ResponseWriter
|
|
R *http.Request
|
|
}
|
|
|
|
// CallFunction allows burst rate limiting client calls with the
|
|
// payloads.
|
|
func CallFunction(ctx context.Context, client Client, payload *Payload, throttle <-chan time.Time) {
|
|
<-throttle // rate limit our client calls
|
|
client.Call(payload)
|
|
}
|
|
|
|
func CreateThrottle(ctx context.Context, burstLimit int, rateLimit time.Duration) <-chan time.Time {
|
|
throttle := make(chan time.Time, burstLimit)
|
|
for i := 0; i < burstLimit; i++ {
|
|
throttle <- time.Now()
|
|
}
|
|
|
|
go func() {
|
|
ticker := time.NewTicker(rateLimit)
|
|
defer ticker.Stop()
|
|
for t := range ticker.C {
|
|
select {
|
|
case throttle <- t:
|
|
{
|
|
fmt.Println("Add bucket to throttle")
|
|
}
|
|
case <-ctx.Done():
|
|
{
|
|
fmt.Println("Ticker done")
|
|
return // exit goroutine when surrounding function returns
|
|
}
|
|
default:
|
|
{
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
return throttle
|
|
}
|