File Coverage

blib/lib/Schedule/AdaptiveThrottler.pm

Criterion	Covered	Total	%
statement	91	95	95.7
branch	45	68	66.1
condition	11	24	45.8
subroutine	9	9	100.0
pod	3	3	100.0
total	159	199	79.9

package Schedule::AdaptiveThrottler;

use warnings;
use strict;

our $VERSION = '0.06';
our $DEBUG   = 0;
our $QUIET   = 0;

use Scalar::Util qw(reftype blessed);
use Digest::MD5 qw(md5_hex);

require Exporter;
our @ISA       = qw(Exporter);
our @EXPORT_OK = qw(&authorize &set_client);
our @EXPORT    = qw(
    SCHED_ADAPTHROTTLE_AUTHORIZED
    SCHED_ADAPTHROTTLE_BLOCKED
);
our %EXPORT_TAGS = ( ALL => [ @EXPORT_OK, @EXPORT ] );

use constant SCHED_ADAPTHROTTLE_BLOCKED    => 0;
use constant SCHED_ADAPTHROTTLE_AUTHORIZED => 1;

my $memcached_client;    # used for the non-OO form

sub set_client {
    my $client = pop;     # keep ordered
    my $self   = shift;
    die "Invalid storage client object\n"
        if ( !blessed($client)
        || !$client->can('set')
        || !$client->can('get') );
    return $memcached_client = $client if !blessed $self;    # non-OO form
    return $self->{memcached_client} = $client;              # Guess what? OO-form
}

# OO-style
sub new {
    my $class = shift;
    my $params;
    if ( @_ == 1 ) {
        if ( !blessed $_[0] ) {
            $params = shift;
        }
        else {
            $params->{memcached_client} = shift;
        }
    }
    else {
        $params = {@_};
    }
    my $self = bless $params, $class;
    $self->set_client( $params->{memcached_client} )
        if $params->{memcached_client};
    return $self;
}

sub authorize {

    my %params;
    my $self;

    # Call it as a method or a sub, with a hash or a hashref,
    # as a class or an instance
    if ( @_ < 3 ) {
        %params = %{ pop() };
        $self   = shift;
    }
    else {
        $self = shift if ( @_ % 2 );
        %params = @_;
    }
    my $cur_memcached_client
        = blessed $self ? $self->{memcached_client} : $memcached_client;    # can it get uglier?

    my $frozen_time = time;
    my %conditions;

    # Check the conditions

    my $condition_type;
    for my $condition_type_tmp (qw(all either)) {

        if ( exists $params{$condition_type_tmp}
            && reftype $params{$condition_type_tmp} eq 'HASH' )
        {
            %conditions = %{ $params{$condition_type_tmp} };
            die "Conditions improperly defined (must be a hashref)"
                if !%conditions;

            # Check the parameters

            for my $condition_params ( values %conditions ) {
                for my $condition_param_key (qw{max ttl message value}) {

                    # message & value are strings (or just anything for 'value'), the rest are integers
                    die
                        "Condition parameter $condition_param_key is '$condition_params->{$condition_param_key}'"
                        if !$condition_params->{$condition_param_key};
                    die "Condition parameter $condition_param_key must be positive integer"
                        if ( ( $condition_param_key eq 'max' || $condition_param_key eq 'ttl' )
                        && $condition_params->{$condition_param_key} !~ /^[1-9][0-9]*$/ );
                }
            }
            $condition_type = $condition_type_tmp;
            last;    # process either 'all' or 'either', not both
        }
    }
    die "No conditions defined"
        if !$condition_type;

    # if lockout is defined, use the 'lockout/ban' scheme.  if not, we'll use a
    # bucket algorithm
    my $lockout = $params{lockout};
    die "'Lockout' parameter must be positive integer"
        if ( defined $lockout && $lockout !~ /^[1-9][0-9]*$/ );

    my $identifier = $params{identifier};
    die "'Identifier' should be a non-empty string"
        if ( !defined $identifier || length($identifier) < 1 );

    # Loop on the conditions. For 'either', we need to find one that is not yet
    # satisfied, for 'all' we need to find lockouts for all of them

    # Make the memcached keys a identifier + key name + value
    # TODO: Retrieve the records in 1 operation with get_multi
    # @conditions_names = sort keys %conditions;
    # @keys = map { $_ . '#' . $conditions{$_}->{value} } @conditions_names ) {

    my ( $conditions_ok, $conditions_unknown ) = ( 0, 0 );
    my $messages_notok = [];

    while ( my ( $condition_name, $condition ) = each %conditions ) {
        my $memcached_key = $identifier . '#' . $condition_name . '#' . $condition->{value};
        $memcached_key = md5_hex($memcached_key) if length $memcached_key > 249;

        my $record = $cur_memcached_client->get($memcached_key);

        if ( defined $record ) {

            # Do we have a 'block' value in the record, in which case we return
            # a message indicating so. The 'block' record will be automatically
            # removed from memcached at the object's expiry time, so don't
            # touch it.
            if ( $record eq 'block' ) {
                push @$messages_notok, $condition->{message};
                print STDERR "Access already blocked by " . __PACKAGE__ . "\n"
                    if $DEBUG;
            }

            # the object in memcached is a list of timestamps, and nothing else.
            elsif ( reftype $record eq 'ARRAY' ) {
                print STDERR "Current timestamps in \$record:  " . join( '|', @$record ) . "\n"

                    if $DEBUG;
                print STDERR "Current frozen time: $frozen_time" . "\n"
                    if $DEBUG;

                # cleanup the records (remove expired timestamps). This is
                # where it all happens, giving us this "magic sliding time
                # window".
                @$record = grep { $_ > $frozen_time } @$record;
                print STDERR "Currently unexpired timestamps in \$record:  "
                    . join( '|', @$record ) . "\n"

                    if $DEBUG;

                # Since we are about to add a record, if we already have the
                # max number of records, set to blocked. If no lockout time
                # specified, use the bucket algorithm: deny access, but do not
                # update the record. The expired timestamps will be evicted in
                # due time (next access, possibly), giving us more tokens.
                print STDERR "Maximum is "
                    . $condition->{max}
                    . " and current number of timestamps is "
                    . @$record . "\n"

                    if $DEBUG;
                if ( @$record >= $condition->{max} ) {
                    print STDERR "Maximum reached" . "\n"
                        if $DEBUG;
                    if ($lockout) {
                        print STDERR "Setting a timed lock" . "\n"
                            if $DEBUG;
                        $cur_memcached_client->set( $memcached_key, 'block', $lockout );
                    }
                    push @$messages_notok, $condition->{message};
                }

                # Add a timestamp to the list. This is NOT the current
                # timestamp, but a timestamp in the future (a TTL record),
                # which allowws for easy filtering by the grep above. And set
                # the memcached record expiration time at the most recent TTL
                # of the list (for automatic cleanup: the object will be
                # discarded from memcached automatically if it is not updated
                # before the longest TTL)
                else {
                    print STDERR "Adding a timestamp to the list" . "\n"
                        if $DEBUG;
                    push @$record, $frozen_time + $condition->{ttl};
                    $cur_memcached_client->set( $memcached_key, $record, $condition->{ttl} );
                    $conditions_ok++;
                }
            }
            else {    # This should not happen, but catch it if it does.
                $conditions_unknown++;
            }
        }

        # $record is undef, either not accessible, or not yet created
        else {
            print STDERR "No record found, creating a new one" . "\n"
                if $DEBUG;
            my $ret
                = $cur_memcached_client->set( $memcached_key,
                [ $condition->{ttl} + $frozen_time ],
                $condition->{ttl} );
            $conditions_ok++;
        }
    }

    if ( $conditions_unknown && !$QUIET ) {
        warn "Unknown conditions count is over 0, this should not happen";
        print STDERR "Current conditions hash: " . Dumper( \%conditions ) . "\n";
    }

    # If logic was 'either', 1 'notok' or more should block
    # If logic was 'all', we should have 0 'ok' to block
    # TODO: re-work the variable names because the explanation above is a bit
    # tricky although the logic is correct :(
    if ( $condition_type eq 'either' ) {
        return ( @$messages_notok > 0 )
            ? ( SCHED_ADAPTHROTTLE_BLOCKED, $messages_notok )
            : ( SCHED_ADAPTHROTTLE_AUTHORIZED, undef );
    }
    else {    # condition is 'all'
        return ( $conditions_ok == 0 )
            ? ( SCHED_ADAPTHROTTLE_BLOCKED, $messages_notok )
            : ( SCHED_ADAPTHROTTLE_AUTHORIZED, undef );
    }
}

1;

__END__