#!/usr/bin/perl

#  dormando's awesome memcached top utility!

#

#  Copyright 2009 Dormando (dormando@rydia.net).  All rights reserved.

#

#  Use and distribution licensed under the BSD license.  See

#  the COPYING file for full text.

use strict;

use warnings FATAL => 'all';

use AnyEvent;

use AnyEvent::Socket;

use AnyEvent::Handle;

use Getopt::Long;

use YAML qw/Dump Load LoadFile/;

use Term::ReadKey qw/ReadMode ReadKey GetTerminalSize/;

our $VERSION = '0.1';

my $CLEAR     = `clear`;

my @TERM_SIZE = ();

$|++;

my %opts = ();

GetOptions(\%opts, 'help|h', 'config=s');

if ($opts{help}) {

    show_help(); exit;

}

$SIG{INT} = sub {

    ReadMode('normal');

    print "\n";

    exit;

};

# TODO: make this load from central location, and merge in homedir changes.

# then merge Getopt::Long stuff on top of that

# TODO: Set a bunch of defaults and merge in.

my $CONF = load_config();

my %CONS    = ();

my $LAST_RUN = time; # time after the last loop cycle.

my $TIME_SINCE_LAST_RUN = time; # time since last loop cycle.

my $loop_timer;

my $main_cond;

my $prev_stats_results;

my %display_modes = (

    't' => \&display_top_mode,

    '?' => \&display_help_mode,

    'h' => \&display_help_mode,

);

my %column_compute = (

    'hostname' => { stats => [], code => \&compute_hostname},

    'hit_rate' => { stats => ['get_hits', 'get_misses'],

                    code  => \&compute_hit_rate },

    'fill_rate' => { stats => ['bytes', 'limit_maxbytes'],

                    code => \&compute_fill_rate },

);

my %column_format = (

    'hit_rate' => \&format_percent,

    'fill_rate' => \&format_percent,

);

# This can collapse into %column_compute

my %column_format_totals = (

    'hit_rate' => 0,

    'fill_rate' => 0,

);

ReadMode('cbreak');

my $LAST_KEY = '';

my $read_keys = AnyEvent->io (

    fh => \*STDIN, poll => 'r',

    cb => sub {

        $LAST_KEY = ReadKey(-1);

        # If there is a running timer, cancel it.

        # Don't want to interrupt a main loop run.

        # fire_main_loop()'s iteration will pick up the keypress.

        if ($loop_timer) {

            $loop_timer = undef;

            $main_cond->send;

        }

    }

);

# start main loop

fire_main_loop();

### AnyEvent related code.

sub fire_main_loop {

    for (;;) {

        $loop_timer = undef;

        $main_cond = AnyEvent->condvar;

        my $time_taken = main_loop();

        my $delay = $CONF->{delay} - $time_taken;

        $delay = 0 if $delay < 0;

        $loop_timer = AnyEvent->timer(

            after => $delay,

            cb    => $main_cond,

        );

        $main_cond->recv;

    }

}

sub main_loop {

    my $start = AnyEvent->now; # use ->time to find the end.

    maintain_connections();

    my $cv = AnyEvent->condvar;

    # FIXME: Need to dump early if there're no connected conns

    # FIXME: Make this only fetch stats from cons we care to visualize?

    # maybe keep everything anyway to maintain averages?

    my %stats_results = ();

    while (my ($hostname, $con) = each %CONS) {

        $cv->begin;

        call_stats($con, ['', 'items', 'slabs'], sub {

            $stats_results{$hostname} = shift;

            $cv->end;

        });

    }

    $cv->recv;

    # Short circuit since we don't have anything to compare to.

    unless ($prev_stats_results) {

        $prev_stats_results = \%stats_results;

        return $CONF->{delay};

    }

    # Semi-exact global time diff for stats that want to average

    # themselves per-second.

    my $this_run = AnyEvent->time;

    $TIME_SINCE_LAST_RUN = $this_run - $LAST_RUN;

    $LAST_RUN = $this_run;

    # Done all our fetches. Drive the display.

    display_run($prev_stats_results, \%stats_results);

    $prev_stats_results = \%stats_results;

    my $end  = AnyEvent->time;

    my $diff = $LAST_RUN - $start;

    print "loop took: $diff";

    return $diff;

}

