Mercurial > ~astiob > upreckon > hgweb

changeset 80:809b77302b21

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Win32-specific module with memory and CPU time limits The Win32-specific implementation of call() and friends now lives in module win32, looks clean and in addition is able to enforce memory and CPU time limits on NT kernels, in particular on Windows 2000 and up asking the system to terminate the process as soon as or (in the case of CPU time) almost as soon as the limits are broken. According to my observations, malloc() in the limited process does not return NULL when memory usage is close to the limit and instead crashes the process (which Upreckon happily translates into 'memory limit exceeded'). The catch is that the module is not actually used yet; coming soon.
author Oleg Oshmyan <chortos@inbox.lv>
date Wed, 16 Feb 2011 00:01:33 +0000
parents ee8a99dcaaed
children 24752db487c5
files compat.py win32.py
diffstat 2 files changed, 505 insertions(+), 1 deletions(-) [+]
line wrap: on
line diff
--- a/compat.py	Mon Jan 17 11:11:01 2011 +0000
+++ b/compat.py	Wed Feb 16 00:01:33 2011 +0000
@@ -44,7 +44,7 @@
 	import __builtin__ as builtins
 
 pseudobuiltins = ('say', 'basestring', 'range', 'map', 'zip', 'filter', 'next',
-                  'items', 'keys', 'values', 'zip_longest', 'callable')
+                  'items', 'keys', 'values', 'zip_longest', 'callable', 'ceil')
 __all__ = pseudobuiltins + ('ABCMeta', 'abstractmethod', 'CompatBuiltins')
 
 try:
@@ -198,6 +198,13 @@
 except AttributeError:
 	values = dict.values
 
