@@ -173,12 +173,19 @@ func (s *store) Set(ctx context.Context, key string, tokens uint64, interval tim
173173
174174// Burst adds the provided value to the bucket's currently available tokens.
175175func (s * store ) Burst (ctx context.Context , key string , tokens uint64 ) error {
176+ s .dataLock .RLock ()
177+ if b , ok := s .data [key ]; ok {
178+ s .dataLock .RUnlock ()
179+ b .burst (tokens )
180+ return nil
181+ }
182+ s .dataLock .RUnlock ()
183+
176184 s .dataLock .Lock ()
185+ // check again just in case
177186 if b , ok := s .data [key ]; ok {
178- b .lock .Lock ()
179187 s .dataLock .Unlock ()
180- b .availableTokens = b .availableTokens + tokens
181- b .lock .Unlock ()
188+ b .burst (tokens )
182189 return nil
183190 }
184191
@@ -225,18 +232,25 @@ func (s *store) purge() {
225232 case <- ticker .C :
226233 }
227234
228- s .dataLock .Lock ()
235+ s .dataLock .RLock ()
229236 now := fasttime .Now ()
237+ var deletes []string
230238 for k , b := range s .data {
231- b .lock .Lock ()
239+ b .lock .RLock ()
232240 lastTime := b .startTime + (b .lastTick * uint64 (b .interval ))
233- b .lock .Unlock ()
241+ b .lock .RUnlock ()
234242
235243 if now - lastTime > s .sweepMinTTL {
236- delete ( s . data , k )
244+ deletes = append ( deletes , k )
237245 }
238246 }
239- s .dataLock .Unlock ()
247+ s .dataLock .RUnlock ()
248+
249+ for _ , k := range deletes {
250+ s .dataLock .Lock ()
251+ delete (s .data , k )
252+ s .dataLock .Unlock ()
253+ }
240254 }
241255}
242256
@@ -325,6 +339,13 @@ func (b *bucket) take() (tokens uint64, remaining uint64, reset uint64, ok bool,
325339 return
326340}
327341
342+ // burst adds the specified number of tokens to the bucket's available tokens in a thread-safe manner.
343+ func (b * bucket ) burst (tokens uint64 ) {
344+ b .lock .Lock ()
345+ b .availableTokens = b .availableTokens + tokens
346+ b .lock .Unlock ()
347+ }
348+
328349// tick is the total number of times the current interval has occurred between
329350// when the time started (start) and the current time (curr). For example, if
330351// the start time was 12:30pm and it's currently 1:00pm, and the interval was 5
0 commit comments