run program with time limit

coderunner.py

# -*- coding: utf-8 -*-


import os
import time
import signal
import pwd

_DEFAULT_ENV = {"PATH": "/sbin:/usr/sbin:/bin:/usr/bin", "LANG": "en_US.UTF-8"}

def _coderun_exec_impl(cmd, username=None, workdirectory=None, env=None):

	print "INFO: cmd=%r, username=%r, wd=%r, env=%r" % (cmd, username, workdirectory, env,)

	# {{{ change work directory
	if workdirectory is None:
		workdirectory = os.getcwd()
	else:
		try:
			os.chdir(workdirectory)
		except Exception as e:
			print "WARN: cannot switch working directory to [%r]: %r" % (workdirectory, e,)
	# }}} change work directory

	# {{{ switch user account
	if username is not None:
		try:
			udb = pwd.getpwnam(username)

			os.initgroups(username, udb.pw_gid)
			os.setresgid(udb.pw_gid, udb.pw_gid, udb.pw_gid)
			os.setresuid(udb.pw_uid, udb.pw_uid, udb.pw_uid)
		except Exception as e:
			print "WARN: cannot switch user ID as [%r]: %r" % (username, e,)
	# }}} switch user account

	# {{{ setup environment variable
	if env is None:
		env = _DEFAULT_ENV
	# }}} setup environment variable

	os.setpgrp()

	os.execve(cmd[0], cmd, env)

	os._exit(os.EX_SOFTWARE)

	return -1
# ### def _coderun_exec_impl


def coderun(cmd, username=None, workdirectory=None, env=None, limit_exec_time=None):

	child_pid = os.fork()

	if 0 == child_pid:
		_coderun_exec_impl(cmd, username=None, workdirectory=None, env=None)
		os._exit(os.EX_SOFTWARE)
		return None

	start_tstamp = time.time()
	end_tstamp = None

	ret = None
	if limit_exec_time is None:
		ret = os.waitpid(child_pid, 0)
		end_tstamp = time.time()
	else:
		exp_tstamp = start_tstamp + limit_exec_time
		force_finish_step = 0
		while end_tstamp is None:
			ret = os.waitpid(child_pid, os.WNOHANG)
			if (ret is None) or (0 == ret[0]):
				# {{{ run too long
				if exp_tstamp < time.time():
					if 1 == force_finish_step:
						os.killpg(child_pid, signal.SIGINT)
					elif 3 == force_finish_step:
						os.killpg(child_pid, signal.SIGTERM)
					elif 13 == force_finish_step:
						os.killpg(child_pid, signal.SIGKILL)
					elif 20 == force_finish_step:
						break

					force_finish_step = force_finish_step + 1
				# }}} run too long

				time.sleep(3)
			else:
				end_tstamp = time.time()
				break

	retcode = None
	retsig = None
	if (ret is not None):
		v = ret[1]
		retcode = (v >> 8) & 0xFF
		retsig = v & 0xFF

	consume_time = None
	if end_tstamp is not None:
		consume_time = end_tstamp - start_tstamp

	print "INFO: start-time: %r, end-time: %r, used: %r" % (start_tstamp, end_tstamp, consume_time,)

	return (retcode, retsig,)
# ### def coderun



# vim: ts=4 sw4= ai nowrap