sub maintain_connections {

    my $cv    = AnyEvent->condvar;

    $cv->begin (sub { shift->send });

    for my $host (@{$CONF->{servers}}) {

        next if $CONS{$host};

        $cv->begin;

        $CONS{$host} = connect_memcached($host, sub {

            if ($_[0] eq 'err') {

                print "Failed connecting to $host: ", $_[1], "\n";

                delete $CONS{$host};

            }

            $cv->end;

        });

    }

    $cv->end;

    $cv->recv;

}

sub connect_memcached {

    my ($fullhost, $cb)   = @_;

    my ($host, $port) = split /:/, $fullhost;

    my $con; $con = AnyEvent::Handle->new (

        connect => [$host => $port],

        on_connect => sub {

            $cb->('con');

        },

        on_connect_error => sub {

            $cb->('err', $!);

            $con->destroy;

        },

        on_eof   => sub {

            $cb->('err', $!);

            $con->destroy;

        },

    );

    return $con;

}

# Function's getting a little weird since I started optimizing it.

# As of my first set of production tests, this routine is where we spend

# almost all of our processing time.

sub call_stats {

    my ($con, $cmds, $cb) = @_;

    my $stats = {};

    my $num_types = @$cmds;

    my $reader; $reader = sub {

        my ($con, $results) = @_;

        {

            my %temp = ();

            for my $line (split(/\n/, $results)) {

                my ($k, $v) = (split(/\s+/, $line))[1,2];

                $temp{$k} = $v;

            }

            $stats->{$cmds->[0]} = \%temp;

        }

        shift @$cmds;

        unless (@$cmds) {

            # Out of commands to process, return goodies.

            $cb->($stats);

            return;

        }

    };

    for my $cmd (@$cmds) {

        $con->push_write('stats ' . $cmd . "\n");

        $stats->{$cmd} = {};

        $con->push_read(line => "END\r\n", $reader);

    }

}

### Compute routines

sub compute_hostname {

    return $_[0];

}

sub compute_hit_rate {

    my $s = $_[1];

    my $total = $s->{get_hits} + $s->{get_misses};

    return 'NA' unless $total;

    return $s->{get_hits} / $total;

}

sub compute_fill_rate {

    my $s = $_[1];

    return $s->{bytes} / $s->{limit_maxbytes};

}

sub format_column {

    my ($col, $val) = @_;

    my $res;

    $col =~ s/^all_//;

    if ($column_format{$col}) {

        if (ref($column_format{$col}) eq 'CODE') {

            return $column_format{$col}->($val);

        } else {

            return $val .= $column_format{$col};

        }

    } else {

        return format_commas($val);

    }

}

sub column_can_total {

    my $col = shift;

    $col =~ s/^all_//;

    return 1 unless exists $column_format_totals{$col};

    return $column_format_totals{$col};

}

### Display routines

# If there isn't a specific column type computer, see if we just want to

# look at the specific stat and return it.

# If column is a generic type and of 'all_cmd_get' format, return the more

# complete stat instead of the diffed stat.

sub compute_column {

    my ($col, $host, $prev_stats, $curr_stats) = @_;

    my $diff_stats = 1;

    $diff_stats    = 0 if ($col =~ s/^all_//);

    # Really should decide on whether or not to flatten the hash :/

    my $find_stat = sub {

        for my $type (keys %{$_[0]}) {

            return $_[0]->{$type}->{$_[1]} if exists $_[0]->{$type}->{$_[1]};

        }

    };

    my $diff_stat = sub {

        my $stat = shift;

        return 'NA' unless defined $find_stat->($curr_stats, $stat);

        if ($diff_stats) {

            my $diff = eval {

                return ($find_stat->($curr_stats, $stat)

                       - $find_stat->($prev_stats, $stat))

                       / $TIME_SINCE_LAST_RUN;

            };

            return 'NA' if ($@);

            return $diff;

        } else {

            return $find_stat->($curr_stats, $stat);

        }

    };

    if (my $comp = $column_compute{$col}) {

        my %s = ();

        for my $stat (@{$comp->{stats}}) {

            $s{$stat} = $diff_stat->($stat);

        }

        return $comp->{code}->($host, \%s);

    } else {

        return $diff_stat->($col);

    }

    return 'NA';

}

# We have a bunch of stats from a bunch of connections.

# At this point we run a particular display mode, capture the lines, then

# truncate and display them.

sub display_run {

    my $prev_stats = shift;

    my $curr_stats = shift;

    @TERM_SIZE = GetTerminalSize;

    die "cannot detect terminal size" unless $TERM_SIZE[0] && $TERM_SIZE[1];

    if ($LAST_KEY eq 'q') {

        print "\n";

        ReadMode('normal'); exit;

    }

    if ($LAST_KEY ne $CONF->{mode} && exists $display_modes{$LAST_KEY}) {

        $CONF->{prev_mode} = $CONF->{mode};

        $CONF->{mode} = $LAST_KEY;

    } elsif ($CONF->{mode} eq 'h' || $CONF->{mode} eq '?') {

        # Bust out of help mode on any key.

        $CONF->{mode} = $CONF->{prev_mode};

    }

    my $lines = $display_modes{$CONF->{mode}}->($prev_stats, $curr_stats);

    display_lines($lines) if $lines;

}

# Default "top" mode.

# create a set of computed columns as requested by the config.

# this has gotten a little out of hand... needs more cleanup/abstraction.

sub display_top_mode {

    my $prev_stats = shift;

    my $curr_stats = shift;

    my @columns = @{$CONF->{top_mode}->{columns}};

    my @rows    = ();

    my @tot_row = ();

    # Round one.

    for my $host (sort keys %{$curr_stats}) {

        my @row = ();

        for my $colnum (0 .. @columns-1) {

            my $col = $columns[$colnum];

            my $res = compute_column($col, $host, $prev_stats->{$host},

                      $curr_stats->{$host});

            $tot_row[$colnum] += $res if is_numeric($res);

            push @row, $res;

        }

        push(@rows, \@row);

    }

    # Sort rows by sort column (ascending or descending)

    if (my $sort = $CONF->{top_mode}->{sort_column}) {

        my $order  = $CONF->{top_mode}->{sort_order} || 'asc';

        my $colnum = 0;

        for (0 .. @columns-1) { $colnum = $_ if $columns[$_] eq $sort; }

        my @newrows;

        if ($order eq 'asc') {

            if (is_numeric($rows[0]->[$colnum])) {

                @newrows = sort { $a->[$colnum] <=> $b->[$colnum] } @rows;

            } else {

                @newrows = sort { $a->[$colnum] cmp $b->[$colnum] } @rows;

            }

        } else {

            if (is_numeric($rows[0]->[$colnum])) {

                @newrows = sort { $b->[$colnum] <=> $a->[$colnum] } @rows;

            } else {

                @newrows = sort { $b->[$colnum] cmp $a->[$colnum] } @rows;

            }

        }

        @rows = @newrows;

    }

    # Format each column after the sort...

    {

        my @newrows = ();

        for my $row (@rows) {

            my @newrow = ();

            for my $colnum (0 .. @columns-1) {

                push @newrow, is_numeric($row->[$colnum]) ?

                            format_column($columns[$colnum], $row->[$colnum]) :

                            $row->[$colnum];

            }

            push @newrows, \@newrow;

        }

        @rows = @newrows;

    }

    # Create average and total rows.

    my @avg_row = ();

    for my $col (0 .. @columns-1) {

        if (is_numeric($tot_row[$col])) {

            my $countable_rows = 0;

            for my $row (@rows) {

                next unless $row->[$col];

                $countable_rows++ unless $row->[$col] eq 'NA';

            }

            $countable_rows = 1 unless $countable_rows;

            push @avg_row, format_column($columns[$col],

                 sprintf('%.2f', $tot_row[$col] / $countable_rows));

        } else {

            push @avg_row, 'NA';

        }

        $tot_row[$col] = 'NA' unless defined $tot_row[$col];

        $tot_row[$col] = 'NA' unless (column_can_total($columns[$col]));

        $tot_row[$col] = format_column($columns[$col], $tot_row[$col])

                         unless $tot_row[$col] eq 'NA';

    }

    unshift @rows, \@avg_row;

    unshift @rows, ['AVERAGE:'];

    unshift @rows, \@tot_row;

    unshift @rows, ['TOTAL:'];

    # Round two. Pass @rows into a function which returns an array with the

    # desired format spacing for each column.

    unshift @rows, \@columns;

    my $spacing = find_optimal_spacing(\@rows);

    my @display_lines = ();

    for my $row (@rows) {

        my $line = '';

        for my $col (0 .. @$row-1) {

            my $space = $spacing->[$col];

            $line .= sprintf("%-${space}s ", $row->[$col]);

        }

        push @display_lines, $line;

    }

    return \@display_lines;

}

sub display_help_mode {

    my $help = <<"ENDHELP";

dormando's awesome memcached top utility version v$VERSION

This early version requires you to edit the ~/.damemtop/damemtop.yaml

(or /etc/damemtop.yaml) file in order to change options.

See --help for more info.

Hit any key to exit help.

ENDHELP

    my @lines = split /\n/, $help;

    display_lines(\@lines);

    $LAST_KEY = ReadKey(0);

    return;

}

# Takes a set of lines, clears screen, dumps header, trims lines, etc

# MAYBE: mode to wrap lines instead of trim them?

sub display_lines {

    my $lines = shift;

    my $width         = $TERM_SIZE[0];

    my $height_remain = $TERM_SIZE[1];

    unshift @$lines, display_header($width);

    clear_screen() unless $CONF->{no_clear};

    while (--$height_remain && @$lines) {

        # truncate too long lines.

        my $line = shift @$lines;

        $line = substr $line, 0, $width-1;

        print $line, "\n";

    }

}

sub display_header {

    my $topbar = 'damemtop: ' . scalar localtime;

    if ($CONF->{mode} eq 't' && $CONF->{top_mode}->{sort_column}) {

        $topbar .= ' [sort: ' . $CONF->{top_mode}->{sort_column} . ']';

    }

    $topbar .= ' [delay: ' . $CONF->{delay} . 's]';

    return $topbar;

}

### Utilities

# find the optimal format spacing for each column, which is:

# longest length of item in col + 2 (whitespace).

sub find_optimal_spacing {

    my $rows  = shift;

    my @maxes = ();

    my $num_cols = @{$rows->[0]};

    for my $row (@$rows) {

        for my $col (0 .. $num_cols-1) {

            $maxes[$col] = 0 unless $maxes[$col];

            next unless $row->[$col];

            $maxes[$col] = length($row->[$col])

                if length($row->[$col]) > $maxes[$col];

        }

    }

    for my $col (0 .. $num_cols) {

        $maxes[$col] += 1;

    }

    return \@maxes;

}

# doesn't try too hard to identify numbers...

sub is_numeric {

    return 0 unless $_[0];

    return 1 if $_[0] =~ m/^\d+(\.\d*)?(\w+)?$/;

    return 0;

}

sub format_percent {

    return sprintf("%.2f%%", $_[0] * 100);

}

sub format_commas {

    my $num = shift;

    $num = int($num);

    $num =~ s/(^[-+]?\d+?(?=(?>(?:\d{3})+)(?!\d))|\G\d{3}(?=\d))/$1,/g;

    return $num;

}

# Can tick counters/etc here as well.

sub clear_screen {

    print $CLEAR;

}

# tries minimally to find a localized config file.

# TODO: Handle the YAML error and make it prettier.

sub load_config {

    my $config = $opts{config} if $opts{config};

    my $homedir = "$ENV{HOME}/.damemtop/damemtop.yaml";

    if (-e $homedir) {

        $config = $homedir;

    } else {

        $config = '/etc/damemtop.yaml';

    }

    return LoadFile($config);

}

sub show_help {

    print <<"ENDHELP";

dormando's awesome memcached top utility version v$VERSION

This program is copyright (c) 2009 Dormando.

Use and distribution licensed under the BSD license.  See

the COPYING file for full text.

contact: dormando\@rydia.net or memcached\@googlegroups.com.

This early version requires you to edit the ~/.damemtop/damemtop.yaml

(or /etc/damemtop.yaml) file in order to change options.

You may display any column that is in the output of

'stats', 'stats items', or 'stats slabs' from memcached's ASCII protocol.

Start a column with 'all_' (ie; 'all_get_hits') to display the current stat,

otherwise the stat is displayed as an average per second.

Specify a "sort_column" under "top_mode" to sort the output by any column.

Some special "computed" columns exist:

hit_rate (get/miss hit ratio)

fill_rate (% bytes used out of the maximum memory limit)

ENDHELP

    exit;

}