File Coverage

blib/lib/DBD/Multi.pm

Criterion	Covered	Total	%
statement	215	226	95.1
branch	53	72	73.6
condition	11	12	91.6
subroutine	41	44	93.1
pod	0	1	0.0
total	320	355	90.1

package DBD::Multi;
use strict;

use DBI;
DBI->setup_driver('DBD::Multi');

use vars qw[$VERSION $err $errstr $sqlstate $drh];

$VERSION   = '1.00';

$err       = 0;        # DBI::err
$errstr    = "";       # DBI::errstr
$sqlstate  = "";       # DBI::state
$drh       = undef;

sub driver {
    return $drh if $drh;
    my($class, $attr) = @_;
    $class .= '::dr';

    $drh = DBI::_new_drh($class, {
        Name        => 'Multi',
        Version     => $VERSION,
        Err         => \$DBD::Multi::err,
        Errstr      => \$DBD::Multi::errstr,
        State       => \$DBD::Multi::sqlstate,
        Attribution => 'DBD::Multi, pair Networks Inc.',
    });
    DBD::Multi::db->install_method('multi_do_all');
    return $drh;
}

#######################################################################
package DBD::Multi::dr;
use strict;

$DBD::Multi::dr::imp_data_size = 0;
use DBD::File;

sub DESTROY { shift->STORE(Active => 0) }

sub connect {
    my($drh, $dbname, $user, $auth, $attr) = @_;
    my $dbh = DBI::_new_dbh(
      $drh => {
               Name         => $dbname,
               USER         => $user,
               CURRENT_USER => $user,
              },
    );
    my @dsns =   $attr->{dsns} && ref($attr->{dsns}) eq 'ARRAY'
               ? @{$attr->{dsns}}
               : ();

    if ( $dbname =~ /dsn=(.*)/ ) {
        push @dsns, ( -1, [$1, $user, $auth] );
    }

    my $handler = DBD::Multi::Handler->new({
        dsources => [ @dsns ],
    });
    $handler->failed_max($attr->{failed_max})
      if exists $attr->{failed_max};
    $handler->failed_expire($attr->{failed_expire})
      if exists $attr->{failed_expire};

    $dbh->STORE(_handler => $handler);
    $dbh->STORE(handler => $handler); # temporary
    $drh->{_handler} = $handler;
    $dbh->STORE(Active => 1);
    return $dbh;
}

sub data_sources { shift->FETCH('_handler')->all_sources }

#######################################################################
package DBD::Multi::db;
use strict;

$DBD::Multi::db::imp_data_size = 0;

sub prepare {
    my ($dbh, $statement, @attribs) = @_;

    # create a 'blank' sth
    my ($outer, $sth) = DBI::_new_sth($dbh, { Statement => $statement });

    my $handler = $dbh->FETCH('_handler');
    $sth->STORE(_handler => $handler);

    my $_dbh = $handler->dbh;
    my $_sth;
    until ( $_sth ) {
        $_sth = $_dbh->prepare($statement, @attribs);
        unless ( $_sth ) {
            $handler->dbh_failed;
            $_dbh = $handler->dbh;
        }
    }

    $sth->STORE(NUM_OF_PARAMS => $_sth->FETCH('NUM_OF_PARAMS'));
    $sth->STORE(_dbh => $_dbh);
    $sth->STORE(_sth => $_sth);

    return $outer;
}

sub multi_do_all {
    my $dbh = shift;
    my $handler = $dbh->FETCH('_handler');
    return $handler->multi_do_all( @_ );
}

sub disconnect {
    my ($dbh) = @_;
    $dbh->STORE(Active => 0);
    1;
}

sub commit {
    my ($dbh) = @_;
    if ( $dbh->FETCH('Active') ) {
        return $dbh->FETCH('_dbh')->commit if $dbh->FETCH('_dbh');
    }
    return;
}

sub rollback {
    my ($dbh) = @_;
    if ( $dbh->FETCH('Active') ) {
        return $dbh->FETCH('_dbh')->rollback if $dbh->FETCH('_dbh');
    }
    return;
}

sub get_info {
    my($dbh, $info_type) = @_;

    # return info from current connection
    my $handler = $dbh->FETCH('_handler');
    my $_dbh = $handler->dbh;
    return $_dbh->get_info($info_type);
}

sub STORE {
    my ($self, $attr, $val) = @_;
    $self->{$attr} = $val;
}

