#!/usr/bin/perl
#
# Danga's Statistics Server
# This program listens on a port and acts as a very generic data server.  It lets you get
# and set data as well as subscribe to information flows given glob patterns.  Commands:
#
#   get <what>
#       get some particular data.  also gives you the time this data was set.  you can also
#       specify a pseudo-regexp here in the same format as the subscribe command.
#   [q]set [time] <what> <data>
#       set what to be data.  if you specify a time, it will set the time you set it as that
#       time.  also, if you specify qset, you will not be given an OK prompt afterwards.
#   quickset <on|off>
#       if you turn quickset on, sets will be processed without processing the subscriptions,
#       instead waiting until you turn quickset off, at which point all pending subscriptions
#       will be processed.
#   subscribe <pattern>
#       pattern is some combination of letters and asterisks. "foo.*" subscribes to anything
#       that happens to start with "foo.".  you can also subscribe to simply "*" to get all
#       information that gets set.  or, "*.system" to see everything relating to a system,
#       e.g. mysql, memcache, etc.
#   list
#       list your subscription patterns.  if you want to remove subscriptions, you need to
#       reconnect to the server.
#   listall
#       list all active subscriptions.  (see what other people are watching, if you care.)
#   quit
#       disconnect your socket nicely.
#   shutdown
#       kill the statistics server.
#   uptime
#       print out the current server statistics... number of keys and subscriptions and the
#       time the server has been running.
#
# Command line options:
#
#    -p INT     set the port the server should listen on
#    -D         if present, tell the server to daemonize
#
# Copyright 2004, Danga Interactive
#
# Authors:
#   Mark Smith <marksmith@danga.com>
#
# License:
#   undecided.
#

package StatServ;

use strict;
use Getopt::Long;
use IO::Socket::INET;
use POSIX;
use Carp;
use Socket qw(IPPROTO_TCP SO_KEEPALIVE TCP_NODELAY SOL_SOCKET);
use constant TCP_KEEPIDLE  => 4; # Start keeplives after this period
use constant TCP_KEEPINTVL => 5; # Interval between keepalives
use constant TCP_KEEPCNT   => 6; # Number of keepalives before death

# declare our globals/defaults
my $conf_port = 9107;
my $daemonize = 0;
my $uptime = time();
my $quickset = 0;
my $lastcleantime = time();

# internal data, not to be touched externally
my %data; # { "fact-name" => [ [ set-time, data ], [ set-time, data ], ... ] }
my @sortedkeys; # sorted keys (for pretty output)
my %subs; # { "fact-name" => [ $subscriber, $subscriber, ... ] }
my @sublist; # list of subscriptions

# Command-line options will override
Getopt::Long::Configure('bundling');
Getopt::Long::GetOptions(
    'p|port=i' => \$conf_port,
    'D|daemon' => \$daemonize,
);

# establish SERVER socket, bind and listen.
my $server = IO::Socket::INET->new(LocalPort => $conf_port,
                                   Type      => SOCK_STREAM,
                                   Proto     => 'tcp',
                                   Blocking  => 1,
                                   Reuse     => 1,
                                   Listen    => 10)
    or die "Error creating socket: $@\n";

# make socket nonblocking
IO::Handle::blocking($server, 0);

my $accept_handler = sub {
    my $csock = $server->accept();
    return unless $csock;

    IO::Handle::blocking($csock, 0);
    setsockopt($csock, IPPROTO_TCP, TCP_NODELAY, pack("l", 1)) or die;

    # Enable keep alive
    (setsockopt($csock, SOL_SOCKET, SO_KEEPALIVE,   pack("l", 1))  &&
     setsockopt($csock, IPPROTO_TCP, TCP_KEEPIDLE,  pack("l", 30)) &&
     setsockopt($csock, IPPROTO_TCP, TCP_KEEPCNT,   pack("l", 10)) &&
     setsockopt($csock, IPPROTO_TCP, TCP_KEEPINTVL, pack("l", 30)) &&
     1
    ) || die "Couldn't set keep-alive settings on socket (Not on Linux?)";

    my $client = Client->new($csock);
    $client->watch_read(1);
};

# daemonize if we should
daemonize() if $daemonize;

Client->OtherFds(fileno($server) => $accept_handler);
Client->EventLoop();
exit 0;

#############################################################################
###### SUBROUTINES ##########################################################
#############################################################################

sub daemonize {
    my($pid, $sess_id, $i);

    ## Fork and exit parent
    if ($pid = fork) { exit 0; }

    ## Detach ourselves from the terminal
    croak "Cannot detach from controlling terminal"
        unless $sess_id = POSIX::setsid();

    ## Prevent possibility of acquiring a controling terminal
    $SIG{'HUP'} = 'IGNORE';
    if ($pid = fork) { exit 0; }

    ## Change working directory
    chdir "/";

    ## Clear file creation mask
    umask 0;

    ## Close open file descriptors
    close(STDIN);
    close(STDOUT);
    close(STDERR);

    ## Reopen stderr, stdout, stdin to /dev/null
    open(STDIN,  "+>/dev/null");
    open(STDOUT, "+>&STDIN");
    open(STDERR, "+>&STDIN");
}

