maddy

  1/*
  2Maddy Mail Server - Composable all-in-one email server.
  3Copyright © 2019-2020 Max Mazurov <fox.cpp@disroot.org>, Maddy Mail Server contributors
  4
  5This program is free software: you can redistribute it and/or modify
  6it under the terms of the GNU General Public License as published by
  7the Free Software Foundation, either version 3 of the License, or
  8(at your option) any later version.
  9
 10This program is distributed in the hope that it will be useful,
 11but WITHOUT ANY WARRANTY; without even the implied warranty of
 12MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 13GNU General Public License for more details.
 14
 15You should have received a copy of the GNU General Public License
 16along with this program.  If not, see <https://www.gnu.org/licenses/>.
 17*/
 18
 19package limiters
 20
 21import (
 22	"context"
 23	"errors"
 24	"time"
 25)
 26
 27var ErrClosed = errors.New("limiters: Rate bucket is closed")
 28
 29// Rate structure implements a basic rate-limiter for requests using the token
 30// bucket approach.
 31//
 32// Take() is expected to be called before each request. Excessive calls will
 33// block. Timeouts can be implemented using the TakeContext method.
 34//
 35// Rate.Close causes all waiting Take to return false. TakeContext returns
 36// ErrClosed in this case.
 37//
 38// If burstSize = 0, all methods are no-op and always succeed.
 39type Rate struct {
 40	bucket chan struct{}
 41	stop   chan struct{}
 42}
 43
 44func NewRate(burstSize int, interval time.Duration) Rate {
 45	r := Rate{
 46		bucket: make(chan struct{}, burstSize),
 47		stop:   make(chan struct{}),
 48	}
 49
 50	if burstSize == 0 {
 51		return r
 52	}
 53
 54	for i := 0; i < burstSize; i++ {
 55		r.bucket <- struct{}{}
 56	}
 57
 58	go r.fill(burstSize, interval)
 59	return r
 60}
 61
 62func (r Rate) fill(burstSize int, interval time.Duration) {
 63	t := time.NewTimer(interval)
 64	defer t.Stop()
 65	for {
 66		t.Reset(interval)
 67		select {
 68		case <-t.C:
 69		case <-r.stop:
 70			close(r.bucket)
 71			return
 72		}
 73
 74	fill:
 75		for i := 0; i < burstSize; i++ {
 76			select {
 77			case r.bucket <- struct{}{}:
 78			default:
 79				// If there are no Take pending and the bucket is already
 80				// full - don't block.
 81				break fill
 82			}
 83		}
 84	}
 85}
 86
 87func (r Rate) Take() bool {
 88	if cap(r.bucket) == 0 {
 89		return true
 90	}
 91
 92	_, ok := <-r.bucket
 93	return ok
 94}
 95
 96func (r Rate) TakeContext(ctx context.Context) error {
 97	if cap(r.bucket) == 0 {
 98		return nil
 99	}
100
101	select {
102	case _, ok := <-r.bucket:
103		if !ok {
104			return ErrClosed
105		}
106		return nil
107	case <-ctx.Done():
108		return ctx.Err()
109	}
110}
111
112func (r Rate) Release() {
113}
114
115func (r Rate) Close() {
116	close(r.stop)
117}