ripple-api/limit/limit.go

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 == nil {
		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
}