# internal setter that has the logic that actually handles setting key/value
# pairs.  also performs notifications, unless specifically told otherwise.
sub _set {
    my ($time, $key, $val, $nosubs) = @_;
    $nosubs ||= 0;
    
    # push onto array
    unless (defined $data{$key}) {
        $data{$key} = [];
        @sortedkeys = sort keys %data;
    }
    unshift @{$data{$key}}, [ $time, "$val" ];
    pop @{$data{$key}} if @{$data{$key}} > 100;

    # handle subscriptions
    unless ($nosubs) {
        unless (defined $subs{$key}) {
            $subs{$key} = []; # define it and set blank if it's not defined
            foreach my $sub (@sublist) {
                push @{$subs{$key}}, $sub->[0]
                    if $key =~ /$sub->[1]/i;
            }
        }
        $_->set_line("$time $key $val") foreach @{$subs{$key}};
    }   

    # done
    return 1;
}

# internal getter.  gets a value for a single key.
sub _get {
    my $key = shift;
    return wantarray ? () : undef unless $data{$key};

    # now, get data. return all points? or just most recent?
    return $data{$key} if wantarray;
    return undef unless @{$data{$key}};
    return $data{$key}->[0]->[1]; # text of most recent value
}

# internal incrementers.  simply gets a value and sets it to val +/- 1.
# this is used by stats within the server only, and not exposed to the
# outside world as of right now.
sub _incr { _set(time(), $_[0], _get($_[0])+1, 1); }
sub _decr { _set(time(), $_[0], _get($_[0])-1, 1); }

#####################################################################
### C L I E N T   C L A S S
#####################################################################
package Client;

use Danga::Socket ();
use base qw{Danga::Socket};
use fields qw(read_buf t_out_buf do_buffer);

sub new {
    my Client $self = shift;
    $self = fields::new($self) unless ref $self;
    $self->SUPER::new( @_ );

    StatServ::_incr("server.clients.total");
    return $self;
}

# Client
sub event_read {
    my $self = shift;

    my $bref = $self->read(1024);
    return $self->close() unless defined $bref;
    $self->{read_buf} .= $$bref;

    StatServ::_incr("server.clients.event_reads");
    while ($self->{read_buf} =~ s/^(.+?)\r?\n//) {
        my $line = $1;
        $self->process_line( $line );
    }
}

# Client
sub event_err { StatServ::_incr("server.clients.event_errs"); my $self = shift; $self->close; }
sub event_hup { StatServ::_incr("server.clients.event_hups"); my $self = shift; $self->close; }

sub process_line {
    my Client $self = shift;
    my $line = shift;

    StatServ::_incr("server.clients.process_lines");

    # clean out closed clients every few seconds
    my $now = time();
    if ($lastcleantime + 10 < $now) {
        foreach my $key (keys %subs) {
            @{$subs{$key}} = grep { !$_->{closed} } @{$subs{$key}};
        }
        @sublist = grep { !$_->[0]->{closed} } @sublist;
        $lastcleantime = $now;
    }

    if ($line =~ /^(\w+)\s*(.*)/) {
        my ($cmd, $args) = ($1, $2);
        $cmd = lc($cmd);
        my ($one, $two);
        if ($args =~ /^([-\w.:]+)\s+(.+)/) {
            ($one, $two) = ($1, $2);
        }

        # statistics keeping
        StatServ::_incr("server.clients.commands");

        # see if this is a command we know about
        if ($cmd eq 'get') {
            # get something
            StatServ::_incr("server.com_$cmd");
            return $self->err_line('no_args') unless $args;
            if (defined $data{$args} && ref $data{$args} eq 'ARRAY') {
                my $ct = 0;
                $self->begin_output();
                foreach my $row (@{$data{$args}}) {
                    $ct++;
                    $self->ok_line("$ct $row->[0] $row->[1]");
                }
                $self->ok_line('done');
                return $self->end_output();
            } else {
                # see if this should be used as a regular expression
                if ($args =~ /[.*]/) {
                    $args =~ s/\./\\./g;
                    $args =~ s/\*/.*/g;
                    $self->begin_output();
                    foreach my $key (@sortedkeys) {
                        if ($key =~ /$args/i) {
                            my $first = $data{$key}->[0];
                            $self->ok_line("$key $first->[0] $first->[1]");
                        }
                    }
                    $self->ok_line('done');
                    return $self->end_output();
                }
                return $self->err_line("not_found");
            }

        } elsif ($cmd =~ /^(q)?set$/) {
            # set something
            StatServ::_incr("server.com_$cmd");
            return $self->err_line('need_two_args') unless defined $one && defined $two;
            my $quiet = $1 eq 'q' ? 1 : 0;

            # see if $one happens to be a number?  that means they gave us a time to use.
            my $time;
            if ($one =~ /^\d+$/ && $two =~ /^([-\w.:]+)\s+(.+)$/) {
                $time = $one + 0;
                ($one, $two) = ($1, $2);
            }
            $time = time() unless defined $time;

            # push data onto front
            StatServ::_set($time, $one, $two);

            # all done
            return $quiet ? 1 : $self->ok_line();

        } elsif ($cmd =~ /^sub/) {
            # subscribe to something
            # convert "*.mysql" to ".*\.mysql", etc.
            StatServ::_incr("server.com_$cmd");
            return $self->err_line('no_args') unless $args;
            $args =~ s/\./\\./g;
            $args =~ s/\*/.*/g;
            return $self->ok_line($args) # no dupes in @sublist!
                if grep { $args eq $_->[1] && $self == $_->[0] } @sublist;
            push @sublist, [ $self, $args ];

            # now see what subscriptions this matches
            foreach my $key (keys %subs) {
                if ($key =~ /$args/i) {
                    # no dupes!
                    next if grep { $_ == $self } @{$subs{$key}};
                    push @{$subs{$key}}, $self;
                }
            }

            # return okay to the person who set something
            return $self->ok_line($args);

        } elsif ($cmd =~ /^list(all)?$/) {
            # list out your subscriptions
            StatServ::_incr("server.com_$cmd");
            foreach my $sub (@sublist) {
                $self->ok_line("$sub->[0] $sub->[1]")
                    if $1 eq 'all' || $sub->[0] == $self;
            }
            return $self->ok_line('done');

        } elsif ($cmd eq 'shutdown') {
            # kill ourselves
            StatServ::_incr("server.com_$cmd");
            exit 0;

        } elsif ($cmd eq 'uptime') {
            # figure out our uptime
            StatServ::_incr("server.com_$cmd");
            my $c = time() - $uptime;
            my ($d, $h, $m, $s) = (int($c / 86400), int(($c % 86400) / 3600),
                                   int(($c % 3600) / 60), int($c % 60));
            my $l = sprintf 'stats-server up %d days %02d:%02d:%02d', $d, $h, $m, $s;
            $self->begin_output();
            $self->ok_line($l);
            $l = sprintf '%d subscriptions, %d total keys', scalar(@sublist), scalar(@sortedkeys);
            $self->ok_line($l);
            $self->ok_line('done');
            return $self->end_output();

        } elsif ($cmd eq 'quickset') {
            StatServ::_incr("server.com_$cmd");
            return $self->err_line('not_on_or_off') unless $args eq 'on' || $args eq 'off';

            # now set it?
            $quickset = $args eq 'on' ? 1 : 0;
            return 1;

        } elsif ($cmd eq 'quit') {
            # simply quit ... at some point we should clear out subscriptions that they
            # have setup... but until then? nah.
            StatServ::_incr("server.com_$cmd");
            return $self->close;

        }
    }

    # oops, don't know what they wanted
    return $self->err_line('unknown_command');
}

sub ok_line {
    my Client $self = shift;
    my $args = shift;
    if ($self->{do_buffer}) {
        StatServ::_incr("server.lines.ok.buffered");
        $self->{t_out_buf} .= "OK $args\r\n";
    } else {
        StatServ::_incr("server.lines.ok.immediate");
        $self->write("OK $args\r\n");
    }
    return 1;
}

sub set_line {
    my Client $self = shift;
    my $args = shift;
    if ($self->{do_buffer}) {
        StatServ::_incr("server.lines.set.buffered");
        $self->{t_out_buf} .= "SET $args\r\n";
    } else {
        StatServ::_incr("server.lines.set.immediate");
        $self->write("SET $args\r\n");
    }
    return 1;
}

sub err_line {
    my Client $self = shift;
    my $err_code = shift;
    my $err_text = {
        'unknown_command' => "Unknown server command",
        'not_found' => "Information not in record",
        'no_args' => "No arguments provided",
        'need_two_args' => "Two arguments needed",
        'no_subscriptions' => "No subscriptions found",
        'not_on_or_off' => "Argument wasn't of form 'on' or 'off'",
    }->{$err_code};

    if ($self->{do_buffer}) {
        StatServ::_incr("server.lines.err.buffered");
        $self->{t_out_buf} .= "ERR $err_code $err_text\r\n";
    } else {
        StatServ::_incr("server.lines.err.immediate");
        $self->write("ERR $err_code $err_text\r\n");
    }
    return 0;
}

sub begin_output {
    my Client $self = shift;
    StatServ::_incr("server.lines.begin_outputs");
    $self->{do_buffer} = 1;
    return 1;
}

sub end_output {
    my Client $self = shift;
    StatServ::_incr("server.lines.end_outputs");
    $self->{do_buffer} = 0;
    $self->write($self->{t_out_buf});
    $self->{t_out_buf} = '';
    return 1;
}

sub eurl
{
    my $a = $_[0];
    $a =~ s/([^a-zA-Z0-9_\,\-.\/\\\: ])/uc sprintf("%%%02x",ord($1))/eg;
    $a =~ tr/ /+/;
    return $a;
}

# Local Variables:
# mode: perl
# c-basic-indent: 4
# indent-tabs-mode: nil
# End: