File Coverage

blib/lib/Sub/Throttler.pm

Criterion	Covered	Total	%
statement	293	308	95.1
branch	85	96	88.5
condition	19	33	57.5
subroutine	44	45	97.7
pod	10	10	100.0
total	451	492	91.6

package Sub::Throttler;
use 5.010001;
use warnings;
use strict;
use utf8;
use Carp;

our $VERSION = 'v0.2.2';

use Perl6::Export::Attrs;
use Scalar::Util qw( weaken refaddr blessed );


use constant CALLER_SUBROUTINE  => 3;

my (@Throttles, @Tasks, @AsapTasks, %Running);
my $IN_flush = 0;
my $IN_flush_recursion = 0;
my $IN_flush_ignore_recursion = 0;
my $IN_flush_pending = 0;


sub done_cb :Export {
    my ($done, $cb_or_obj_or_class, @p) = @_;
    if (ref $cb_or_obj_or_class eq 'CODE') {
        my $cb = $cb_or_obj_or_class;
        return sub { $done->(); $cb->(@p, @_) };
    }
    elsif (blessed($cb_or_obj_or_class)) {
        my $obj = $cb_or_obj_or_class;
        weaken($obj);
        my $method = shift @p;
        croak 'second param must be $method'
            if !$method || (ref $method && ref $method ne 'CODE');
        return sub { $done->(); $obj && $obj->$method(@p, @_) };
    }
    elsif (defined $cb_or_obj_or_class && !ref $cb_or_obj_or_class) {
        my $class = $cb_or_obj_or_class;
        my $method = shift @p;
        croak 'second param must be $method'
            if !$method || (ref $method && ref $method ne 'CODE');
        return sub { $done->(); $class->$method(@p, @_) };
    }
    else {
        croak 'first param must be $cb or $obj or $class';
    }
}

sub throttle_add :Export {
    my ($throttle, $target) = @_;
    croak 'require 2 params' if 2 != @_;
    croak 'throttle must be an object' if !ref $throttle;
    croak 'target must be CODE' if ref $target ne 'CODE';
    push @Throttles, [$throttle, $target];
    return $throttle;
}

sub throttle_del :Export {
    my ($throttle) = @_;
    @Throttles = grep { $throttle && $_->[0] != $throttle } @Throttles;
    throttle_flush();
    return;
}

sub throttle_flush :Export {
    if ($IN_flush) {
        if (!$IN_flush_ignore_recursion) {
            $IN_flush_recursion = 1;
        }
        return;
    }
    $IN_flush = 1;
    $IN_flush_recursion = 0;
    $IN_flush_ignore_recursion = 0;
    $IN_flush_pending = 0;

    for my $tasks (\@AsapTasks, \@Tasks) {
        my @tasks = @{$tasks};
        @{$tasks} = ();
        my @delayed;
TASK:
        for my $task (@tasks) {
            my ($done, $name, $this, $code, @params) = @{$task};
            my $id = refaddr $done;
            if (!defined $this) {
                $done->();  # release $done
                next;
            }
            my %acquired;
            for (@Throttles) {
                my ($throttle, $target) = @{$_};
                my $resources = $target->($this, $name, @params);
                next if !defined $resources;
                die "Sub::Throttler: target returns not a HASHREF: $resources\n"
                    if ref $resources ne 'HASH';
                next if !keys %{$resources};
                my $acquired = 0;
                while (my ($key, $quantity) = each %{$resources}) {
                    die "Sub::Throttler: target returns bad quantity for '$key': $quantity\n"
                        if ref $quantity;
                    if ($throttle->try_acquire($id, $key, $quantity)) {
                        $acquired++;
                    }
                    else {
                        last;
                    }
                }
                if ($acquired == keys %{$resources}) {
                    $acquired{$throttle} = $throttle;
                }
                else {
                    $IN_flush_ignore_recursion = 1;
                    if ($acquired) {
                        $throttle->release_unused($id);
                    }
                    for (values %acquired) {
                        $_->release_unused($id);
                    }
                    $IN_flush_ignore_recursion = 0;
                    push @delayed, $task;
                    next TASK;
                }
            }
            $Running{$id} = [values %acquired];
            _run_task($this, $code, $done, @params);
        }
        @{$tasks} = (@delayed, @{$tasks}); # while _run_task() new tasks may be added
    }

    $IN_flush = 0;
    goto &throttle_flush if $IN_flush_recursion;
    return;
}

sub throttle_it :Export {
    return _it(0, 0, @_);
}

sub throttle_it_asap :Export {
    return _it(0, 1, @_);
}

sub throttle_it_sync :Export {
    return _it(1, 0, @_);
}

sub throttle_me :Export {
    return _me(\@Tasks, \@_);
}

sub throttle_me_asap :Export {
    return _me(\@AsapTasks, \@_);
}

