File Coverage

blib/lib/Consumer/NonBlock.pm

Criterion	Covered	Total	%
statement	154	175	88.0
branch	50	76	65.7
condition	20	32	62.5
subroutine	23	28	82.1
pod	11	13	84.6
total	258	324	79.6

package Consumer::NonBlock;
use strict;
use warnings;

our $VERSION = '0.002';

use IO::Handle;
use File::Spec;

use File::Path qw/remove_tree/;
use File::Temp qw/tempdir/;
use Fcntl qw/:flock/;
use Carp qw/croak confess/;
use Time::HiRes qw/sleep/;

use Object::HashBase qw{









    +fh
    +batch
    +batch_item
    +buffer
};

sub weaken { $_[0]->{+IS_WEAK} = 1 };

sub init {
    my $self = shift;

    croak "Must be either a reader or a writer" unless $self->{+IS_READER} || $self->{+IS_WRITER};
    croak "Must only be a reader or a writer, not both" if $self->{+IS_READER} && $self->{+IS_WRITER};

    my $dir = $self->{+DIR} or croak "'dir' is a required attribute";
    croak "Invalid directory '$dir'" unless -d $dir;

    if ($self->{+IS_WRITER}) {
        my $data = $self->_update(mode => '+>', open => 1, batch_size => $self->{+BATCH_SIZE} // 100);
        $self->{+BATCH_SIZE} //= $data->{batch_size};
    }
    else {
        my $data = $self->_update();
        $self->{+BATCH_SIZE} //= $data->{batch_size};
        croak "Could not find batch size" unless defined $self->{+BATCH_SIZE};
    }

    delete $self->{+FH};

    $self->{+BATCH}      = 0;
    $self->{+BATCH_ITEM} = 0;
}

sub pair {
    my $class = shift;
    my (%params) = @_;

    my $dir;
    my $base = $params{base_dir} // ($params{shm} ? '/dev/shm' : ());
    if ($base) {
        croak "base dir '$base' is not valid" unless -d $base;
        $dir = tempdir("ConsumerNonBlock-$$-XXXXXX", CLEANUP => 0, DIR => $base);
    }
    else {
        $dir = tempdir("ConsumerNonBlock-$$-XXXXXX", CLEANUP => 0, TMPDIR => 1);
    }

    my $batch_size = $params{+BATCH_SIZE} // 100;

    $params{+BATCH_SIZE} = $batch_size;
    $params{+DIR} = $dir;

    my $writer = $class->new(%params, is_writer => 1);
    my $reader = $class->new(%params, is_reader => 1);

    return ($reader, $writer);
}

sub reader {
    my $class = shift;
    my ($dir, %params) = @_;
    return $class->new(%params, DIR() => $dir, is_reader => 1);
}

sub reader_from_env {
    my $class = shift;
    my $dir = $ENV{CONSUMER_NONBLOCK_DIR} or croak 'The $CONSUMER_NONBLOCK_DIR env var is not set';
    $class->reader($dir, @_);
}

sub writer {
    my $class = shift;
    my ($dir, %params) = @_;
    return $class->new(batch_size => 100, %params, DIR() => $dir, is_writer => 1);
}

sub _data_file {
    my $self_or_class = shift;
    my ($dir) = @_;

    unless ($dir) {
        croak "Cannot call data_file() on a class, must use an instance or pass in a directory"
            unless ref $self_or_class;

        $dir = $self_or_class->dir;
    }

    return File::Spec->catfile($dir, 'data');
}

sub _update {
    my $self = shift;
    my %data = @_;
    my $write = @_;

    my $df = $self->_data_file;

    my $mode = delete($data{mode}) || '+<';

    open(my $fh, $mode, $df) or confess "Could not open file '$df': $!";
    flock($fh, $write ? LOCK_EX : LOCK_SH) or confess "Could not lock file '$df': $!";

    while (my $line = <$fh>) {
        my ($key, $val) = ($line =~ m/^([^:]+):(.+)$/);
        next unless $key;
        $data{$key} //= $val;
    }

    if ($write) {
        seek($fh, 0, 0);
        print $fh "$_\:$data{$_}\n" for keys %data;
    }

    $fh->flush();

    flock($fh, LOCK_UN) or confess "Could not unlock file '$df': $!";
    close($fh);

    return \%data;
}

sub set_env_var {
    my $self = shift;
    $ENV{CONSUMER_NONBLOCK_DIR} = $self->{+DIR};
}

sub write {
    my $self = shift;
    my $count = 0;

    for my $item (@_) {
        next unless defined $item;

        for my $line (split /\n/, $item) {
            $count++;
            $self->write_line($line);
        }
    }

    return $count;
}

sub _check_batch_boundary {
    my $self = shift;
    my (%params) = @_;

    return if $self->{+BATCH_ITEM} < $self->{+BATCH_SIZE};

    delete $self->{+FH};
    $self->{+BATCH_ITEM} = 0;

    unlink(File::Spec->catfile($self->{+DIR}, $self->{+BATCH}))
        if $params{delete};

    $self->{+BATCH} += 1;

    return;
}

sub _batch_fh {
    my $self = shift;
    my ($mode) = @_;

    return $self->{+FH} if $self->{+FH};

    my $file = File::Spec->catfile($self->{+DIR}, $self->{+BATCH});
    my $fh;

    # Only open it if it exists, or we are creating it.
    return unless -e $file || $mode eq '>';

    open($fh, $mode, $file);
    return $self->{+FH} = $fh;
}

sub write_raw {
    my $self = shift;
    my ($raw) = @_;

    croak "No data to write" unless defined $raw;

    $self->_check_batch_boundary();
    my $fh = $self->_batch_fh('>');

    print $fh $raw;
    $fh->flush();
    $self->{+BATCH_ITEM}++;

    return;
}

sub write_line {
    my $self = shift;
    my ($line) = @_;

    croak "No line to write" unless defined $line;

    chomp($line);

    $self->_check_batch_boundary();
    my $fh = $self->_batch_fh('>');

    print $fh $line, "\n";
    $self->{+BATCH_ITEM}++;

    return;
}

sub read_line {
    my $self = shift;

    my $buffer = \$self->{+BUFFER};

    my $loop = 0;
    while (1) {
        # There must be a better way, maybe INotify?
        sleep(0.02) if $loop;
        $loop ||= 1;

        $self->_check_batch_boundary(delete => 1);

        my ($fh, $line);

        if ($fh = $self->_batch_fh('<')) {
            seek($fh, 0, 1) if $fh->eof;
            $line = <$fh>;
            if (defined $line) {
                unless (chomp($line)) {
                    $$buffer = defined($$buffer) ? $$buffer . $line : $line;
                    next;
                }
            }
        }

        if ($fh && defined $line) {
            $self->{+BATCH_ITEM}++;
            if ($$buffer) {
                delete $self->{+BUFFER};
                return $$buffer . $line if $$buffer;
            }

            return $line;
        }
        else {
            my $data = $self->_update();
            next if $data->{open};

            unlink(File::Spec->catfile($self->{+DIR}, $self->{+BATCH}));

            return undef;
        }
    }
}

sub read_lines {
    my $self = shift;
    my @out;

    while (my $line = $self->read_line) {
        push @out => $line;
    }

    return @out;
}

sub DESTROY {
    my $self = shift;

    return if $self->is_weak;

    if ($self->{+IS_WRITER}) {
        $self->_update(open => 0);
        return;
    }

    if ($self->{+IS_READER}) {
        remove_tree($self->{+DIR}, {safe => 1, keep_root => 0});
        return;
    }
}

sub close { $_[0] = undef }

1;

__END__