ngx-req-distr

#!/usr/bin/env perl

# Copyright (C) Yichun Zhang (agentzh)

use 5.006001;
use strict;
use warnings;

use Getopt::Std qw( getopts );

my %opts;

getopts("a:dhm:c", \%opts)
    or die usage();

if ($opts{h}) {
    print usage();
    exit;
}

my $master_pid = $opts{m}
    or die "No nginx master process pid specified by the -m option\n";

if ($master_pid !~ /^\d+$/) {
    die "Bad -m option value \"$master_pid\": not look like a pid\n";
}

my $stap_args = $opts{a} || '';

my $analyze_conn = $opts{c};

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

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

my $nginx_path = readlink $exec_file;

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";
}

my $child_pids = get_child_processes($master_pid);
if (!@$child_pids) {
    push @$child_pids, $master_pid;
}

my $condition = gen_pid_test_condition($child_pids);

#warn "Nginx worker processes: @$child_pids\n";

my $stap_src;

my $preamble = <<_EOC_;
probe begin {
    printf("Tracing @$child_pids ($nginx_path)...\\nHit Ctrl-C to end.\\n")
}
_EOC_

if ($analyze_conn) {
    my $output_code = '';
    for my $pid (@$child_pids) {
        $output_code .= <<_EOC_;
    printf("worker $pid:\\t%d reqs,\\t%d conns\\n", \@count(reqs[$pid]), \@count(conns[$pid]))
_EOC_
    }

    $stap_src = <<_EOC_;
$preamble
global reqs
global conns

probe process("$nginx_path").function("ngx_http_create_request")!,
      process("$nginx_path").function("ngx_http_init_request")
{
    my_pid = pid()
    if ($condition) {
        reqs[my_pid] <<< 1
    }
}

probe process("$nginx_path").function("ngx_http_init_connection")
{
    my_pid = pid()
    if ($condition) {
        conns[my_pid] <<< 1
    }
}

probe end {
    printf("\\n")
$output_code}
_EOC_

} else {
    my $output_code = '';
    for my $pid (@$child_pids) {
        $output_code .= <<_EOC_;
    printf("worker $pid:\\t%d reqs\\n", \@count(reqs[$pid]))
_EOC_
    }

    $stap_src = <<_EOC_;
$preamble
global reqs

probe process("$nginx_path").function("ngx_http_create_request")!,
      process("$nginx_path").function("ngx_http_init_request")
{
    my_pid = pid()
    if ($condition) {
        reqs[my_pid] <<< 1
    }
}

probe end {
    printf("\\n")
$output_code}
_EOC_
}

if ($opts{d}) {
    print $stap_src;
    exit;
}

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

print $in $stap_src;

close $in;

sub usage {
    return <<'_EOC_';
Usage:
    ngx-req-distr [optoins]

Options:
    -a <args>           Pass extra arguments to the stap utility.
    -c                  analyze connections at the same time.
    -d                  Dump out the systemtap script source.
    -h                  Print this usage.
    -m <pid>            Specify the nginx master process pid.

Examples:
    ngx-req-distr -m 12345
    ngx-req-distr -m 12345 -a '-DMAXACTION=100000'
_EOC_
}

sub get_child_processes {
    my $pid = shift;
    my @files = glob "/proc/[0-9]*/stat";
    my @children;
    for my $file (@files) {
        #print "file: $file\n";
        if ($file =~ m{^/proc/$pid/}) {
            next;
        }

        open my $in, $file or next;
        my $line = <$in>;
        close $in;
        if ($line =~ /^(\d+) \S+ \S+ (\d+)/) {
            my ($child, $parent) = ($1, $2);
            if ($parent eq $pid) {
                push @children, $child;
            }
        }
    }

    @children = sort { $a <=> $b } @children;
    return \@children;
}

sub gen_pid_test_condition {
    my $pids = shift;
    my @c;
    for my $pid (@$pids) {
        push @c, "my_pid == $pid";
    }
    return '(' . join(" || ", @c) . ')';
}