sub throttle_me_sync :Export {
    my ($done, $failed);

    my ($this, @params);
    my $func = (caller 1)[CALLER_SUBROUTINE];
    croak 'impossible to throttle anonymous function' if !defined &{$func};
    my ($pkg, $name) = $func =~ /\A(.*)::(.*)\z/ms;
    my $is_method = eval { local $SIG{__DIE__}; $_[0]->isa($pkg) };
    if ($is_method) {
        ($this, @params) = @_;
        $done = Sub::Throttler::__done->new($this.q{->}.$name);
    }
    else {
        ($this, @params) = (q{}, @_);
        $name = $func;
        $done = Sub::Throttler::__done->new($func);
    }

    my @old = ($IN_flush, $IN_flush_ignore_recursion);
    ($IN_flush, $IN_flush_ignore_recursion) = (1, 1);
    my $id = refaddr $done;
ACQUIRE_ALL:
    {
        my %acquired;
        for (@Throttles) {
            my ($throttle, $target) = @{$_};
            my $resources = $target->($this, $name, @params);
            next if !defined $resources;
            die "Sub::Throttler: target returns not a HASHREF: $resources\n"
                if ref $resources ne 'HASH';
            next if !keys %{$resources};
            while (my ($key, $quantity) = each %{$resources}) {
                die "Sub::Throttler: target returns bad quantity for '$key': $quantity\n"
                    if ref $quantity;
                if ($throttle->try_acquire($id, $key, $quantity)) {
                    $acquired{$throttle} = $throttle;
                }
                else {
                    eval { ## no critic (RequireCheckingReturnValueOfEval)
                        local $SIG{__DIE__};
                        $throttle->acquire($id, $key, $quantity);
                        $acquired{$throttle} = $throttle;
                    };
                    $failed = $@;
                    for (values %acquired) {
                        $_->release_unused($id);
                    }
                    if ($failed) {
                        last ACQUIRE_ALL;
                    } else {
                        redo ACQUIRE_ALL;
                    }
                }
            }
        }
        $Running{$id} = [values %acquired];
    }
    ($IN_flush, $IN_flush_ignore_recursion) = @old;
    # while waiting for resources needed for this sync call some resources
    # needed for queued async calls may be released, but in this case
    # throttle_flush() wasn't called because it was blocked, so let's
    # ensure it will be called no late than this sync call will $done->()
    $IN_flush_pending = 1;

    if ($failed) {
        croak $failed;
    } else {
        return $done;
    }
}

sub _done { ## no critic (ProhibitUnusedPrivateSubroutines)
    my ($id, $is_used) = @_;
    $is_used ||= 1 == @_;
    for my $throttle (@{ delete $Running{$id} // [] }) {
        if ($is_used) {
            $throttle->release($id);
        } else {
            $throttle->release_unused($id);
        }
    }
    if ($IN_flush_pending) {
        throttle_flush();
    }
    return;
}

sub _it {
    my ($is_sync, $is_asap, $func) = @_;
    croak 'require function name' if !$func || ref $func;
    if ($func !~ /::/ms) {
        $func = caller(1) . q{::} . $func;
    }
    croak 'no such function: '.$func if !defined &{$func};
    my $orig = \&{$func};
    ## no critic (ProhibitNoWarnings ProhibitStringyEval RequireCheckingReturnValueOfEval ProhibitImplicitNewlines RequireCarping)
    no warnings 'redefine';
    eval 'sub '.$func.' {
        if (!'.$is_sync.' && @_ && ref $_[-1] eq "CODE") {
            my $done = &throttle_me'.($is_asap ? '_asap' : q{}).' || return;
            my $cb = pop;
            $orig->(@_, done_cb($done, $cb));
            return;
        } elsif (wantarray) {
            my $done = &throttle_me_sync;
            my @res = &$orig;
            $done->();
            return @res;
        } else {
            my $done = &throttle_me_sync;
            my $res = &$orig;
            $done->();
            return $res;
        }
    }; 1' or die $@;
    return $orig;
}

sub _me {
    my ($queue, $args) = @_;
    for (0, 1) {
        if (ref $args->[$_] eq 'Sub::Throttler::__done') {
            return splice @{$args}, $_, 1;
        }
    }
    my $func = (caller 2)[CALLER_SUBROUTINE];
    croak 'impossible to throttle anonymous function' if !defined &{$func};
    my $code = \&{$func};
    my ($pkg, $name) = $func =~ /\A(.*)::(.*)\z/ms;
    my $is_method = eval { local $SIG{__DIE__}; $args->[0]->isa($pkg) };
    if ($is_method) {
        my $self = shift @{$args};
        my $done = Sub::Throttler::__done->new($self.q{->}.$name);
        push @{$queue}, [$done, $name, $self, $code, @{$args}];
        if (ref $self) {
            weaken $queue->[-1][2];
        }
    }
    else {
        my $done = Sub::Throttler::__done->new($func);
        push @{$queue}, [$done, $func, q{}, $code, @{$args}];
    }
    throttle_flush();
    return;
}

# should be used only from tests
sub _reset { ## no critic (ProhibitUnusedPrivateSubroutines)
    $IN_flush = 0;
    $IN_flush_recursion = 0;
    $IN_flush_ignore_recursion = 0;
    @Throttles = @Tasks = @AsapTasks = %Running = ();
    return;
}

sub _run_task {
    my ($this, $code, $done, @params) = @_;
    no strict 'refs';
    if ($this) {
        $this->$code($done, @params);
    } else {
        $code->($done, @params);
    }
    return;
}


package Sub::Throttler::__done;    ## no critic (ProhibitMultiplePackages)
use Carp;

use Scalar::Util qw( refaddr );

my (%Check, %Name);

sub new {
    my (undef, $name) = @_;
    my $id;
    my $done = bless sub {
        if ($Check{$id}) {
            croak "Sub::Throttler: $name: \$done->() already called";
        }
        $Check{$id}=1;
        Sub::Throttler::_done($id, @_); ## no critic(ProtectPrivateSubs)
    }, __PACKAGE__;
    $id = refaddr $done;
    $Name{$id} = $name;
    return $done;
}

sub DESTROY {
    my $done = shift;
    my $id   = refaddr $done;
    my $name = delete $Name{$id};
    if (!delete $Check{$id}) {
        carp "Sub::Throttler: $name: \$done->() was not called";
    }
    return;
}


1; # Magic true value required at end of module
__END__