File Coverage

blib/lib/AnyEvent/Beanstalk.pm

Criterion	Covered	Total	%
statement	33	293	11.2
branch	0	150	0.0
condition	0	162	0.0
subroutine	11	55	20.0
pod	30	33	90.9
total	74	693	10.6

package AnyEvent::Beanstalk;
{
  $AnyEvent::Beanstalk::VERSION = '1.123530';
}

use strict;
use warnings;

use constant DEBUG => $ENV{AE_BEANSTALK_DEBUG};
use Scalar::Util qw(blessed);

use AnyEvent;
use AnyEvent::Handle;
use AnyEvent::Socket;
use AnyEvent::Beanstalk::Job;
use AnyEvent::Beanstalk::Stats;

use base qw(Class::Accessor::Fast);

__PACKAGE__->mk_accessors(
  qw< decoder delay encoder on_error on_connect priority server socket ttr >    ##
);

my $YAML_CLASS = do {
  local ($SIG{__DIE__}, $SIG{__WARN__});
      eval { require YAML::XS }   ? 'YAML::XS'
    : eval { require YAML::Syck } ? 'YAML::Syck'
    : eval { require YAML }       ? 'YAML'
    :                               die $@;
};
my $YAML_LOAD = $YAML_CLASS->can('Load');
my $YAML_DUMP = $YAML_CLASS->can('Dump');


sub new {
  my $proto = shift;
  my %arg   = @_;

  bless(
    { delay      => $arg{delay}      || 0,
      ttr        => $arg{ttr}        || 120,
      priority   => $arg{priority}   || 10_000,
      encoder    => $arg{encoder}    || $YAML_DUMP,
      decoder    => $arg{decoder}    || $YAML_LOAD,
      server     => $arg{server}     || undef,
      debug      => $arg{debug}      || 0,
      on_error   => $arg{on_error}   || undef,
      on_connect => $arg{on_connect} || undef,
    },
    ref($proto) || $proto
  );
}


sub run_cmd {
  my $self = shift;

  $self->{_cmd_cb} or return $self->connect(@_);
  $self->{_cmd_cb}->(@_);
}


sub quit { shift->disconnect }
sub reserve_pending { shift->{_reserve_pending} || 0 }


sub disconnect {
  my $self = shift;
  my $condvar = delete $self->{_condvar};
  delete @{$self}{grep {/^_[a-z]/} keys %$self};
  if ($condvar) {
    $_->send for values %$condvar;
  }
  return;
}


sub _error {
  my $self = shift;
  $self->disconnect;
  ($self->on_error || sub { die @_ })->(@_);
}


sub reconnect {
  my $self = shift;

  my $using = $self->{__using} || 'default';
  $self->use(
    $using,
    sub {
      $self->_error("Can't use '$using'") unless @_ and $_[0] eq 'USING';
    }
  );

  my $watching = $self->{__watching} || {default => 1};
  $self->watch_only(
    keys %$watching,
    sub {
      $self->_error("Error watching tubes") unless @_ and $_[0] eq 'WATCHING';
    }
  );
}

my %EXPECT = qw(
  put                  INSERTED
  use                  USING
  reserve              RESERVED
  reserve-with-timeout RESERVED
  delete               DELETED
  release              RELEASED
  bury                 BURIED
  touch                TOUCHED
  watch                WATCHING
  ignore               WATCHING
  peek                 FOUND
  peek-ready           FOUND
  peek-delayed         FOUND
  peek-buried          FOUND
  kick                 KICKED
  kick-job             KICKED
  stats-job            OK
  stats-tube           OK
  stats                OK
  list-tubes           OK
  list-tube-used       USING
  list-tubes-watched   OK
  pause-tube           PAUSED
);

