File Coverage

blib/lib/Data/YUID/Client.pm

Criterion	Covered	Total	%
statement	27	107	25.2
branch	0	48	0.0
condition	0	29	0.0
subroutine	9	16	56.2
pod	2	5	40.0
total	38	205	18.5

# $Id$

package Data::YUID::Client;
use strict;

use fields qw( servers select_timeout hosts connect_timeout );
use Carp;
use Errno qw( EINPROGRESS EWOULDBLOCK EISCONN );
use IO::Socket::INET;
use Socket qw( MSG_NOSIGNAL );
use URI::Escape ();

use constant DEFAULT_PORT => 9001;
use constant DEFAULT_TIMEOUT => {
    connect => 1.0, # seconds
    select  => 1.0, # seconds
};

our $FLAG_NOSIGNAL = 0;
eval { $FLAG_NOSIGNAL = MSG_NOSIGNAL };

my $Active_Sock;
my $Socket_Used_Count = 0; #yuids served per socket before we refresh it, to force rebalancing
my $Max_Socket_Reused = 1000 + rand(1000);

sub new {
    my Data::YUID::Client $client = shift;
    my %args = @_;
    $client = fields::new($client) unless ref $client;

    croak "servers must be an arrayref if specified"
        unless !exists $args{servers} || ref $args{servers} eq 'ARRAY';
    $client->{servers} = $args{servers} || [];
    for (qw(select_timeout connect_timeout)) {
        my $value = exists $args{$_} ? $args{$_} : DEFAULT_TIMEOUT->{$_};
        $client->{$_} = $value;
    }
    $client->{hosts} = $args{servers};

    $client;
}

sub _oneline {
    my Data::YUID::Client $client = shift;
    my ($sock, $line) = @_;
    my $res;
    my ($ret, $offset) = (undef, 0);

    # state: 0 - writing, 1 - reading, 2 - done
    my $state = defined $line ? 0 : 1;

    # the bitsets for select
    my ($rin, $rout, $win, $wout);
    my $nfound;

    my $copy_state = -1;
    local $SIG{'PIPE'} = "IGNORE" unless $FLAG_NOSIGNAL;

    # the select loop
    while(1) {
        if ($copy_state!=$state) {
            last if $state==2;
            ($rin, $win) = ('', '');
            vec($rin, fileno($sock), 1) = 1 if $state==1;
            vec($win, fileno($sock), 1) = 1 if $state==0;
            $copy_state = $state;
        }
        $nfound = select($rout=$rin, $wout=$win, undef,
                         $client->{select_timeout});
        last unless $nfound;

        if (vec($wout, fileno($sock), 1)) {
            $res = send($sock, $line, $FLAG_NOSIGNAL);
            next
                if not defined $res and $!==EWOULDBLOCK;
            unless (($res || 0) > 0) {
                _close_sock($sock);
                return undef;
            }
            if ($res == length($line)) { # all sent
                $state = 1;
            } else { # we only succeeded in sending some of it
                substr($line, 0, $res, ''); # delete the part we sent
            }
        }

        if (vec($rout, fileno($sock), 1)) {
            $res = sysread($sock, $ret, 255, $offset);
            next
                if !defined($res) and $!==EWOULDBLOCK;
            if ($res == 0) { # catches 0=conn closed or undef=error
                _close_sock($sock);
                return undef;
            }
            $offset += $res;
            if (rindex($ret, "\r\n") + 2 == length($ret)) {
                $state = 2;
            }
        }
    }

    unless ($state == 2) {
        _close_sock($sock);
        return undef;
    }

    return $ret;
}

sub get_id {
    my Data::YUID::Client $client = shift;
    my($ns) = @_;
    my $id;
    while (!$id && (my $sock = $client->get_sock)) {
        my $cmd = sprintf "getid ns=%s\r\n", URI::Escape::uri_escape($ns || '');
        my $res = $client->_oneline($sock, $cmd) or next;
        ($id) = $res =~ /^ok\s+id=(\d+)/i;
    }
    $id;
}

sub _close_sock {
    my($sock) = @_;
    undef $Active_Sock
        if ($Active_Sock && fileno($sock) == fileno($Active_Sock));
    close $sock;
}

#class method, for request cleanup
sub close_active_sock {
    close $Active_Sock;
}

sub connect_to_server {
    my Data::YUID::Client $client = shift;
    my($host) = @_;
    my($ip, $port) = split /:/, $host;
    $port ||= DEFAULT_PORT;
    my $sock = IO::Socket::INET->new(
            PeerAddr        => $ip,
            PeerPort        => $port,
            Proto           => 'tcp',
            Type            => SOCK_STREAM,
            ReuseAddr       => 1,
            Blocking        => 0,
            Timeout         => $client->{connect_timeout},
        ) or return;
    $sock;
}

sub get_sock {
    my Data::YUID::Client $client = shift;

    ++$Socket_Used_Count;

    if ($Active_Sock && $Active_Sock->connected && ($Socket_Used_Count % $Max_Socket_Reused != 0)) {
        return $Active_Sock;
    }else{
        my $hosts = $client->{hosts};
        my @try_hosts = @$hosts;
        my $sock;
        while (@try_hosts){
            my $host = splice(@try_hosts, int(rand(@try_hosts)), 1);
            $sock = $client->connect_to_server($host);
            last if $sock;
        }
        return $Active_Sock = $sock;
    }
}

1;
__END__