From 1a5dacd60904a20abf13b1995a4bbb46fccbf739 Mon Sep 17 00:00:00 2001 From: Dylan William Hardison Date: Thu, 22 Dec 2016 16:18:23 -0500 Subject: Bug 1299855 - Implement token-bucket rate limiting on top of memcached (#23) * cleanup --- Bugzilla/Memcached.pm | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) (limited to 'Bugzilla') diff --git a/Bugzilla/Memcached.pm b/Bugzilla/Memcached.pm index a1b8a5ac7..6ada1adf8 100644 --- a/Bugzilla/Memcached.pm +++ b/Bugzilla/Memcached.pm @@ -13,6 +13,7 @@ use warnings; use Bugzilla::Error; use Scalar::Util qw(blessed); +use List::Util qw(sum); use Bugzilla::Util qw(trick_taint); use URI::Escape; use Encode; @@ -20,6 +21,7 @@ use Sys::Syslog qw(:DEFAULT); # memcached keys have a maximum length of 250 bytes use constant MAX_KEY_LENGTH => 250; +use constant RATE_LIMIT_PREFIX => "rate:"; sub _new { my $invocant = shift; @@ -157,6 +159,26 @@ sub clear { } } +sub should_rate_limit { + my ($self, $name, $rate_max, $rate_seconds, $tries) = @_; + my $prefix = RATE_LIMIT_PREFIX . $name . ':'; + my $memcached = $self->{memcached}; + + $tries //= 3; + + for (0 .. $tries) { + my $now = time; + my ($key, @keys) = map { $prefix . ( $now - $_ ) } 0 .. $rate_seconds; + $memcached->add($key, 0, $rate_seconds+1); + my $tokens = $memcached->get_multi(@keys); + my $cas = $memcached->gets($key); + $tokens->{$key} = $cas->[1]++; + return 1 if sum(values %$tokens) >= $rate_max; + return 0 if $memcached->cas($key, @$cas, $rate_seconds+1); + } + return 1; +} + sub clear_all { my ($self) = @_; return unless $self->{memcached}; -- cgit v1.2.3-24-g4f1b