#!/usr/bin/env perl
use 5.010;
use strict;

##############################################################################
#                                                                            #
# Copyright 2011 New Dream Network LLC dba DreamHost ("DreamHost").          #
# All rights reserved.                                                       #
#                                                                            #
# The intellectual and technical concepts contained herein are proprietary   #
# to "DreamHost", and are protected by trade secret and/or copyright law.    #
# Disclosure, use, or reproduction without explicit authorization of         #
# "DreamHost" is prohibited.                                                 #
#                                                                            #
##############################################################################

use Data::Dumper;
use Time::HiRes qw( time sleep );
use POSIX qw( :fcntl_h );
use Getopt::Long;
use JSON;

use constant {
	SYSTEM_UID     => 1000,
	SYSTEM_GID     => 1000,
	NOBODY_UID     => 65534,

	# stat() fields
	st_dev         => 0,
	st_ino         => 1,
	st_mode        => 2,
	st_nlink       => 3,
	st_uid         => 4,
	st_gid         => 5,
	st_rdev        => 6,
	st_size        => 7,
	st_atime       => 8,
	st_mtime       => 9,
	st_ctime       => 10,
	st_blksize     => 11,
	st_blocks      => 12,

	# /proc/$pid/stat fields (details: man proc)
	proc_ppid      => 0,
	proc_pgrp      => 1,
	proc_session   => 2,
	proc_tty       => 3,
	proc_tpgid     => 4,
	proc_flags     => 5,
	proc_minflt    => 6,
	proc_cminflt   => 7,
	proc_majflt    => 8,
	proc_cmajflt   => 9,
	proc_utime     => 10,
	proc_stime     => 11,
	proc_cutime    => 12,
	proc_cstime    => 13,
	proc_priority  => 14,
	proc_nice      => 15,
	proc_threads   => 16,
	proc_itrealval => 17,
	proc_starttime => 18,
	proc_vsize     => 19,
	proc_rss       => 20,
	proc_rsslim    => 21,
	proc_startcode => 22,
	proc_endcode   => 23,
	proc_startstk  => 24,
	proc_endstk    => 25,
	proc_kstkesp   => 26,
	proc_kstkeip   => 27,
	proc_signal    => 28,
	proc_sigblock  => 29,
	proc_sigignore => 30,
	proc_sigcatch  => 31,
	proc_wchan     => 32,
	proc_nswap     => 33,
	proc_cnswap    => 34,
	proc_exitsig   => 35,
	proc_processor => 36,
	proc_rtprio    => 37,
	proc_policy    => 38,
	proc_delaybio  => 39,
	proc_vcputime  => 40,
	proc_cvcputime => 41,
};

my $fork       = 1;
my $fake       = 0;

my $interval   = 10;
my $config     = "/usr/local/dh/etc/procwatch.cfg";
my $person_cfg = "/usr/local/dh/etc/procwatch.users";
my $logfile    = "/var/log/procwatch.log";
my $pidfile    = "/var/run/procwatch.pid";

my $highCpuMax = 90;        # processes are niced after this many seconds CPU
my $lifeMax    = 86400 * 7; # processes are killed after living this long

my $uidRamMax  = 150;       # Maximum megabytes RSS per user
my $gidRamMax  = undef;     # Same, but per group (= account)

my $uidProcMax = 25;        # Maximum processes per user
my $gidProcMax = undef;     # Per group

my $niceLevel  = 10;        # Priority to nice processes to

# Process names to particularly avoid killing
my @protectedProcs = qw(
	bash dash csh ksh sh tcsh zsh
	sshd sftp-server proftpd su
	vi vim emacs screen
);

# UIDs and GIDs under 1000 are automatically exempt
my @exemptUids = (
	11139, # mysql
	11142, # Gdnscache
	11143, # Gdnslog
	65534, # nobody
);
my @exemptGids = (
);

# Processes with these names are given extra memory (2x)
my @bonusRam = qw(
	dispatch.fcgi ruby ruby1.8 trac.fcgi rsync git
);