sub DESTROY { shift->disconnect }

#######################################################################
package DBD::Multi::st;
use strict;

$DBD::Multi::st::imp_data_size = 0;

use vars qw[@METHODS @FIELDS];
@METHODS = qw[
    bind_param
    bind_param_inout
    bind_param_array
    execute_array
    execute_for_fetch
    fetch
    fetchrow_arrayref
    fetchrow_array
    fetchrow_hashref
    fetchall_arrayref
    fetchall_hashref
    bind_col
    bind_columns
    dump_results
];

@FIELDS = qw[
    NUM_OF_FIELDS
    CursorName
    ParamValues
    RowsInCache
];

sub execute {
    my $sth  = shift;
    my $_sth = $sth->FETCH('_sth');
    my $params =   @_
                 ? $sth->{f_params} = [ @_ ]
                 : $sth->{f_params};

    $sth->finish if $sth->FETCH('Active');
    $sth->{Active} = 1;
    my $rc = $_sth->execute(@{$params});

    for my $field ( @FIELDS ) {
        my $value = $_sth->FETCH($field);
        $sth->STORE($field => $value)
          unless    ! defined $value
                 || defined $sth->FETCH($field);
    }

    return $rc;
}

sub FETCH {
    my ($sth, $attrib) = @_;
    $sth->{'_sth'}->FETCH($attrib) || $sth->{$attrib};
}

sub STORE {
    my ($self, $attr, $val) = @_;
    $self->{$attr} = $val;
}

sub rows { shift->FETCH('_sth')->rows }

sub finish {
    my ($sth) = @_;
    $sth->STORE(Active => 0);
    return $sth->FETCH('_sth')->finish;
}

foreach my $method ( @METHODS ) {
    no strict;
    *{$method} = sub { shift->FETCH('_sth')->$method(@_) };
}

#######################################################################
package DBD::Multi::Handler;
use strict;

use base qw[Class::Accessor::Fast];
use Sys::SigAction qw(timeout_call);
use List::Util qw(shuffle);

=begin ImplementationNotes

