Skip to content
Snippets Groups Projects
Commit 83dda4e8 authored by Andy Yates's avatar Andy Yates
Browse files

[ENSCORESW-736]. Removing the old throttle extensions.

We used to extend existing code to provide decent throttling. Lets DIY it and we can have a far nicer way of doing it.
parent bd78a4ed
No related branches found
No related tags found
No related merge requests found
package Plack::Middleware::Throttle::Backend::Memcached;
use Moose;
use Plack::Util;
has 'driver' => ( isa => 'Str', is => 'ro', required => 1);
has 'args' => ( isa => 'HashRef', is => 'ro', required => 1);
has 'expire' => ( isa => 'Int', is => 'ro', required => 1);
has '_backend' => ( isa => 'Ref', is => 'ro', builder => '_build_backend', lazy => 1);
sub _build_backend {
my ($self) = @_;
my $driver = $self->driver();
my $args = $self->args();
Plack::Util::load_class($driver);
return $driver->new(%{$args});
}
sub incr {
my ($self, $key, $value) = @_;
$value ||= 1;
my $backend = $self->_backend();
my $expire = $self->expire();
my $v = $backend->incr($key, $value, $expire);
if(! $v) {
#Try to add the value 1 if it didn't exist before
$v = $backend->add($key, $value, $expire);
if (! $v) {
#Means someone else called the add() so we can use increment again
$v = $backend->incr($key, $value, $expire);
}
}
return $v;
}
sub get {
my ($self, $key) = @_;
my $backend = $self->_backend();
return $backend->get($key);
}
sub set {
my ($self, $key, $value) = @_;
my $backend = $self->_backend();
my $expire = $self->expire();
return $backend->set($key, $value, $expire);
}
__PACKAGE__->meta->make_immutable;
1;
\ No newline at end of file
package Plack::Middleware::Throttle::Hour;
use Moose;
use DateTime;
extends 'Plack::Middleware::Throttle::ReLimiter';
sub cache_key {
my ( $self, $env ) = @_;
$self->client_identifier($env) . "_"
. DateTime->now->strftime("%Y-%m-%d-%H");
}
sub reset_time {
my $dt = DateTime->now;
3600 - (( 60 * $dt->minute ) + $dt->second);
}
1;
\ No newline at end of file
package Plack::Middleware::Throttle::Minute;
use Moose;
use DateTime;
extends 'Plack::Middleware::Throttle::ReLimiter';
sub cache_key {
my ($self, $env) = @_;
return $self->client_identifier($env) . q{_}
. DateTime->now->strftime("%Y-%m-%d-%H-%M");
}
sub reset_time {
my $dt = DateTime->now;
return period() - $dt->second;
}
sub period {
return 60;
}
__PACKAGE__->meta->make_immutable;
1;
\ No newline at end of file
package Plack::Middleware::Throttle::ReLimiter;
use Moose;
extends 'Plack::Middleware::Throttle::Limiter';
# Provides a way of controlling the retry after env variable by
# adding a constant amount of time. This is useful if you are using
# the per-second limiter as a way of controlling per second bursts
has 'retry_after_addition' => ( isa => 'Num', is => 'ro', lazy => 1, default => 0 );
use Plack::Util qw//;
use Scalar::Util qw//;
sub _create_backend {
my ( $self, $backend ) = @_;
if (! defined $backend ) {
Plack::Util::load_class("Plack::Middleware::Throttle::Backend::Hash");
return Plack::Middleware::Throttle::Backend::Hash->new;
}
return $backend if defined $backend && Scalar::Util::blessed $backend;
die "backend must be a cache object";
}
sub client_identifier {
my ($self, $env) = @_;
if($env->{HTTP_X_FORWARDED_HOST}) {
return $self->key_prefix."_".$env->{HTTP_X_FORWARDED_HOST};
}
return $self->SUPER::client_identifier($env);
}
override 'over_rate_limit' => sub {
my ($self) = @_;
my $res = super();
my $headers = $res->[1];
my $reset_time = $self->reset_time();
$reset_time += $self->retry_after_addition();
Plack::Util::header_set( $headers, 'Retry-After', $reset_time );
return $res;
};
1;
package Plack::Middleware::Throttle::Second;
use Moose;
use Time::HiRes qw//;
extends 'Plack::Middleware::Throttle::ReLimiter';
sub cache_key {
my ($self, $env) = @_;
return sprintf('%s_%d', $self->client_identifier($env),CORE::time());
}
sub reset_time {
my ($self) = @_;
my ($seconds, $microseconds) = Time::HiRes::gettimeofday;
#Current millis from microseconds
my $millis = $microseconds/1000;
#Period allowed in milliseconds
my $period = ($self->period()*1000);
#The time remaining in seconds before we will allow more requests
my $diff = ($period - $millis)/1000;
return $diff;
}
sub period {
return 1;
}
1;
\ No newline at end of file
use strict;
use warnings;
use Test::More;
use Plack::Util qw/headers/;
use Plack::Builder;
use Plack::Test;
use HTTP::Request::Common;
use Time::HiRes qw/sleep/;
use Plack::Middleware::Throttle::Second;
use Plack::Middleware::Throttle::Hour;
use Plack::Middleware::Throttle::Backend::Hash;
sub get {
my ($cb) = @_;
my $req = GET "http://localhost/";
my $res = $cb->($req);
return $res;
}
sub success {
my ($cb) = @_;
my $res = get($cb);
is($res->code, 200, 'http response is 200');
return;
}
{
my $handler = builder {
enable "Throttle::Hour",
max => 4,
backend => Plack::Middleware::Throttle::Backend::Hash->new();
enable "Throttle::Second",
max => 2,
backend => Plack::Middleware::Throttle::Backend::Hash->new();
sub {
[ '200', [ 'Content-Type' => 'text/html' ], ['hello world'] ]
};
};
test_psgi
app => $handler,
client => sub {
my $cb = shift;
note 'Not limited';
success($cb) for 1..2;
{
note 'Now expecting failure';
my $res = get($cb);
my $retry_after = $res->header('Retry-After');
cmp_ok($retry_after, '<', 1, 'Checking Retry-After is smaller than a second');
sleep($retry_after);
}
note 'Retrying after the alotted time has passed';
success($cb);
{
note 'Now we will hit into the hour rate limiter';
my $res = get($cb);
my $retry_after = $res->header('Retry-After');
# Should always be a second as we only work in those time-units
cmp_ok($retry_after, '>=', 1, 'Checking Retry-After is larger than 1 second');
# Can't be more than an hour ever
cmp_ok($retry_after, '<=', (60*60), 'Checking Retry-After is smaller than 60 minutes');
}
};
}
# Testing bursts with an additional fudge factor making clients wait for
# time not equal to the rate-limiter's true limit (but imposed by a higher level)
{
my $handler = builder {
enable "Throttle::Second",
max => 2,
retry_after_addition => 2,
backend => Plack::Middleware::Throttle::Backend::Hash->new();
sub {
[ '200', [ 'Content-Type' => 'text/html' ], ['hello world'] ]
};
};
test_psgi
app => $handler,
client => sub {
my $cb = shift;
note 'Not limited';
success($cb) for 1..2;
{
note 'Now expecting failure';
my $res = get($cb);
my $retry_after = $res->header('Retry-After');
cmp_ok($retry_after, '<', 3, 'Checking Retry-After is smaller than 3 seconds');
cmp_ok($retry_after, '>', 1, 'Checking Retry-After is larger than a second');
}
};
}
done_testing();
\ No newline at end of file
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment