65 lines
1.3 KiB
Go
65 lines
1.3 KiB
Go
|
package common
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
type CustomRateLimiter struct {
|
||
|
|
tokens chan struct{}
|
||
|
|
requestQueue chan context.Context
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewCustomRateLimiter(requestsPerSecond int) *CustomRateLimiter {
|
||
|
|
rl := &CustomRateLimiter{
|
||
|
|
tokens: make(chan struct{}, requestsPerSecond),
|
||
|
|
requestQueue: make(chan context.Context, 20000), // Adjust the size accordingly
|
||
|
|
}
|
||
|
|
|
||
|
|
go rl.refillTokens(requestsPerSecond)
|
||
|
|
go rl.processQueue()
|
||
|
|
|
||
|
|
return rl
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *CustomRateLimiter) refillTokens(requestsPerSecond int) {
|
||
|
|
ticker := time.NewTicker(time.Second / time.Duration(requestsPerSecond))
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case <-ticker.C:
|
||
|
|
// Try to add a token, skip if full
|
||
|
|
select {
|
||
|
|
case rl.tokens <- struct{}{}:
|
||
|
|
default:
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *CustomRateLimiter) processQueue() {
|
||
|
|
for ctx := range rl.requestQueue {
|
||
|
|
select {
|
||
|
|
case <-rl.tokens:
|
||
|
|
if ctx.Err() == nil {
|
||
|
|
// Token acquired and context not cancelled, proceed
|
||
|
|
// Simulate processing by calling a dummy function
|
||
|
|
// processRequest() or similar
|
||
|
|
}
|
||
|
|
case <-ctx.Done():
|
||
|
|
// Context cancelled, skip
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (rl *CustomRateLimiter) Wait(ctx context.Context) error {
|
||
|
|
rl.requestQueue <- ctx
|
||
|
|
select {
|
||
|
|
case <-ctx.Done():
|
||
|
|
// Request cancelled
|
||
|
|
return ctx.Err()
|
||
|
|
case <-rl.tokens:
|
||
|
|
// Acquired token, proceed
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
}
|