ngx-lua-shdict

#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)
# Copyright (C) Guanlan Dai

use 5.006001;
use strict;
use warnings;

use Getopt::Long qw( GetOptions );

GetOptions("a=s",       \(my $stap_args),
           "d",         \(my $dump_src),
           "h",         \(my $help),
           "lua51",     \(my $lua51),
           "luajit20",  \(my $luajit20),
           "w",         \(my $trace_write),
           "f",         \(my $find_value),
           "p=i",       \(my $pid),
           "raw",       \(my $raw),
           "dict=s",    \(my $dict),
           "key=s",     \(my $key))
    or die usage();

if ($help) {
    print usage();
    exit;
}

if (!$pid) {
    die "No nginx process pid specified by the -p option\n";
}

if (!defined $stap_args) {
    $stap_args = '';
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXACTION=/) {
    $stap_args .= " -DMAXACTION=100000"
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXMAPENTRIES=/) {
    $stap_args .= " -DMAXMAPENTRIES=5000"
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXBACKTRACE=/) {
    $stap_args .= " -DMAXBACKTRACE=200"
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXSTRINGLEN=2048/) {
    $stap_args .= " -DMAXSTRINGLEN=2048"
}

if ($stap_args !~ /(?:^|\s)-D\s*MAXSKIPPED=1024/) {
    $stap_args .= " -DMAXSKIPPED=1024"
}

if ($^O ne 'linux') {
    die "Only linux is supported but I am on $^O.\n";
}

my $nginx_file = "/proc/$pid/exe";
if (!-f $nginx_file) {
    die "Nginx process $pid is not running or ",
        "you do not have enough permissions.\n";
}

my $nginx_path = readlink $nginx_file;
my $lua_path = get_lua_path();

my $ver = `stap --version 2>&1`;
if (!defined $ver) {
    die "Systemtap not installed or its \"stap\" utility is not visible to the PATH environment: $!\n";
}

if ($ver =~ /version\s+(\d+\.\d+)/i) {
    my $v = $1;
    if ($v < 2.1) {
        die "ERROR: at least systemtap 2.1 is required but found $v\n";
    }

} else {
    die "ERROR: unknown version of systemtap:\n$ver\n";
}

if (!$key) {
    die "ERROR: please specify a key\n";
}

my $stap_src;
my $preamble;

