-
Notifications
You must be signed in to change notification settings - Fork 5
/
autothrust.go
88 lines (72 loc) · 2.31 KB
/
autothrust.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
package main
import (
"net/http"
"strconv"
"sync"
"time"
log "github.com/sirupsen/logrus"
)
// Throttle is the interface to get throttlig handler
type Throttle interface {
NewThrottle(handler http.Handler) http.Handler
}
// AutoThrustFactory news the AutoThrust instance
type AutoThrustFactory struct {
logger *log.Logger
defaultRetryTryAfterInSeconds int
}
// NewThrottle constructs a http handler for throttling
func (f *AutoThrustFactory) NewThrottle(handlerToWrap http.Handler) http.Handler {
return &AutoThrust{f.logger, handlerToWrap, sync.Map{}, f.defaultRetryTryAfterInSeconds}
}
// AutoThrust is a middleware handler that does request throttling
type AutoThrust struct {
logger *log.Logger
handler http.Handler
cache sync.Map
defaultRetryTryAfterInSeconds int
}
type statusRecorder struct {
http.ResponseWriter
status int
}
func (rec *statusRecorder) WriteHeader(code int) {
rec.status = code
rec.ResponseWriter.WriteHeader(code)
}
type cacheItem struct {
until time.Time
}
// ServeHTTP handles the request by passing it to the real
// handler and logging the request details
func (t *AutoThrust) ServeHTTP(w http.ResponseWriter, r *http.Request) {
if c, ok := t.cache.Load(r.RequestURI); ok {
shouldAfter := c.(*cacheItem).until
if time.Now().After(shouldAfter) {
t.cache.Delete(r.RequestURI)
} else {
t.logger.Infof("Request %s will be throttled until %v", r.RequestURI, shouldAfter)
// Add 5 seconds delay to avoid client side spin
time.Sleep(5 * time.Second)
w.WriteHeader(429)
return
}
}
rec := statusRecorder{w, 0}
t.handler.ServeHTTP(&rec, r)
if rec.status == 429 {
// Default retry after is 3 minutes
retryAfter := time.Now().Add(time.Duration(t.defaultRetryTryAfterInSeconds) * time.Second)
retryAfterStr := w.Header().Get("Retry-After")
if retryAfterStr != "" {
if retrySec, _ := strconv.Atoi(retryAfterStr); retrySec > 0 {
retryAfter = time.Now().Add(time.Duration(retrySec) * time.Second)
} else if retryTime, err := time.Parse(time.RFC1123, retryAfterStr); err == nil {
retryAfter = retryTime
}
}
t.logger.Infof("Setting retry after to %s for request %v", retryAfter, r.RequestURI)
ci := &cacheItem{until: retryAfter}
t.cache.Store(r.RequestURI, ci)
}
}