sub connect {
  my $self = shift;

  my $cv;
  if (@_) {
    $cv = AE::cv;
    $self->{_condvar}{$cv} = $cv;
    push @{$self->{_connect_queue}}, [@_, $cv];
  }

  return $cv if $self->{_sock};

  my ($host, $port) = parse_hostport($self->server || '127.0.0.1', 11300);
  $self->{_sock} = tcp_connect $host, $port, sub {
    $self->server("$host\:$port");
    my $fh = shift
      or return $self->_error("Can't connect to beanstalk server: $!");

    $self->{__using} = 'default';
    $self->{__watching} = {default => 1};
    my $on_connect = $self->on_connect;
    $on_connect->() if $on_connect;
    $self->{_socket} = $fh;

    my $hd = AnyEvent::Handle->new(
      fh       => $fh,
      on_error => sub { $_[0]->destroy; $self->_error($_[2]) },
      on_eof   => sub { $_[0]->destroy; $self->_error("EOF") },
    );

    $self->{_cmd_cb} = sub {
      my $command = lc shift;

      my ($cv, $cb);
      {
        no warnings;
        $cv = pop if @_ && blessed($_[-1]) eq 'AnyEvent::CondVar';
        $cb = pop if @_ && ref $_[-1]      eq 'CODE';
      }

      my $value = $command eq 'put' ? pop(@_) . "\015\012" : '';
      my @argv  = @_;
      my $send  = join(" ", $command, @argv) . "\015\012" . $value;

      warn "Sending [$send]\n" if DEBUG;

      $cv ||= AE::cv;
      $self->{_condvar}{$cv} = $cv;
      $cv->cb(
        sub {
          my $cv  = shift;
          my @res = $cv->recv;
          $cb->(@res);
        }
      ) if $cb;

      $self->{_reserve_pending}++ if $command =~ /^reserve/;

      $hd->push_write($send);
      $hd->push_read(
        line => sub {
          my ($hd, $result) = @_;
          warn "got line <$result> for command [$send]\n" if DEBUG;
          my @resp = split(/\s+/, $result);
          my $resp = uc shift @resp;

          $self->{_reserve_pending}-- if $command =~ /^reserve/;

          unless ($resp eq $EXPECT{$command}) {
                delete $self->{_condvar}{$cv};
                $cv->send(undef, $result);
                return;
          }

          if ($resp =~ /^ (?: RESERVED | FOUND ) $/x) {
            my ($id, $bytes) = @resp;
            $hd->unshift_read(
              chunk => $bytes + 2,
              sub {
                my ($hd, $chunk) = @_;
                my $job = AnyEvent::Beanstalk::Job->new(
                  id     => $id,
                  client => $self,
                  data   => substr($chunk, 0, -2),
                );
                delete $self->{_condvar}{$cv};
                $cv->send($job, $result);
              }
            );
          }
          elsif ($resp =~ /^ (?: INSERTED | BURIED ) $/x) {
            my $id  = shift @resp;
            my $job = AnyEvent::Beanstalk::Job->new(
              id     => $id,
              client => $self,
              data   => substr($value, 0, -2),
              buried => $resp eq 'BURIED' ? 1 : 0,
            );
            delete $self->{_condvar}{$cv};
            $cv->send($job, $result);
          }
          elsif ($resp eq 'OK') {
            my $bytes = shift @resp;
            $hd->unshift_read(
              chunk => $bytes + 2,
              sub {
                my ($hd, $chunk) = @_;
                warn "got '$chunk'\n" if DEBUG;
                my $yaml = $YAML_LOAD->($chunk);
                delete $self->{_condvar}{$cv};
                $yaml = AnyEvent::Beanstalk::Stats->new($yaml) if $command =~ /^stats/;
                $cv->send($yaml,$result);
              }
            );
          }
          elsif($resp =~ /^ (?: RELEASED | TOUCHED ) $/x) {
            my ($id, $pri, $delay) = @argv;
            delete $self->{_condvar}{$cv};
            my $job = AnyEvent::Beanstalk::Job->new(
              id     => $id,
              client => $self,
              ( $command eq 'release'
                ? (
                  priority => $pri,
                  delay    => $delay,
                  )
                : ()
              )
            );
            $cv->send($job, $result);
          }
          elsif($resp =~ /^ (?: USING | WATCHING ) $/x) {
            delete $self->{_condvar}{$cv};
            my $retval = shift @resp;
            $cv->send($retval, $result);
          }
          else {
            delete $self->{_condvar}{$cv};
            $cv->send(1, $result);
          }
        }
      );

      return $cv;
    };

    for my $queue (@{$self->{_connect_queue} || []}) {
      $self->{_cmd_cb}->(@$queue);
    }
    delete $self->{_connect_queue};
  };

  return $cv;
}


sub watch_only {
  my $self = shift;
  my $cb   = pop if @_ and ref($_[-1]) eq 'CODE';
  my $cv   = AE::cv;
  $self->{_condvar}{$cv} = $cv;

  $cv->cb(
    sub {
      my $cv  = shift;
      my @res = $cv->recv;
      $cb->(@res);
    }
  ) if $cb;

  unless (@_) {
    delete $self->{_condvar}{$cv};
    $cv->send('NOT_IGNORED');
    return $cv;
  }

  my %tubes = map { ($_ => 1) } @_;
  my $done = sub {
    %tubes = ();
    delete $self->{_condvar}{$cv};
    $cv->send(@_);
  };
  $self->list_tubes_watched(
    sub {
      my ($tubes,$r) = @_;
      return $done->(@_) unless $r and $r =~ /^OK\b/;
      foreach my $t (@$tubes) {
        $tubes{$t} = 0 unless delete $tubes{$t};
      }
      $done->() if !keys %tubes;
      foreach my $t (sort { $tubes{$b} <=> $tubes{$a} } keys %tubes) {
        my $cmd = $tubes{$t} ? 'watch' : 'ignore';
        $self->run_cmd(
          $cmd, $t,
          sub {
            my ($r) = @_;
            return unless keys %tubes;
            delete $tubes{$t};
            return $done->(@_)
              if !@_
                or $r ne 'WATCHING'
                or !keys %tubes;
          }
        );
      }
    }
  );

  $cv;
}


sub put {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $opt  = shift || {};

  my $pri   = exists $opt->{priority} ? $opt->{priority} : $self->priority;
  my $ttr   = exists $opt->{ttr}      ? $opt->{ttr}      : $self->ttr;
  my $delay = exists $opt->{delay}    ? $opt->{delay}    : $self->delay;
  my $data =
      exists $opt->{data}   ? $opt->{data}
    : exists $opt->{encode} ? $self->encoder->($opt->{encode})
    :                         '';

  $pri   = int($pri   || 0) || 1;
  $ttr   = int($ttr   || 0) || 1;
  $delay = int($delay || 0) || 0;
  $data = '' unless defined $data;

  utf8::encode($data) if utf8::is_utf8($data);    # need bytes

  $self->run_cmd('put' => $pri, $delay, $ttr, length($data), $data, @cb);
}


sub stats {
  my $self = shift;
  my @cb = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  $self->run_cmd('stats' => @cb);
}


sub stats_tube {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $tube = shift;
  $self->run_cmd('stats-tube' => $tube, @cb);
}


sub stats_job {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift || 0;
  $self->run_cmd('stats-job' => $id, @cb);
}


sub kick {
  my $self  = shift;
  my @cb    = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $bound = shift || 1;
  $self->run_cmd('kick' => $bound, @cb);
}


sub kick_job {
  my $self  = shift;
  my @cb    = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id    = shift || 0;
  $self->run_cmd('kick-job' => $id, @cb);
}

sub use {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $tube = shift;
  $self->run_cmd(
    'use' => $tube,
    sub {
      $self->{__using} = $_[0] if @_ and $_[1] =~ /^USING/;
      $cb[0]->(@_) if @cb;
    }
  );
}


sub reserve {
  my $self    = shift;
  my @cb      = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $timeout = shift;

  my @cmd = defined($timeout) ? ('reserve-with-timeout' => $timeout) : "reserve";
  $self->run_cmd(@cmd, @cb);
}


sub delete {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift || 0;
  $self->run_cmd('delete' => $id, @cb);
}


sub touch {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift || 0;
  $self->run_cmd('touch' => $id, @cb);
}


sub release {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift || 0;
  my $opt  = shift || {};

  my $pri   = exists $opt->{priority} ? $opt->{priority} : $self->priority;
  my $delay = exists $opt->{delay}    ? $opt->{delay}    : $self->delay;

  $self->run_cmd('release' => $id, $pri, $delay, @cb);
}


sub bury {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift;
  my $opt  = shift || {};

  my $pri = exists $opt->{priority} ? $opt->{priority} : $self->priority;

  $self->run_cmd('bury' => $id, $pri, @cb);
}


sub watch {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $tube = shift;

  $self->run_cmd(
    'watch' => $tube,
    sub {
      $self->{__watching}{$tube} = 1 if @_ and $_[1] =~ /^WATCHING/;
      $cb[0]->(@_) if @cb;
    }
  );
}


sub ignore {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $tube = shift;

  $self->run_cmd(
    'ignore' => $tube,
    sub {
      delete $self->{__watching}{$tube} if @_ and $_[1] =~ /^WATCHING/;
      $cb[0]->(@_) if @cb;
    }
  );
}


sub peek {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift;

  $self->run_cmd('peek' => $id, @cb);
}


sub peek_ready {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift;

  $self->run_cmd('peek-ready' => @cb);
}


sub peek_delayed {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift;

  $self->run_cmd('peek-delayed' => @cb);
}


sub peek_buried {
  my $self = shift;
  my @cb   = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $id   = shift;

  $self->run_cmd('peek-buried' => @cb);
}


sub list_tubes {
  my $self = shift;
  my @cb = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  $self->run_cmd('list-tubes' => @cb);
}


sub list_tube_used {
  my $self = shift;
  my @cb = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  $self->run_cmd('list-tube-used' => @cb);
}


sub list_tubes_watched {
  my $self = shift;
  my @cb = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  $self->run_cmd('list-tubes-watched' => @cb);
}


sub pause_tube {
  my $self  = shift;
  my @cb    = (@_ and ref($_[-1]) eq 'CODE') ? splice(@_, -1) : ();
  my $tube  = shift;
  my $delay = shift || 0;
  $self->run_cmd('pause-tube' => $tube, $delay, @cb);
}


sub watching {
  my $self = shift;
  return unless $self->{_sock};
  my $watching = $self->{__watching} or return;
  return keys %$watching;
}


sub using {
  my $self = shift;
  return $self->{__using};
}


sub sync {
  my $self = shift;

  while ($self->{_condvar} and my ($cv) = values %{$self->{_condvar}}) {
    $cv->recv;
  }
}

1;
__END__