+from math import ceil
+if not isinstance(ceil(0), int):
+	def ceil(x):
+		y = int(x)
+		if y < x: y += 1
+		return y
+
 try:
 	# Python 3
 	from itertools import zip_longest
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/win32.py	Wed Feb 16 00:01:33 2011 +0000
@@ -0,0 +1,497 @@
+# Copyright (c) 2011 Chortos-2 <chortos@inbox.lv>
+
+from __future__ import division, with_statement
+
+try:
+	from compat import *
+except ImportError:
+	import __main__
+	__main__.import_error(sys.exc_info()[1])
+
+from __main__ import clock
+from ctypes import *
+from ctypes.wintypes import *
+from subprocess import Popen
+
+# Defaults that may be overwritten by values from _subprocess
+INFINITE = -1
+STD_INPUT_HANDLE = -10
+WAIT_OBJECT_0 = 0
+
+try:
+	from _subprocess import *
+except ImportError:
+	pass
+
+try:
+	from numbers import Integral
+except ImportError:
+	Integral = int, long
+
+try:
+	from collections import namedtuple
+except ImportError:
+	from operator import itemgetter
+	class ProcessTimes(tuple):
+		__slots__ = ()
+		__new__ = lambda cls, kernel, user: tuple.__new__(cls, (kernel, user))
+		__getnewargs__ = lambda self: tuple(self)
+		kernel, user = (property(itemgetter(i)) for i in (0, 1))
+else:
+	ProcessTimes = namedtuple('ProcessTimes', 'kernel user')
+
+__all__ = 'call', 'kill', 'terminate'
+
+
+# Automatically convert _subprocess handle objects into low-level HANDLEs
+# and replicate their functionality for our own use
+try:
+	_subprocess_handle = type(GetCurrentProcess())
+except NameError:
+	_subprocess_handle = Integral
+class Handle(object):
+	@staticmethod
+	def from_param(handle):
+		if isinstance(handle, (_subprocess_handle, Integral)):
+			return HANDLE(int(handle))
+		elif isinstance(handle, Handle):
+			return HANDLE(handle.handle)
+		elif isinstance(handle, HANDLE):
+			return handle
+		else:
+			raise TypeError('cannot convert %s to a handle' %
+			                type(handle).__name__)
+	
+	__slots__ = 'handle'
+	
+	def __init__(self, handle):
+		if isinstance(handle, Integral):
+			self.handle = handle
+		elif isinstance(handle, HANDLE):
+			self.handle = handle.value
+		elif isinstance(handle, Handle):
+			self.handle = handle.handle
+		elif isinstance(handle, _subprocess_handle):
+			handle = HANDLE(int(handle))
+			flags = DWORD()
+			try:
+				if windll.kernel32.GetHandleInformation(handle, byref(flags)):
+					flags = flags.value
+				else:
+					flags = 0
+			except AttributeError:
+				# Available on NT 3.51 and up, NT line only
+				flags = 0
+			proc = HANDLE(int(GetCurrentProcess()))
+			handle = DuplicateHandle(proc, handle, proc, 0, flags & 1, 2)
+			self.handle = handle.Detach()
+		else:
+			raise TypeError("Handle() argument must be a handle, not '%s'" %
+			                type(name).__name__)
+	
+	def __int__(self):
+		return int(self.handle)
+	
+	def Detach(self):
+		handle = self.handle
+		self.handle = None
+		return handle
+	
+	# This is also __del__, so only locals are accessed
+	def Close(self, _CloseHandle=windll.kernel32.CloseHandle, _HANDLE=HANDLE):
+		if self.handle:
+			_CloseHandle(_HANDLE(self.handle))
+			self.handle = None
+	__del__ = Close
+
+CHAR = c_char
+INVALID_HANDLE_VALUE = HANDLE(-1).value
+LPDWORD = POINTER(DWORD)
+LPFILETIME = POINTER(FILETIME)
+SIZE_T = ULONG_PTR = WPARAM
+ULONGLONG = c_ulonglong
+
+try:
+	unicode
+except NameError:
+	LPCTSTR = LPCWSTR
+	unisuffix = 'W'
+else:
+	LPCTSTR = LPCSTR
+	unisuffix = 'A'
+
+
+prototype = WINFUNCTYPE(BOOL, Handle,
+                        LPFILETIME, LPFILETIME, LPFILETIME, LPFILETIME)
+flags = ((1, 'process'),
+         (2, 'creation'), (2, 'exit'), (2, 'kernel'), (2, 'user'))
+try:
+	GetProcessTimes = prototype(('GetProcessTimes', windll.kernel32), flags)
+except AttributeError:
+	# Available on NT 3.5 and up, NT line only
+	GetProcessTimes = None
+else:
+	def errcheck(result, func, args):
+		if not result: raise WinError()
+		ftimes = [t.dwHighDateTime << 32 | t.dwLowDateTime for t in args[3:]]
+		kernel = ftimes[0] / 10000000
+		user   = ftimes[1] / 10000000
+		return ProcessTimes(kernel, user)
+	GetProcessTimes.errcheck = errcheck
+
+
+class PROCESS_MEMORY_COUNTERS(Structure):
+	_fields_ = (('cb', DWORD),
+	            ('PageFaultCount', DWORD),
+	            ('PeakWorkingSetSize', SIZE_T),
+	            ('WorkingSetSize', SIZE_T),
+	            ('QuotaPeakPagedPoolUsage', SIZE_T),
+	            ('QuotaPagedPoolUsage', SIZE_T),
+	            ('QuotaPeakNonPagedPoolUsage', SIZE_T),
+	            ('QuotaNonPagedPoolUsage', SIZE_T),
+	            ('PagefileUsage', SIZE_T),
+	            ('PeakPagefileUsage', SIZE_T))
+
+prototype = WINFUNCTYPE(BOOL, Handle, POINTER(PROCESS_MEMORY_COUNTERS), DWORD)
+flags = ((1, 'process'), (2, 'counters'),
+         (5, 'cb', sizeof(PROCESS_MEMORY_COUNTERS)))
+try:
+	GetProcessMemoryInfo = prototype(('GetProcessMemoryInfo', windll.psapi),
+	                                 flags)
+except AttributeError:
+	# Available on NT 4.0 and up, NT line only
+	GetProcessMemoryInfo = None
+else:
+	def errcheck(result, func, args):
+		if not result: raise WinError()
+		return args
+	GetProcessMemoryInfo.errcheck = errcheck
+
+
+class _uChar_union(Union):
+	_fields_ = (('UnicodeChar', WCHAR),
+	            ('AsciiChar', CHAR))
+
+class KEY_EVENT_RECORD(Structure):
+	_fields_ = (('bKeyDown', BOOL),
+	            ('wRepeatCount', WORD),
+	            ('wVirtualKeyCode', WORD),
+	            ('wVirtualScanCode', WORD),
+	            ('uChar', _uChar_union),
+	            ('dwControlKeyState', DWORD))
+
+RIGHT_ALT_PRESSED  = 0x001
+LEFT_ALT_PRESSED   = 0x002
+RIGHT_CTRL_PRESSED = 0x004
+LEFT_CTRL_PRESSED  = 0x008
+SHIFT_PRESSED      = 0x010
+NUMLOCK_ON         = 0x020
+SCROLLLOCK_ON      = 0x040
+CAPSLOCK_ON        = 0x080
+ENHANCED_KEY       = 0x100
+
+class _Event_union(Union):
+	_fields_ = ('KeyEvent', KEY_EVENT_RECORD),
+
+class INPUT_RECORD(Structure):
+	_fields_ = (('EventType', WORD),
+	            ('Event', _Event_union))
+
+KEY_EVENT                = 0x01
+MOUSE_EVENT              = 0x02
+WINDOW_BUFFER_SIZE_EVENT = 0x04
+MENU_EVENT               = 0x08
+FOCUS_EVENT              = 0x10
+
+prototype = WINFUNCTYPE(BOOL, Handle, POINTER(INPUT_RECORD), DWORD, LPDWORD)
+flags = (1, 'input'), (2, 'buffer'), (5, 'length', 1), (2, 'number_read')
+ReadConsoleInput = prototype(('ReadConsoleInputA', windll.kernel32), flags)
+def errcheck(result, func, args):
+	if not result: raise WinError()
+	return args[1] if args[3] else None
+ReadConsoleInput.errcheck = errcheck
+
+
+prototype = WINFUNCTYPE(BOOL, Handle)
+flags = (1, 'input'),
+FlushConsoleInputBuffer = prototype(('FlushConsoleInputBuffer',
+                                     windll.kernel32), flags)
+def errcheck(result, func, args):
+	if not result: raise WinError()
+FlushConsoleInputBuffer.errcheck = errcheck
+
+
+prototype = WINFUNCTYPE(BOOL, Handle, DWORD)
+flags = (1, 'console'), (1, 'mode')
+SetConsoleMode = prototype(('SetConsoleMode', windll.kernel32), flags)
+def errcheck(result, func, args):
+	if not result: raise WinError()
+SetConsoleMode.errcheck = errcheck
+
+ENABLE_PROCESSED_INPUT = 0x001
+ENABLE_LINE_INPUT      = 0x002
+ENABLE_ECHO_INPUT      = 0x004
+ENABLE_WINDOW_INPUT    = 0x008
+ENABLE_MOUSE_INPUT     = 0x010
+ENABLE_INSERT_MODE     = 0x020
+ENABLE_QUICK_EDIT_MODE = 0x040
+ENABLE_EXTENDED_FLAGS  = 0x080
+
+ENABLE_PROCESSED_OUTPUT   = 1
+ENABLE_WRAP_AT_EOL_OUTPUT = 2
+
+
+prototype = WINFUNCTYPE(HANDLE, c_void_p, LPCTSTR)
+flags = (5, 'attributes'), (1, 'name')
+try:
+	CreateJobObject = prototype(('CreateJobObject'+unisuffix, windll.kernel32),
+	                            flags)
+except AttributeError:
+	# Available on 2000 and up, NT line only
+	CreateJobObject = lambda name: None
+else:
+	def errcheck(result, func, args):
+		if not result: raise WinError()
+		return Handle(result)
+	CreateJobObject.errcheck = errcheck
+
+
+prototype = WINFUNCTYPE(BOOL, Handle, Handle)
+flags = (1, 'job'), (1, 'handle')
+try:
+	AssignProcessToJobObject = prototype(('AssignProcessToJobObject',
+	                                      windll.kernel32), flags)
+except AttributeError:
+	# Available on 2000 and up, NT line only
+	AssignProcessToJobObject = lambda job, handle: None
+else:
+	def errcheck(result, func, args):
+		if not result: raise WinError()
+	AssignProcessToJobObject.errcheck = errcheck
+
+
+class JOBOBJECT_BASIC_LIMIT_INFORMATION(Structure):
+	_fields_ = (('PerProcessUserTimeLimit', LARGE_INTEGER),
+	            ('PerJobUserTimeLimit', LARGE_INTEGER),
+	            ('LimitFlags', DWORD),
+	            ('MinimumWorkingSetSize', SIZE_T),
+	            ('MaximumWorkingSetSize', SIZE_T),
+	            ('ActiveProcessLimit', DWORD),
+	            ('Affinity', ULONG_PTR),
+	            ('PriorityClass', DWORD),
+	            ('SchedulingClass', DWORD))
+
+JOB_OBJECT_LIMIT_WORKINGSET                 = 0x0001
+JOB_OBJECT_LIMIT_PROCESS_TIME               = 0x0002
+JOB_OBJECT_LIMIT_JOB_TIME                   = 0x0004
+JOB_OBJECT_LIMIT_ACTIVE_PROCESS             = 0x0008
+JOB_OBJECT_LIMIT_AFFINITY                   = 0x0010
+JOB_OBJECT_LIMIT_PRIORITY_CLASS             = 0x0020
+JOB_OBJECT_LIMIT_PRESERVE_JOB_TIME          = 0x0040
+JOB_OBJECT_LIMIT_SCHEDULING_CLASS           = 0x0080
+JOB_OBJECT_LIMIT_PROCESS_MEMORY             = 0x0100
+JOB_OBJECT_LIMIT_JOB_MEMORY                 = 0x0200
+JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION = 0x0400
+JOB_OBJECT_LIMIT_BREAKAWAY_OK               = 0x0800
+JOB_OBJECT_LIMIT_SILENT_BREAKAWAY_OK        = 0x1000
+JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE          = 0x2000
+JOB_OBJECT_LIMIT_SUBSET_AFFINITY            = 0x4000
+
+class IO_COUNTERS(Structure):
+	_fields_ = (('ReadOperationCount', ULONGLONG),
+	            ('WriteOperationCount', ULONGLONG),
+	            ('OtherOperationCount', ULONGLONG),
+	            ('ReadTransferCount', ULONGLONG),
+	            ('WriteTransferCount', ULONGLONG),
+	            ('OtherTransferCount', ULONGLONG))
+
+class JOBOBJECT_EXTENDED_LIMIT_INFORMATION(Structure):
+	_fields_ = (('BasicLimitInformation', JOBOBJECT_BASIC_LIMIT_INFORMATION),
+	            ('IoInfo', IO_COUNTERS),
+	            ('ProcessMemoryLimit', SIZE_T),
+	            ('JobMemoryLimit', SIZE_T),
+	            ('PeakProcessMemoryUsed', SIZE_T),
+	            ('PeakJobMemoryUsed', SIZE_T))
+
+prototype = WINFUNCTYPE(BOOL, HANDLE, c_int, c_void_p, DWORD)
+flags = (1, 'job'), (1, 'infoclass'), (1, 'info'), (1, 'infosize')
+try:
+	_setjobinfo = prototype(('SetInformationJobObject',windll.kernel32), flags)
+except AttributeError:
+	# Available on 2000 and up, NT line only
+	SetInformationJobObject = lambda job, infoclass, info: None
+else:
+	def errcheck(result, func, args):
+		if not result: raise WinError()
+	_setjobinfo.errcheck = errcheck
+	def SetInformationJobObject(job, infoclass, info):
+		return _setjobinfo(job, infoclass, info, sizeof(info))
+
+(
+	JobObjectBasicAccountingInformation,
+	JobObjectBasicLimitInformation,
+	JobObjectBasicProcessIdList,
+	JobObjectBasicUIRestrictions,
+	JobObjectSecurityLimitInformation,   
+	JobObjectEndOfJobTimeInformation,
+	JobObjectAssociateCompletionPortInformation,
+	JobObjectBasicAndIoAccountingInformation,
+	JobObjectExtendedLimitInformation,
+	JobObjectJobSetInformation,
+	MaxJobObjectInfoClass
+) = range(1, 12)
+
+
+prototype = WINFUNCTYPE(DWORD, DWORD, POINTER(HANDLE), BOOL, DWORD)
+flags = (1, 'count'), (1, 'handles'), (1, 'wait_all'), (1, 'milliseconds')
+_wait_multiple = prototype(('WaitForMultipleObjects', windll.kernel32), flags)
+def errcheck(result, func, args):
+	if result == WAIT_FAILED: raise WinError()
+	return args
+_wait_multiple.errcheck = errcheck
+def WaitForMultipleObjects(handles, wait_all, timeout):
+	n = len(handles)
+	handles = (Handle.from_param(handle) for handle in handles)
+	timeout = ceil(timeout * 1000)
+	return _wait_multiple(n, (HANDLE * n)(*handles), wait_all, timeout)
+
+# WAIT_OBJECT_0 defined at the top of the file
+WAIT_ABANDONED_0 = 0x00000080
+WAIT_TIMEOUT     = 0x00000102
+WAIT_FAILED      = 0xFFFFFFFF
+
+
+try:
+	_wait_single = WaitForSingleObject
+except NameError:
+	prototype = WINFUNCTYPE(DWORD, Handle, DWORD)
+	flags = (1, 'handle'), (1, 'milliseconds')
+	_wait_single = prototype(('WaitForSingleObject', windll.kernel32), flags)
+	def errcheck(result, func, args):
+		if result == WAIT_FAILED: raise WinError()
+		return args
+	_wait_single.errcheck = errcheck
+def WaitForSingleObject(handle, timeout):
+	return _wait_single(handle, ceil(timeout * 1000))
+
+
+try:
+	GetStdHandle
+except NameError:
+	prototype = WINFUNCTYPE(HANDLE, DWORD)
+	flags = (1, 'which'),
+	GetStdHandle = prototype(('GetStdHandle', windll.kernel32), flags)
+	def errcheck(result, func, args):
+		if result == INVALID_HANDLE_VALUE: raise WinError()
+		return args if result else None
+	GetStdHandle.errcheck = errcheck
+
+
+try:
+	TerminateProcess
+except NameError:
+	prototype = WINFUNCTYPE(BOOL, Handle, UINT)
+	flags = (1, 'process'), (1, 'exitcode')
+	TerminateProcess = prototype(('TerminateProcess', windll.kernel32), flags)
+	def errcheck(result, func, args):
+		if not result: raise WinError()
+	TerminateProcess.errcheck = errcheck
+
+
+# Do not show error messages due to errors in the program being tested
+try:
+	errmode = ctypes.windll.kernel32.GetErrorMode()
+except AttributeError:
+	# GetErrorMode is available on Vista/2008 and up
+	errmode = ctypes.windll.kernel32.SetErrorMode(0)
+ctypes.windll.kernel32.SetErrorMode(errmode | 0x8003)
+
+stdin = GetStdHandle(STD_INPUT_HANDLE)
+try:
+	SetConsoleMode(stdin, ENABLE_PROCESSED_INPUT)
+except WindowsError:
+	console_input = False
+else:
+	console_input = True
+	FlushConsoleInputBuffer(stdin)
+
+def kill(process):
+	try:
+		process.terminate()
+	except AttributeError:
+		TerminateProcess(process._handle)
+terminate = kill
+
+def call(*args, **kwargs):
+	case = kwargs.pop('case')
+	job = CreateJobObject(None)
+	flags = 0
+	if case.maxcputime:
+		flags |= JOB_OBJECT_LIMIT_PROCESS_TIME
+	if case.maxmemory:
+		flags |= JOB_OBJECT_LIMIT_PROCESS_MEMORY
+	limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION(
+		JOBOBJECT_BASIC_LIMIT_INFORMATION(
+			PerProcessUserTimeLimit=ceil((case.maxcputime or 0)*10000000),
+			LimitFlags=flags,
+		),
+		ProcessMemoryLimit=ceil((case.maxmemory or 0)*1048576),
+	)
+	SetInformationJobObject(job, JobObjectExtendedLimitInformation, limits)
+	try:
+		case.process = Popen(*args, **kwargs)
+	except OSError:
+		raise CannotStartTestee(sys.exc_info()[1])
+	case.time_started = clock()
+	AssignProcessToJobObject(job, case.process._handle)
+	if not console_input:
+		if case.maxwalltime:
+			if (WaitForSingleObject(case.process._handle, case.maxwalltime) !=
+			    WAIT_OBJECT_0):
+				raise TimeLimitExceeded
+		else:
+			case.process.wait()
+	else:
+		handles = stdin, case.process._handle
+		if case.maxwalltime:
+			time_end = clock() + case.maxwalltime
+			while case.process.poll() is None:
+				remaining = time_end - clock()
+				if remaining > 0:
+					if (WaitForMultipleObjects(handles, False, remaining) ==
+					    WAIT_OBJECT_0):
+						ir = ReadConsoleInput(stdin)
+						if (ir and
+						    ir.EventType == 1 and
+						    ir.Event.KeyEvent.bKeyDown and
+						    ir.Event.KeyEvent.wVirtualKeyCode == 27):
+							raise CanceledByUser
+				else:
+					raise TimeLimitExceeded
+		else:
+			while case.process.poll() is None:
+				if (WaitForMultipleObjects(handles, False, INFINITE) ==
+				    WAIT_OBJECT_0):
+					ir = ReadConsoleInput(stdin)
+					if (ir and
+					    ir.EventType == 1 and
+					    ir.Event.KeyEvent.bKeyDown and
+					    ir.Event.KeyEvent.wVirtualKeyCode == 27):
+						raise CanceledByUser
+	case.time_stopped = clock()
+	if case.maxcputime and GetProcessTimes:
+		try:
+			times = GetProcessTimes(case.process._handle)
+		except WindowsError:
+			pass
+		else:
+			if times.kernel + times.user > case.maxcputime:
+				raise TimeLimitExceeded
+	if case.maxmemory and GetProcessMemoryInfo:
+		try:
+			counters = GetProcessMemoryInfo(case.process._handle)
+		except WindowsError:
+			pass
+		else:
+			if counters.PeakPagefileUsage > case.maxmemory * 1048576:
+				raise MemoryLimitExceeded
\ No newline at end of file