File Coverage

blib/lib/Data/WeightedRoundRobin.pm

Criterion	Covered	Total	%
statement	91	97	93.8
branch	33	40	82.5
condition	14	20	70.0
subroutine	12	12	100.0
pod	7	7	100.0
total	157	176	89.2

package Data::WeightedRoundRobin;

use strict;
use warnings;
our $VERSION = '0.06';

our $DEFAULT_WEIGHT = 100;
our $BTREE_BORDER = 10;

use Scope::Guard qw(guard);

sub new {
    my ($class, $list, $args) = @_;
    $args ||= {};
    my $self = bless {
        rrlist         => [],
        weights        => 0,
        list_num       => 0,
        default_weight => $args->{default_weight} || $DEFAULT_WEIGHT,
        btree_border   => $args->{btree_border} || $BTREE_BORDER,
    }, $class;
    $self->set($list) if $list;
    return $self;
}

sub _normalize {
    my ($self, $data) = @_;
    return unless defined $data;

    my ($key, $value, $weight);

    # { value => 'foo', weight => 1 }
    if (ref $data eq 'HASH') {
        ($key, $value, $weight) = @$data{qw/key value weight/};
        return unless defined $value;
        return if defined $weight && $weight < 0;
        $key = $value unless defined $key;
        $weight = $self->{default_weight} unless defined $weight;
    }
    # foo
    else {
        # \{ foo => 'bar' }
        if (ref $data eq 'REF' && ref $$data eq 'HASH') {
            $data = $$data;
        }
        $key = $value = $data;
        $weight = $self->{default_weight};
    }

    return { key => $key, value => $value, weight => $weight };
}

sub set {
    my ($self, $list) = @_;
    return unless $list;

    my $normalized = {};
    for my $data (@$list) {
        $data = $self->_normalize($data) || next;
        $normalized->{$data->{key}} = $data;
    }

    my $rrlist = [];
    my $weights = 0;
    for my $key (sort keys %$normalized) {
        unshift @$rrlist, {
            key    => $key,
            value  => $normalized->{$key}{value},
            range  => $weights,
            weight => $normalized->{$key}{weight},
        };
        $weights += $normalized->{$key}{weight};
    }

    $self->{rrlist}   = $rrlist;
    $self->{weights}  = $weights;
    $self->{list_num} = scalar @$rrlist;

    return 1;
}

sub add {
    my ($self, $value) = @_;
    my $rrlist = $self->{rrlist};
    $value = $self->_normalize($value) || return;

    my $added = 1;
    for my $data (@$rrlist) {
        if ($data->{key} eq $value->{key}) {
            $added = 0;
            last;
        }
    }

    if ($added) {
        push @$rrlist, $value;
        $self->set($rrlist);
    }

    return $added;
}

sub replace {
    my ($self, $value) = @_;
    my $rrlist = $self->{rrlist};
    $value = $self->_normalize($value) || return;

    my $replaced = 0;
    for my $data (@$rrlist) {
        if ($data->{key} eq $value->{key}) {
            $data = $value;
            $replaced = 1;
            last;
        }
    }

    if ($replaced) {
        $self->set($rrlist);
    }

    return $replaced;
}

sub remove {
    my ($self, $value) = @_;
    my $rrlist = $self->{rrlist};

    my $removed = 0;
    my $newlist = [];
    for my $data (@$rrlist) {
        unless ($data->{key} eq $value) {
            push @$newlist, $data;
        }
        else {
            $removed = 1;
        }
    }

    if ($removed) {
        $self->set($newlist);
    }

    return $removed;
}

sub next {
    my ($self, $key) = @_;
    my ($rrlist, $weights, $list_num) = @$self{qw/rrlist weights list_num/};
    return unless $list_num; # empty data
    my ($start, $end) = (0, $list_num - 1);

    # if all weight is 0, choose random
    return $rrlist->[int rand $list_num]->{value} if $weights == 0;

    my $rweight = rand($weights);
    if ($list_num < $self->{btree_border}) {
        # linear
        for my $rr (@$rrlist) {
            return $rr->{value} if $rweight >= $rr->{range};
        }
    }
    else {
        # b-tree
        while ($start < $end) {
            my $mid = int(($start + $end) / 2);
            if ($rrlist->[$mid]{range} <= $rweight) {
                $end = $mid;
            }
            else {
                $start = $mid + 1;
            }
        }
        return $rrlist->[$start]{value};
    }
}

sub save {
    my $self = shift;
    my $orig_rrlist = $self->{rrlist};
    guard { $self->set($orig_rrlist) };
}

1;
__END__