dsources - This thing changes from an arrayref to a hashref during construction.  :(

  Initially, when data is passed in during construction, it's an arrayref
  containing the 'dsns' param from the user's connect() call.

  Later, when _configure_dsources gets called, it turns into a multi-dimension
  hashref:

       $dsources->{$pri}->{$dsource_id} = 1;

  The first key is the priority number, the second key is the data source index
  number.  The value is always just a true value.

nextid - A counter.  Stores the index number of the next data source to be added.

all_dsources - A hashref.  Maps index number to the connect data.

current_dsource - The most recent chosen datasource index number.

used - A hashref.  Keys are index numbers.  Values are true when the datasource
has been previously assigned and we want to prefer other datasources of the
same priority (for round-robin load distribution).

failed - A hashref.   Keys are index numbers.   Values are counters indicating
how many times the data source has failed.

failed_last - A hashref.   Keys are index number.   Values are unix timestamp
indicating the most recent time a data source failed.

failed_max - A scalar value.   Number of times a datasource may fail before we
stop trying it.

failed_expire - A scalar value.   Number of seconds since we stopped trying a
datasource before we'll try it again.

timeout - A scalar value.   Number of seconds we try to connect to a datasource
before giving up.

=end ImplementationNotes

=cut

__PACKAGE__->mk_accessors(qw[
    dsources
    nextid
    all_dsources
    current_dsource
    used
    failed
    failed_last
    failed_max
    failed_expire
    timeout
]);

sub new {
    my ($class, $args) = @_;
    my $self     = $class->SUPER::new($args);
    $self->nextid(0) unless defined $self->nextid;
    $self->all_dsources({});
    $self->used({});
    $self->failed({});
    $self->failed_last({});
    $self->failed_max(3) unless defined $self->failed_max;
    $self->failed_expire(60*5) unless defined $self->failed_expire;
    $self->timeout( 5 ) unless defined $self->timeout;
    $self->_configure_dsources;
    return $self;
}

sub all_sources {
    my ($self) = @_;
    return values %{$self->all_dsources};
}

sub add_to_pri {
    my ($self, $pri, $dsource) = @_;
    my $dsource_id = $self->nextid;
    my $dsources   = $self->dsources;
    my $all        = $self->all_dsources;

    $all->{$dsource_id} = $dsource;
    $dsources->{$pri}->{$dsource_id} = 1;

    $self->nextid($dsource_id + 1);
}

sub dbh {
    my $self = shift;
    my $dbh = $self->_connect_dsource;
    return $dbh if $dbh;
    $self->dbh_failed;
    $self->dbh;
}

sub dbh_failed {
    my ($self) = @_;

    my $current_dsource = $self->current_dsource;
    $self->failed->{$current_dsource}++;
    $self->failed_last->{$current_dsource} = time;
}

sub _purge_old_failures {
    my ($self) = @_;
    my $now = time;
    my @all = keys %{$self->all_dsources};

    foreach my $dsource ( @all ) {
        next unless $self->failed->{$dsource};
        if ( ($now - $self->failed_last->{$dsource}) > $self->failed_expire ) {
            delete $self->failed->{$dsource};
            delete $self->failed_last->{$dsource};
        }
    }
}

sub _pick_dsource {
    my ($self) = @_;
    $self->_purge_old_failures;
    my $dsources = $self->dsources;
    my @pri      = sort { $a <=> $b } keys %{$dsources};

    foreach my $pri ( @pri ) {
        my $dsource = $self->_pick_pri_dsource($dsources->{$pri});
        if ( defined $dsource ) {
            $self->current_dsource($dsource);
            return;
        }
    }

    $self->used({});
    return $self->_pick_dsource
      if (grep {$self->failed->{$_} >= $self->failed_max} keys(%{$self->failed})) < keys(%{$self->all_dsources});
    die("All data sources failed!");
}

### _pick_pri_dsource
#  Given a list of data sources, all of the same priority, choose one to use.
#  Passed in a structure like this:
#  _pick_pri_dsource( $self, { 3 => 1, 4 => 1, 5 => 1 } )
#  Where the keys 3,4,5 are dsource id's, and the values are always true.
sub _pick_pri_dsource {
    my ($self, $dsources) = @_;
    my @dsources = sort { $a <=> $b } keys %{$dsources};  # Indexes of data sources to try.
    my @used     = grep { exists $self->used->{$_} } @dsources; # List of data sources already used.
    my @failed   = grep { exists($self->failed->{$_}) && $self->failed->{$_} >= $self->failed_max } @dsources;  # List of data sources that won't be tried.

    # We've used them all and they all failed. Escallate.
    return if @used == @dsources && @failed == @dsources;

    # We've used them all but some are good. Purge and reuse.
    delete @{$self->used}{@dsources} if @used == @dsources;

    foreach my $dsource ( shuffle @dsources ) {
        next if    $self->failed->{$dsource}
                && $self->failed->{$dsource} >= $self->failed_max;
        next if $self->used->{$dsource};

        $self->used->{$dsource} = 1;
        return $dsource;
    }
    return;
}

sub _configure_dsources {
    my ($self) = @_;
    my $dsources = $self->dsources;
    $self->dsources({});

    while ( my $pri = shift @{$dsources} ) {
        my $dsource = shift @{$dsources} or last;
        $self->add_to_pri($pri => $dsource);
    }
}

sub _connect_dsource {
    my ($self, $dsource) = @_;
    unless ( $dsource ) {
        $self->_pick_dsource;
        $dsource = $self->all_dsources->{$self->current_dsource};
    }

    # Support ready-made handles
    return $dsource if UNIVERSAL::isa($dsource, 'DBI::db');

    # Support code-refs which return handles
    if (ref $dsource eq 'CODE') {
        my $handle = $dsource->();
        return $handle if UNIVERSAL::isa($handle, 'DBI::db');
        return undef; # Connect by coderef failed.
    }

    my $dbh;
    local $ENV{DBI_AUTOPROXY};
    if (timeout_call( $self->timeout, sub { $dbh = DBI->connect_cached(@{$dsource}) } )) {
        #warn "Timeout[", $self->current_dsource, "] at ", time, "\n";
    }
    return $dbh;
}

sub connect_dsource {
    my ($self, $dsource) = @_;
    $self->_connect_dsource($dsource);
}

sub multi_do_all {
    my ($self, $code) = @_;

    my @all = values %{$self->all_dsources};

    foreach my $source ( @all ) {
        my $dbh = $self->connect_dsource($source);
        next unless $dbh;
        if ( $dbh->{handler} ) {
            $dbh->{handler}->multi_do_all($code, $source);
            next;
        }
        $code->($dbh);
    }
}

1;
__END__