if ($trace_write) {
    if ($dict) {
        die "ERROR: specify a dict when tracing writes is not supported\n";
    }
    my $o = qq{\@cast(o, "TValue", "$lua_path")};
    my $L = qq{\@cast(L, "lua_State", "$lua_path")};
    my $gcr = qq{\@cast(gcr, "GCRef", "$lua_path")};
    my $gcobj = qq{\@cast(gcobj, "GCobj", "$lua_path")};
    my $str = qq{\@cast(str, "GCstr", "$lua_path")};
    my $sizeof_TValue = qq{\&\@cast(0, "TValue", "$lua_path")[1]};
    my $sizeof_GCstr = qq{\&\@cast(0, "GCstr", "$lua_path")[1]};
    my $TL_TSTR = "4294967291";
    my $NGX_HTTP_LUA_SHDICT_ADD = '0x0001';
    my $NGX_HTTP_LUA_SHDICT_REPLACE = '0x0002';
    my $NGX_HTTP_LUA_SHDICT_SAFE_STORE = '0x0004';
    my $LJ_TISNUM = '0xfffeffffu';
    $preamble = <<_EOC_;
probe begin {
    printf("Tracing %d ($nginx_path)...\\n\\n", target())
    printf("Hit Ctrl-C to end.\\n\\n")
}
_EOC_

    $stap_src = <<_EOC_;
$preamble

function gcref(gcr)
{
    return $gcr->gcptr32
}

function gcval(o)
{
    return gcref(\&$o->gcr)
}

function strdata(s)
{
    return s + $sizeof_GCstr
}

probe process("$nginx_path").function("ngx_http_lua_shdict_set_helper")
{
    if (pid() == target()) {
        flags = \$flags
        if (flags & $NGX_HTTP_LUA_SHDICT_ADD) {
            op = "add"

        } else if (flags & $NGX_HTTP_LUA_SHDICT_REPLACE) {
            op = "replace"

        } else {
            op = "set"
        }

        if (flags & $NGX_HTTP_LUA_SHDICT_SAFE_STORE) {
            op = "safe_" . op
        }

        L = \$L
        o = $L->base + $sizeof_TValue * (2 - 1)
        if (o < $L->top && $o->it == $TL_TSTR) {
            gcobj = gcval(o)
            str = \&$gcobj->str
            //printf("gmatch regex: %s\\n", user_string_n(strdata(str), $str->len))
            key = user_string_n(strdata(str), $str->len)
            if (key == "$key") {
                printf("%s %s", op, key)

                o = $L->base + $sizeof_TValue * (4 - 1)
                if (o < $L->top) {
                    exptime = user_long(&$o->n)
                    printf(" exptime=%d", exptime)
                }

                printf("\\n")
            }
        }
    }
}
_EOC_

} else {
    if (!$dict) {
        die "ERROR: please specify a dict zone\n";
    }

    my $hash = crc32($key);
    my $dict_name_len = length $dict;
    my $key_len = length $key;
    my $quoted_zone = quote_str($dict);
    my $LUA_TBOOLEAN = 1;
    my $LUA_TNUMBER = 3;
    my $LUA_TSTRING = 4;
    my $part = qq{\@cast(part, "ngx_list_part_t")};
    my $node = qq{\@cast(node, "ngx_rbtree_node_t", "$nginx_path")};
    my $sd = qq{\@cast(sd, "ngx_http_lua_shdict_node_t", "$nginx_path")};
    my $sh = qq{\@cast(sh, "ngx_http_lua_shdict_shctx_t", "$nginx_path")};

    my $print_result_func_def;
    my $print_time_elasped_func_def;
    my $set_timer_begin;

    if ($raw) {
        $preamble = "";

        $set_timer_begin = "";

        $print_time_elasped_func_def = "";

        $print_result_func_def = <<_EOC_;
function print_result(sd) {
    if (sd) {
        value_type = $sd->value_type;
        data = $sd->data + $key_len;

        if (value_type == $LUA_TSTRING) {
            print(user_string(data))

        } else if (value_type == $LUA_TNUMBER) {
            print(user_long(data))

        } else if (value_type == $LUA_TBOOLEAN) {
            print(user_long(data))

        } else {
            printf("ERROR: bad value type found for key $key in shared_dict $dict\\n")
        }

    } else {
        exit()
    }
}
_EOC_

    } else {
        $preamble = <<_EOC_;
probe begin {
    printf("Tracing %d ($nginx_path)...\\n\\n", target())
}
_EOC_

    $set_timer_begin = qq{begin = local_clock_us()};

    $print_time_elasped_func_def = <<_EOC_;
elapsed = local_clock_us() - begin
printf("\\n%d microseconds elapsed in the probe handler.\\n", elapsed)
_EOC_

    $print_result_func_def = <<_EOC_;
function print_result(sd) {
    if (sd) {
        value_type = $sd->value_type;
        data = $sd->data + $key_len;
        if (value_type == $LUA_TSTRING) {
            printf("type: LUA_TSTRING\\n")
            printf("value: %s\\n", text_str(user_string(data)))
            printf("expires: %d\\n", $sd->expires)
            printf("flags: 0x%x\\n", $sd->user_flags)

        } else if (value_type == $LUA_TNUMBER) {
            printf("type: LUA_TNUMBER\\n")
            printf("value: %d\\n", user_long(data))
            printf("expires: %d\\n", $sd->expires)
            printf("flags: 0x%x\\n", $sd->user_flags)

        } else if (value_type == $LUA_TBOOLEAN) {
            printf("type: LUA_TBOOLEAN\\n")
            if ( user_long(data) == 0) {
                printf("value: false\\n")

            } else if ( user_long(data) == 1) {
                printf("value: true\\n")

            } else {
                printf("ERROR: value error\\n")
            }

            printf("expires: %d\\n", $sd->expires)
            printf("flags: 0x%x\\n", $sd->user_flags)

        } else {
            error("Bad value type found for key $key in shared_dict $dict\\n")
        }

    } else {
        printf("ERROR: key not found\\n");
    }
}
_EOC_

    }
    chop $preamble;

    $stap_src = <<_EOC_;
$preamble

function ngx_http_lua_shdict_lookup(zone, hash, key) {
    ctx = \@cast(zone, "ngx_shm_zone_t", "$nginx_path")->data
    sh = \@cast(ctx, "ngx_http_lua_shdict_ctx_t", "$nginx_path")->sh
    node = $sh->rbtree->root
    sentinel = $sh->rbtree->sentinel

    while (node != sentinel) {

        if (hash < $node->key) {
            node = $node->left
            continue
        }

        if (hash > $node->key) {
            node = $node->right
            continue
        }

        sd = &$node->color
        data = $sd->data
        data_str = text_str(user_string(data))
        data_key = substr(data_str, 0 , $key_len)

        if (data_key == "$key"){
            return sd

        } else {
            if (data_key < "$key") {
                node = $node->left

            } else {
                node = $node->right
            }
        }
    }
    return 0
}

$print_result_func_def

probe process("$nginx_path").function("ngx_process_events_and_timers"),
      process("$nginx_path").function("ngx_http_init_request")!,
      process("$nginx_path").function("ngx_http_init_connection")
{
    if (pid() == target()) {

        zone_found = 0

        $set_timer_begin

        part = &\@var("ngx_cycle\@ngx_cycle.c")->shared_memory->part
        zone = $part->elts

        for (i = 0; ; i++) {

            if (i >= $part->nelts) {
                if ($part->next == 0) {
                    break
                }

                part = $part->next
                zone = $part->elts
                i = 0
            }

            shm = &\@cast(zone, "ngx_shm_zone_t")[i]->shm
            name = &\@cast(shm, "ngx_shm_t")->name

            if (\@cast(name, "ngx_str_t")->len != $dict_name_len) {
                continue;
            }

            zone_name = user_string_n(\@cast(name, "ngx_str_t")->data, $dict_name_len)
            if (zone_name != $quoted_zone) {
                continue;
            }

            sd = ngx_http_lua_shdict_lookup(zone, $hash, "$key")
            print_result(sd)

            zone_found = 1
            break
        }

        if (!zone_found) {
            printf("dict \\"%s\\" not found.\\n", $quoted_zone)
        }

        $print_time_elasped_func_def

        exit()

    } /* pid() == target() */
}
_EOC_

}

if ($dump_src) {
    print $stap_src;
    exit;
}

open my $in, "|stap $stap_args -x $pid -"
    or die "Cannot run stap: $!\n";

print $in $stap_src;

close $in;

sub usage {
    return <<'_EOC_';
Usage:
    ngx-lua-shdict [optoins]

Options:
    -p <pid>            Specify the nginx worker process pid.
    -a <args>           Pass extra arguments to the stap utility.
    -d                  Dump out the systemtap script source.
    -f                  Find value of specify key.
    -w                  Trace writes to the dict.
    --lua51             The target Nginx is using the standard Lua 5.1 interpreter.
    --luajit20          The target Nginx is using the LuaJIT 2.0.
    --dict <dict>       Specify the dict.
    --key  <key>        Specify the key.
    --raw               Raw output.

Examples:
     ngx-lua-shdict -p 12345
     ngx-lua-shdict -p 12345 -f --dict dogs --key Jim --luajit20
     ngx-lua-shdict -p 12345 -f --dict dogs --key Jim --luajit20 --raw
     ngx-lua-shdict -p 12345 -w --key Jim --luajit
_EOC_
}

sub quote_str {
    my $s = shift;
    $s =~ s/\\/\\\\/g;
    $s =~ s/"/\\"/g;
    $s =~ s/\n/\\n/g;
    $s =~ s/\t/\\t/g;
    $s =~ s/\r/\\r/g;
    return qq{"$s"};
}

sub get_lua_path {
    my $lua_path;

    if (!defined $lua51 && !defined $luajit20) {
        die "Neither --lua51 nor --luajit20 options are specified.\n";
    }

    my $maps_file = "/proc/$pid/maps";
    open my $in, $maps_file
        or die "Cannot open $maps_file for reading: $!\n";

    while (<$in>) {
        if (m{\S+\bliblua-(\d+\.\d+)\.so(?:\.\d+)*$}) {
            my ($path, $ver) = ($&, $1);

            if ($luajit20) {
                die "The --luajit20 option is specified but seen standard Lua library: $path\n";
            }

            if ($ver ne '5.1') {
                die "Nginx server $pid uses a Lua $ver library ",
                    "but only Lua 5.1 is supported.\n";
            }

            $lua_path = $path;
            last;

        } elsif (m{\S+\bliblua\.so(?:\.\d+)*$}) {
            my $path = $&;

            if ($luajit20) {
                die "The --luajit20 option is specified but seen standard Lua library: $path\n";
            }

            $lua_path = $path;
            last;

        } elsif (m{\S+\blibluajit-(\d+\.\d+)\.so(?:\.\d+)*$}) {
            my ($path, $ver) = ($&, $1);

            if ($lua51) {
                die "The --lua51 option is specified but seen the LuaJIT library: $path\n";
            }

            if ($ver ne '5.1') {
                die "Nginx server $pid uses a Lua $ver compatible LuaJIT library ",
                    "but only Lua 5.1 is supported.\n";
            }

            $lua_path = $path;
            last;
        }
    }

    close $in;

    if (!defined $lua_path) {
        #warn "FALL BACK TO NGINX PATH";
        $lua_path = $nginx_path;
    }

    return $lua_path;
}

sub crc32 {
    my ($input, $init_value, $polynomial) = @_;

    $init_value = 0 unless (defined $init_value);
    $polynomial = 0xedb88320 unless (defined $polynomial);

    my @lookup_table;

    for (my $i = 0; $i < 256; $i++) {
        my $x = $i;
        for (my $j = 0; $j < 8; $j++) {
            if ($x & 1) {
                $x = ($x >> 1) ^ $polynomial;

            } else {
                $x = $x >> 1;
            }
        }
        push @lookup_table, $x;
    }
    my $crc = $init_value ^ 0xffffffff;
    foreach my $x (unpack ('C*', $input)) {
        $crc = (($crc >> 8) & 0xffffff) ^ $lookup_table[ ($crc ^ $x) & 0xff ];
    }
    $crc = $crc ^ 0xffffffff;

    return $crc;
}