forked from RealistikOsu/RealistikAPI
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Create `limit' package, remove need of login_attempts.go
- Loading branch information
Showing
4 changed files
with
128 additions
and
133 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,122 @@ | ||
package limit | ||
|
||
import ( | ||
"fmt" | ||
"sync" | ||
"time" | ||
) | ||
|
||
// Request is a Request with DefaultLimiter. | ||
func Request(u string, perMinute int) { DefaultLimiter.Request(u, perMinute) } | ||
|
||
// NonBlockingRequest is a NonBlockingRequest with DefaultLimiter. | ||
func NonBlockingRequest(u string, perMinute int) bool { | ||
return DefaultLimiter.NonBlockingRequest(u, perMinute) | ||
} | ||
|
||
// DefaultLimiter is the RateLimiter used by the package-level | ||
// Request and NonBlockingRequest. | ||
var DefaultLimiter = &RateLimiter{ | ||
Map: make(map[string]chan struct{}), | ||
Mutex: &sync.RWMutex{}, | ||
} | ||
|
||
// RateLimiter is a simple rate limiter. | ||
type RateLimiter struct { | ||
Map map[string]chan struct{} | ||
Mutex *sync.RWMutex | ||
} | ||
|
||
// Request is a simple request. Blocks until it can make the request. | ||
func (s *RateLimiter) Request(u string, perMinute int) { | ||
s.request(u, perMinute, true) | ||
} | ||
|
||
// NonBlockingRequest checks if it can do a request. If it can't, it returns | ||
// false, else it returns true if the request succeded. | ||
func (s *RateLimiter) NonBlockingRequest(u string, perMinute int) bool { | ||
return s.request(u, perMinute, false) | ||
} | ||
|
||
func (s *RateLimiter) request(u string, perMinute int, blocking bool) bool { | ||
s.check() | ||
s.Mutex.RLock() | ||
c, exists := s.Map[u] | ||
s.Mutex.RUnlock() | ||
if !exists { | ||
c = makePrefilledChan(perMinute) | ||
s.Mutex.Lock() | ||
// Now that we have exclusive read and write-access, we want to | ||
// make sure we don't overwrite an existing channel. Otherwise, | ||
// race conditions and panic happen. | ||
if cNew, exists := s.Map[u]; exists { | ||
c = cNew | ||
s.Mutex.Unlock() | ||
} else { | ||
s.Map[u] = c | ||
s.Mutex.Unlock() | ||
<-c | ||
go s.filler(u, perMinute) | ||
} | ||
} | ||
return rcv(c, blocking) | ||
} | ||
|
||
// rcv receives from a channel, but if blocking is true it waits til something | ||
// is received and always returns true, otherwise if it can't receive it | ||
// returns false. | ||
func rcv(c chan struct{}, blocking bool) bool { | ||
if blocking { | ||
<-c | ||
return true | ||
} | ||
select { | ||
case <-c: | ||
return true | ||
default: | ||
return false | ||
} | ||
} | ||
|
||
func (s *RateLimiter) filler(el string, perMinute int) { | ||
defer func() { | ||
r := recover() | ||
if r != nil { | ||
fmt.Println(r) | ||
} | ||
}() | ||
|
||
s.Mutex.RLock() | ||
c := s.Map[el] | ||
s.Mutex.RUnlock() | ||
for { | ||
select { | ||
case c <- struct{}{}: | ||
time.Sleep(time.Minute / time.Duration(perMinute)) | ||
default: // c is full | ||
s.Mutex.Lock() | ||
close(c) | ||
delete(s.Map, el) | ||
s.Mutex.Unlock() | ||
return | ||
} | ||
} | ||
} | ||
|
||
// check makes sure the map and the mutex are properly initialised. | ||
func (s *RateLimiter) check() { | ||
if s.Map == nil { | ||
s.Map = make(map[string]chan struct{}) | ||
} | ||
if s.Mutex { | ||
s.Mutex = new(sync.RWMutex) | ||
} | ||
} | ||
|
||
func makePrefilledChan(l int) chan struct{} { | ||
c := make(chan struct{}, l) | ||
for i := 0; i < l; i++ { | ||
c <- struct{}{} | ||
} | ||
return c | ||
} |