# Processes with these names are given tons of memory (3x)
my @superBonusRam = qw(
	gem bundle cpan cc c++ cc1 cc1plus gcc g++ ffmpeg
);

# Processes matching this description or with this ancestry are exempt
my @exemptionCriteriaSets = (
	{
		system_owned  => 1,
		cmdline_start => ['python', '/usr/local/dh/bin/dupwrap.py'],
	},
);

# Variables used at runtime
my ($uptime, $now, %matchers, %exemptUids, %exemptGids);

# metrics state
my $metrics_state = {};
my $metrics_prefix = "procwatch";
my $metrics_state_file = "/usr/local/dh/etc/procwatch/metrics.state";

# Run!
clear_metric_state();
main();

exit 0;


# ----------------------------------------------------------------------------


# Logging functions

sub msg_debug ($@) {
	my ($fmt, @args) = @_;
	print STDERR "$now procwatch3 DEBUG: " . sprintf($fmt, @args) . "\n";
}

sub msg_info ($@) {
	my ($fmt, @args) = @_;
	print STDERR "$now procwatch3 INFO: " . sprintf($fmt, @args) . "\n";
}

sub msg_warn ($@) {
	my ($fmt, @args) = @_;
	print STDERR "$now procwatch3 ERROR: " . sprintf($fmt, @args) . "\n";
}

sub increment_counter {
	my $uid = shift;
	my $metric = shift;
	my ($user)  = getpwuid($uid);
	my $person = lookup_person_by_uid($uid);
	$metric = join("_",($metrics_prefix, $metric));
	# only returning metrics by person, if we tagged every kill with the user
	# and person.. our metrics db gets huge
	$metrics_state->{$person}->{$metric} += 1
}

sub clear_metric_state {
	truncate $metrics_state_file, 0;
}

sub save_metric_state {
	#print Dumper($metrics_state);
	eval {
		my $json_out = JSON->new->utf8->encode ($metrics_state);
		open(FH, '>', $metrics_state_file) or die $!;
		print FH $json_out;
		close(FH);
	};
	msg_warn "failed to update state file %s", $metrics_state_file, "$!" if $@;
}

sub lookup_person_by_uid {
	my $id = shift;
	my $person = 'not_found';
	my $users = parseProcFile($person_cfg) || return $person;
	$person = $users->{$id} if $users->{$id};
	return $person;
}


# Read in a file. Faster than standard perl open/read.

sub slurp {
	my ($path) = @_;
	my $fd = POSIX::open($path, O_RDONLY);
	return if !defined $fd;
	POSIX::read($fd, my $data, 1048576); # up to 1 MB
	POSIX::close($fd);
	return $data;
}


# Read in a file consisting of lines of the format "Key: Value", like those
# used commonly in procfs.

sub parseProcFile {
	my ($path) = @_;
	my $content = slurp($path);
	my %h = $content =~ /^\s*(\S+):\s+(.*)$/gm;
	return \%h;
}


# Returns a regex which matches any of the strings passed as arguments.

sub match_any {
	my $body = join('|', map { quotemeta $_ } @_);
	return qr{^(?:$body)$};
}


# Returns a hash containing information on the specified pid, or undef if it
# couldn't be loaded.

sub statProc {
	my ($pid) = @_;

	my $stat = slurp("/proc/$pid/stat") or return;
	my ($state, $rest) = $stat =~ m{^\d+ \(([^)]*)\) \S+ (.*)$} or return;
	return if $state eq 'Z'; # zombie processes don't matter
	my @info = split /\s+/, $rest;

	my $status = parseProcFile("/proc/$pid/status") or return;

	my $tty;
	# Get actual tty number from the device number. Look up "new_tty_dev" in
	# the Linux kernel source for the gory details.
	if (($info[proc_tty] & 0xff00) == 0x8800) {
		$tty = (($info[proc_tty] & 255) | ($info[proc_tty] >> 12));
	}

	# UID/GID under 1000 generally indicate the process is running suid/sgid,
	# which we don't care about much really... we mainly just want to know
	# who the user responsible is.
	my ($uid) = grep { $_ >= SYSTEM_UID } split /\s+/, $status->{Uid};
	my ($gid) = grep { $_ >= SYSTEM_GID } split /\s+/, $status->{Gid};

	my $h = {
		pid    => $pid,
		ppid   => $status->{PPid},
		name   => $status->{Name},
		tty    => $tty,
		cpu    => ($info[proc_utime] + $info[proc_stime]) / 100.0, # convert jiffies -> seconds
		nice   => $info[proc_nice],
		start  => $info[proc_starttime],
		ram    => $status->{VmRSS} / 1024.0, # convert KB -> MB
		uid    => $uid,
		gid    => $gid,
	};

	return $h;
}


