@@ -2,10 +2,13 @@ package httprateredis
22
33import (
44 "context"
5+ "errors"
56 "fmt"
7+ "net"
68 "os"
79 "path/filepath"
810 "strconv"
11+ "sync/atomic"
912 "time"
1013
1114 "github.com/go-chi/httprate"
@@ -16,14 +19,11 @@ func WithRedisLimitCounter(cfg *Config) httprate.Option {
1619 if cfg .Disabled {
1720 return httprate .WithNoop ()
1821 }
19- rc , err := NewRedisLimitCounter (cfg )
20- if err != nil {
21- panic (err )
22- }
22+ rc , _ := NewRedisLimitCounter (cfg )
2323 return httprate .WithLimitCounter (rc )
2424}
2525
26- func NewRedisLimitCounter (cfg * Config ) (httprate. LimitCounter , error ) {
26+ func NewRedisLimitCounter (cfg * Config ) (* redisCounter , error ) {
2727 if cfg == nil {
2828 cfg = & Config {}
2929 }
@@ -39,20 +39,15 @@ func NewRedisLimitCounter(cfg *Config) (httprate.LimitCounter, error) {
3939 if cfg .PrefixKey == "" {
4040 cfg .PrefixKey = "httprate"
4141 }
42-
43- c , err := newClient (cfg )
44- if err != nil {
45- return nil , err
42+ if cfg .CommandTimeout == 0 {
43+ cfg .CommandTimeout = 50 * time .Millisecond
4644 }
47- return & redisCounter {
48- client : c ,
49- prefixKey : cfg .PrefixKey ,
50- }, nil
51- }
5245
53- func newClient (cfg * Config ) (* redis.Client , error ) {
54- if cfg .Client != nil {
55- return cfg .Client , nil
46+ rc := & redisCounter {
47+ prefixKey : cfg .PrefixKey ,
48+ }
49+ if ! cfg .FallbackDisabled {
50+ rc .fallbackCounter = httprate .NewLocalLimitCounter (cfg .WindowLength )
5651 }
5752
5853 var maxIdle , maxActive = cfg .MaxIdle , cfg .MaxActive
@@ -64,80 +59,112 @@ func newClient(cfg *Config) (*redis.Client, error) {
6459 }
6560
6661 address := fmt .Sprintf ("%s:%d" , cfg .Host , cfg .Port )
67- c : =redis .NewClient (& redis.Options {
62+ rc . client = redis .NewClient (& redis.Options {
6863 Addr : address ,
6964 Password : cfg .Password ,
7065 DB : cfg .DBIndex ,
7166 PoolSize : maxActive ,
7267 MaxIdleConns : maxIdle ,
7368 ClientName : cfg .ClientName ,
74- })
7569
76- status := c .Ping (context .Background ())
77- if status == nil || status .Err () != nil {
78- return nil , fmt .Errorf ("httprateredis: unable to dial redis host %v" , address )
79- }
70+ DialTimeout : cfg .CommandTimeout ,
71+ ReadTimeout : cfg .CommandTimeout ,
72+ WriteTimeout : cfg .CommandTimeout ,
73+ MinIdleConns : 1 ,
74+ MaxRetries : - 1 ,
75+ })
8076
81- return c , nil
77+ return rc , nil
8278}
8379
8480type redisCounter struct {
85- client * redis.Client
86- windowLength time.Duration
87- prefixKey string
81+ client * redis.Client
82+ windowLength time.Duration
83+ prefixKey string
84+ isRedisDown atomic.Bool
85+ fallbackCounter httprate.LimitCounter
8886}
8987
90- var _ httprate.LimitCounter = & redisCounter {}
88+ var _ httprate.LimitCounter = ( * redisCounter )( nil )
9189
9290func (c * redisCounter ) Config (requestLimit int , windowLength time.Duration ) {
9391 c .windowLength = windowLength
92+ if c .fallbackCounter != nil {
93+ c .fallbackCounter .Config (requestLimit , windowLength )
94+ }
9495}
9596
9697func (c * redisCounter ) Increment (key string , currentWindow time.Time ) error {
9798 return c .IncrementBy (key , currentWindow , 1 )
9899}
99100
100- func (c * redisCounter ) IncrementBy (key string , currentWindow time.Time , amount int ) error {
101- ctx := context .Background ()
102- conn := c .client
103-
101+ func (c * redisCounter ) IncrementBy (key string , currentWindow time.Time , amount int ) (err error ) {
102+ if c .fallbackCounter != nil {
103+ if c .isRedisDown .Load () {
104+ return c .fallbackCounter .IncrementBy (key , currentWindow , amount )
105+ }
106+ defer func () {
107+ if err != nil {
108+ // On redis network error, fallback to local in-memory counter.
109+ var netErr net.Error
110+ if errors .As (err , & netErr ) || errors .Is (err , redis .ErrClosed ) {
111+ go c .fallback ()
112+ err = c .fallbackCounter .IncrementBy (key , currentWindow , amount )
113+ }
114+ }
115+ }()
116+ }
117+
118+ ctx := context .Background () // Note: We use timeouts set up on the Redis client directly.
104119 hkey := c .limitCounterKey (key , currentWindow )
105120
106- pipe := conn .TxPipeline ()
121+ pipe := c . client .TxPipeline ()
107122 incrCmd := pipe .IncrBy (ctx , hkey , int64 (amount ))
108123 expireCmd := pipe .Expire (ctx , hkey , c .windowLength * 3 )
109- _ , err := pipe .Exec (ctx )
124+
125+ _ , err = pipe .Exec (ctx )
110126 if err != nil {
111127 return fmt .Errorf ("httprateredis: redis transaction failed: %w" , err )
112128 }
113-
114129 if err := incrCmd .Err (); err != nil {
115130 return fmt .Errorf ("httprateredis: redis incr failed: %w" , err )
116131 }
117-
118132 if err := expireCmd .Err (); err != nil {
119133 return fmt .Errorf ("httprateredis: redis expire failed: %w" , err )
120134 }
121135
122136 return nil
123137}
124138
125- func (c * redisCounter ) Get (key string , currentWindow , previousWindow time.Time ) (int , int , error ) {
126- ctx := context .Background ()
127- conn := c .client
139+ func (c * redisCounter ) Get (key string , currentWindow , previousWindow time.Time ) (curr int , prev int , err error ) {
140+ if c .fallbackCounter != nil {
141+ if c .isRedisDown .Load () {
142+ return c .fallbackCounter .Get (key , currentWindow , previousWindow )
143+ }
144+ defer func () {
145+ if err != nil {
146+ // On redis network error, fallback to local in-memory counter.
147+ var netErr net.Error
148+ if errors .As (err , & netErr ) || errors .Is (err , redis .ErrClosed ) {
149+ go c .fallback ()
150+ curr , prev , err = c .fallbackCounter .Get (key , currentWindow , previousWindow )
151+ }
152+ }
153+ }()
154+ }
155+
156+ ctx := context .Background () // Note: We use timeouts set up on the Redis client directly.
128157
129158 currKey := c .limitCounterKey (key , currentWindow )
130159 prevKey := c .limitCounterKey (key , previousWindow )
131160
132- values , err := conn .MGet (ctx , currKey , prevKey ).Result ()
161+ values , err := c . client .MGet (ctx , currKey , prevKey ).Result ()
133162 if err != nil {
134163 return 0 , 0 , fmt .Errorf ("httprateredis: redis mget failed: %w" , err )
135164 } else if len (values ) != 2 {
136165 return 0 , 0 , fmt .Errorf ("httprateredis: redis mget returned wrong number of keys: %v, expected 2" , len (values ))
137166 }
138167
139- var curr , prev int
140-
141168 // MGET always returns slice with nil or "string" values, even if the values
142169 // were created with the INCR command. Ignore error if we can't parse the number.
143170 if values [0 ] != nil {
@@ -152,6 +179,28 @@ func (c *redisCounter) Get(key string, currentWindow, previousWindow time.Time)
152179 return curr , prev , nil
153180}
154181
182+ func (c * redisCounter ) IsRedisDown () bool {
183+ return c .isRedisDown .Load ()
184+ }
185+
186+ func (c * redisCounter ) fallback () {
187+ // Fallback to in-memory counter.
188+ wasAlreadyDown := c .isRedisDown .Swap (true )
189+ if wasAlreadyDown {
190+ return
191+ }
192+
193+ // Try to re-connect to redis every 50ms.
194+ for {
195+ err := c .client .Ping (context .Background ()).Err ()
196+ if err == nil {
197+ c .isRedisDown .Store (false )
198+ return
199+ }
200+ time .Sleep (50 * time .Millisecond )
201+ }
202+ }
203+
155204func (c * redisCounter ) limitCounterKey (key string , window time.Time ) string {
156205 return fmt .Sprintf ("%s:%d" , c .prefixKey , httprate .LimitCounterKey (key , window ))
157206}
0 commit comments