# Return an array of statProc structures for every non-system process running.

sub crawlProcs {
	my @procs;
	opendir my $dh, "/proc" or die "WTF: couldn't open /proc";
	while (my $pid = readdir $dh) {
		next if $pid =~ /\D/;
		my @stat = stat "/proc/$pid" or next;
		next if $stat[st_uid] < SYSTEM_UID || $stat[st_uid] == NOBODY_UID;
		next if defined $exemptUids{$stat[st_uid]} || defined $exemptGids{$stat[st_gid]};
		my $info = statProc($pid);
		push @procs, $info if $info;
	}
	return @procs;
}


# Take an array of statProc structures and return those that are not exempt

sub filterExemptProcs {
	my @allProcs = @_;
	my @filteredProcs;

	# If a proc is known to be exempt or not exempt, its pid will be a key
	# here.  Therefore, $knownProcDecisions->{$pid} will be 1 if we have
	# decided to exempt it, 0 if we have decided to not exempt it, or undefined
	# if we have not yet made that decision.
	#
	# This assumes that it is sufficient to decide on exemption once per pid
	# during a procwatch run, which could produce incorrect results if a pid is
	# reused by a different process during the course of a run.  It is believed
	# that given runs are fast enough that this is not currently a problem.
	#
	# This also assumes that a traversal of process ancestry will either end in
	# an exempted process or at init.
	my $knownProcDecisions = {
		1 => 0,  # if we hit pid 1 (init), don't exempt the subtree that got to it
	};

	# For each process in @allProcs, come up with a decision as to whether or
	# not it (and probably some amount of process tree above it) should be
	# exempt from killing of processes.  This will allow us to define
	# procExempt: 1 for exempt, 0 for being added to filteredProcs
	for my $proc (@allProcs) {
		my @traversedPids;

		my $currentProc = $proc;
		while (!defined $knownProcDecisions->{$currentProc->{pid}}) {
			if (matchExemptProc($currentProc)) {
				# exempt this process (and in a moment all of the decendants we
				# traversed to get to it)
				$knownProcDecisions->{$currentProc->{pid}} = 1;
			} else {
				# move on to this process' parent, if it is still running,
				# keeping track of this process' pid so we can make this
				# decision for this process later as well.  If we can't get
				# info on the parent process, we're not going anywhere good, so
				# we can make our decision immediately
				my $parentProc = statProc($currentProc->{ppid});
				if ($parentProc && keys %$parentProc) {
					push @traversedPids, $currentProc->{pid};
					$currentProc = $parentProc;
				} else {
					$knownProcDecisions->{$currentProc->{pid}} = 0;
				}
			}
		}

		my $procExempt = $knownProcDecisions->{$currentProc->{pid}};
		# exempt or don't exempt all the pids we worked our way through to make
		# the decision
		for my $pid (@traversedPids) {
			$knownProcDecisions->{$pid} = $procExempt;
		}

		if ($procExempt) {
			msg_info "%s: exempted", procRepr($proc);
		} else {
			push @filteredProcs, $proc;
		}
	}

	return @filteredProcs;
}


# match a process with full exemption criteria, return 1 if it should be
# exempted, nothing otherwise

sub matchExemptProc {
	my ($proc) = @_;

	# The process must match the criteria provided in each
	# $exemptionCriteriaSet, given $pid = $proc->{pid}:
	# - system_owned: whether the process should be a system process or a user
	# process
	# - exe_symlink: if specified, the full path to which /proc/$pid/exe must
	# be a link
	# - cmdline_start: if specified, the first n tokens of /proc/$pid/cmdline,
	# split on null bytes
	for my $exemptionCriteria (@exemptionCriteriaSets) {
		my $shouldBeSystem = $exemptionCriteria->{system_owned};
		# Because of the way statProc filters for uid, a truly system owned
		# process will not have one at all, so if we should be system-run we
		# need to not have a uid here, or if we should not be system-run, we
		# need to have a uid
		next if !($shouldBeSystem xor $proc->{uid});

		my $pid = $proc->{pid};
		if (my $cmdline_start_ref = $exemptionCriteria->{cmdline_start}) {
			my @cmdline_start   = @$cmdline_start_ref;
			my $cmdline         = slurp("/proc/$pid/cmdline") or return;
			my @cmdline         = split /\0/, $cmdline;
			my $cmdline_matches = 1;
			for my $i (0..$#cmdline_start) {
				if ($cmdline[$i] ne $cmdline_start[$i]) {
					$cmdline_matches = 0;
					last;
				}
			}
			next if !$cmdline_matches;
		}

		if ($exemptionCriteria->{exe_symlink}) {
			next if readlink "/proc/$pid/exe" ne $exemptionCriteria->{exe_symlink};
		}

		return 1;
	}
	return;
}


# Return a nice string representation of a given uid / gid

sub reprUid {
	my ($uid) = @_;
	state %uidCache;
	return $uidCache{$uid} ||= scalar(getpwuid $uid) || $uid;
}

sub reprGid {
	my ($gid) = @_;
	state %gidCache;
	return $gidCache{$gid} ||= scalar(getgrgid $gid) || $gid;
}


# Return a nice string representation of a given process structure

sub procRepr {
	my ($proc) = @_;
	my $name = $proc->{name};
	$name =~ s/[^\x20-\x7e]/?/g; # sanitize
	my $user  = reprUid($proc->{uid});
	my $group = reprGid($proc->{gid});
	my $repr = sprintf "PID %d (%s) %s:%s - %.1fMB ram, %.2f sec cpu",
		$proc->{pid}, $name,
		$user, $group,
		$proc->{ram},
		$proc->{cpu};
	$repr .= " (nice $proc->{nice})" if $proc->{nice} != 0;
	$repr .= " [$proc->{mark}]" if $proc->{mark};
	return $repr;
}


# Return a number corresponding to how much we should avoid killing this
# process if its user/group is low on memory. Higher numbers = less likely to
# be killed.

sub procOomValue {
	my ($proc) = @_;
	my $value = 0;

	# Protected processes get killed last
	if ($proc->{name} =~ $matchers{protected}) {
		$proc->{mark} ||= "protected";
		$value += 5000;
	}

	# Interactive processes get an automatic boost to importance
	if (defined $proc->{tty}) {
		$proc->{mark} ||= "interactive";
		$value += 1500; # automatic boost
		my $age = 50 * log($uptime - $proc->{start});
		$value += $age;
	}

	# Processes using more RAM should be killed earlier
	$value -= $proc->{ram};

	# Extra logic for PHP:
	#     If there's a socket at fd 0 (FCGI socket) and nothing in fd 3
	#     (usually a FCGI connection), then this is probably an idle worker,
	#     which we can kill with little impact.
	if (($proc->{name} =~ /^php/)
			&& (readlink("/proc/$proc->{pid}/fd/0") =~ /^socket:/)
			&& (!-e "/proc/$proc->{pid}/fd/3")) {
		$proc->{mark} = "idle php";
		$value -= 3000;
	}

	#msg_debug "proc %s got value %.3f", procRepr($proc), $value;
	return $value;
}


# Kill a process. Arguments are used to construct the error message.

sub killProc {
	my ($proc, $group, $limitType, $excess, $limitVal, $score) = @_;
	my $stderr = readlink "/proc/$proc->{pid}/fd/2";
	if ($fake || kill 9, $proc->{pid}) {
		msg_info "%s: killed for %s %s (total %.1f exceeds limit %.1f; score %.2f)",
			procRepr($proc), $group, $limitType, $excess, $limitVal, $score;
		increment_counter(
			$proc->{uid},
			sprintf("killed_%s_%s_total", $group, $limitType)
		);
		increment_counter(
			$proc->{uid},
			sprintf("killed_interactive_%s_%s_total", $group, $limitType)
		) if $proc->{'mark'} eq 'interactive';

		$proc->{killed}++;
		if (!$fake && $stderr =~ m{^/dev/pts/}) {
			my $fd = POSIX::open($stderr, O_WRONLY | O_NONBLOCK);
			if ($fd >= 0) {
				my $msg = sprintf(
					"Yikes! One of your processes (%s, pid %d) " .
					"was just killed for excessive resource usage.\n" .
					"Please contact DreamHost Support for details.",
					$proc->{name}, $proc->{pid}
				);
				$msg = "\n\n\e[1;33m" . $msg . "\e[0m\n\n\n";
				POSIX::write($fd, $msg, length $msg);
				POSIX::close($fd);
			}
		}
		return 1;
	} else {
		msg_warn "%s: couldn't kill: %s", procRepr($proc), "$!";
		return 1; # It's probably just already exited, which is OK
	}
}


# Renice a process a bit.

sub reniceProc {
	my ($proc) = @_;
	if ($proc->{nice} >= $niceLevel) {
		msg_warn "%s: already niced!", procRepr($proc);
		return;
	}
	if ($fake || setpriority(0, $proc->{pid}, $niceLevel)) {
		msg_info "%s: reniced to +%d", procRepr($proc), $niceLevel;

		increment_counter(
			$proc->{uid},
			"reniced_total"
        );
		return 1;
	} else {
		msg_info "%s: couldn't renice: %s", procRepr($proc), "$!";
		return;
	}
}


# Given an arrayref of processes, kill some if there are too many.

sub killForCount {
	my ($procs, $group) = @_;
	my $limit = ($group eq 'uid') ? $uidProcMax : $gidProcMax;
	return if !defined $limit;
	return if @$procs < $limit;

	my $count = @$procs;
	my %order = map { $_ => procOomValue($_) } @$procs;

	for my $proc (sort { $order{$a} <=> $order{$b} } @$procs) {
		if (killProc($proc, $group, 'count', $count, $limit, $order{$proc})) {
			$count -= 1;

			return if $count < $limit;
		}
	}
}


# Given an arrayref of processes, kill some if they're using too much memory.

sub killForRam {
	my ($procs, $group, $id) = @_;
	my $limit = ($group eq 'uid') ? $uidRamMax : $gidRamMax;
	return if !defined $limit;

	my $total = 0;
	for my $proc (@$procs) {
		my $ram = $proc->{ram};
		$ram /= 2 if $proc->{name} =~ $matchers{bonusRam};
		$ram /= 3 if $proc->{name} =~ $matchers{superBonusRam};
		$total += $ram;
	}

	return if $total < $limit;

	my %order = map { $_ => procOomValue($_) } @$procs;

	for my $proc (sort { $order{$a} <=> $order{$b} } @$procs) {
		if (killProc($proc, $group, 'ram', $total, $limit, $order{$proc})) {
			$total -= $proc->{ram};
			return if $total < $limit;
		}
	}

	msg_warn "Whoa, weird. We didn't manage to kill enough processes there.";
}


sub run {
	$now = scalar localtime;
	$uptime = (split / /, slurp("/proc/uptime"))[0] * 100;

	my @procs = crawlProcs();
	@procs = filterExemptProcs(@procs);

	my %procByUid;
	for my $proc (@procs) {
		push @{$procByUid{$proc->{uid}}}, $proc unless $proc->{killed};
	}
	for my $uid (keys %procByUid) {
		killForCount($procByUid{$uid}, 'uid', $uid);
		killForRam($procByUid{$uid}, 'uid', $uid);
	}

	my %procByGid;
	for my $proc (@procs) {
		push @{$procByGid{$proc->{gid}}}, $proc unless $proc->{killed};
	}
	for my $gid (keys %procByGid) {
		killForCount($procByGid{$gid}, 'gid', $gid);
		killForRam($procByGid{$gid}, 'gid', $gid);
	}

	for my $proc (@procs) {
		next if $proc->{killed};
		my $age = ($uptime - $proc->{start}) / 100.0;
		if (defined $lifeMax && $age > $lifeMax) {
			killProc($proc, 'old', 'age', $age, $lifeMax, '0');
		} elsif (defined $highCpuMax && $proc->{nice} < $niceLevel && $proc->{cpu} > $highCpuMax) {
			reniceProc($proc);
		}
	}
}


sub usage {
	print STDERR q{
Options are:
  --no-fork           Run in the foreground (e.g, for debugging).
  --fake              Don't actually kill/renice anything.
};
	exit 1;
}


sub readConfig {
	my $cfg = parseProcFile($config);

	my %mapScalar = (
		interval   => \$interval,
		logfile    => \$logfile,
		pidfile    => \$pidfile,
		highCpuMax => \$highCpuMax,
		lifeMax    => \$lifeMax,
		uidRamMax  => \$uidRamMax,
		gidRamMax  => \$gidRamMax,
		uidProcMax => \$uidProcMax,
		gidProcMax => \$gidProcMax,
		niceLevel  => \$niceLevel,
	);

	my %mapArray = (
		protectedProcs => \@protectedProcs,
		exemptUids     => \@exemptUids,
		exemptGids     => \@exemptGids,
		bonusRam       => \@bonusRam,
		superBonusRam  => \@superBonusRam,
	);

	while (my ($key, $var) = each %mapScalar) {
		$$var = $cfg->{$key} if exists $cfg->{$key};
	}

	while (my ($key, $var) = each %mapArray) {
		push @$var, split /\s+/, $cfg->{$key} if exists $cfg->{$key};
	}

	$matchers{protected}     = match_any(@protectedProcs);
	$matchers{bonusRam}      = match_any(@bonusRam);
	$matchers{superBonusRam} = match_any(@superBonusRam);

	for my $uid (@exemptUids) { $exemptUids{$uid} = 1; }
	for my $gid (@exemptGids) { $exemptGids{$gid} = 1; }
}

sub main {
	my $P = Getopt::Long::Parser->new(config => ['no_ignore_case']);
	$P->getoptions(
		'fork!'      => \$fork,
		'fake!'      => \$fake,
		'help'       => \&usage,
		'config=s'   => \$config,
	) or usage();

	die "Can't run procwatch as non-root. That's just a bad idea.\n" if $<;

	# Set umask to prevent logs from being world readable
	umask 0077;

	if (my @stat = stat "/usr/local/dh/etc/procwatch/on") {
		if (time - $stat[9] > 86400 * 15) {
			print STDERR "/dh/etc/procwatch/on is over two weeks old -- exiting.\n";
			exit 0;
		}
	} else {
		print STDERR "/dh/etc/procwatch/on doesn't exist -- exiting.\n";
		exit 0;
	}

	readConfig();

	if ($fork) {
		if (open my $fd, '<', $pidfile) {
			my $pid = slurp($pidfile);
			chomp $pid;
			my $cmd = slurp("/proc/$pid/cmdline");
			if ($cmd =~ /procwatch/) {
				print STDERR "Procwatch already running (pid $pid) -- exiting.\n";
				exit 1;
			}
		}
		chdir "/";

		exit if fork > 0;
		open STDIN, '<', '/dev/null';
		open STDOUT, '>>', $logfile;
		open STDERR, '>>', $logfile;

		open my $pid, '>', $pidfile;
		print $pid "$$\n";
		close $pid;
	}

	$now = scalar localtime;

	msg_info("procwatch3 starting up");

	for (;;) {
		eval {
			run();
			save_metric_state();
		};
		if (my $err = $@) {
			chomp $err;
			msg_warn("Error: $err");
		}

		sleep $interval;
	}
}

1;
