# HG changeset patch
# User Oleg Oshmyan <chortos@inbox.lv>
# Date 1294940855 -7200
# Node ID 69eadc60f4e29288fe8218c90e5bfb5c1fbbb90d
# Parent  0e5ae28e0b2bdbe8e28c5e4a164f43ef2ad3e1ba
Memory limit is now applied to the RSS when os.wait4 is available

diff -r 0e5ae28e0b2b -r 69eadc60f4e2 problem.py
--- a/problem.py	Sat Jan 08 16:03:35 2011 +0200
+++ b/problem.py	Thu Jan 13 19:47:35 2011 +0200
@@ -1,5 +1,5 @@
 #! /usr/bin/env python
-# Copyright (c) 2010 Chortos-2 <chortos@inbox.lv>
+# Copyright (c) 2010-2011 Chortos-2 <chortos@inbox.lv>
 
 from __future__ import division, with_statement
 
@@ -143,6 +143,8 @@
 					verdict = 'canceled by the user'
 				except testcases.TimeLimitExceeded:
 					verdict = 'time limit exceeded'
+				except testcases.MemoryLimitExceeded:
+					verdict = 'memory limit exceeded'
 				except testcases.WrongAnswer:
 					e = sys.exc_info()[1]
 					if e.comment:
diff -r 0e5ae28e0b2b -r 69eadc60f4e2 testcases.py
--- a/testcases.py	Sat Jan 08 16:03:35 2011 +0200
+++ b/testcases.py	Thu Jan 13 19:47:35 2011 +0200
@@ -1,5 +1,5 @@
 #! /usr/bin/env python
-# Copyright (c) 2010 Chortos-2 <chortos@inbox.lv>
+# Copyright (c) 2010-2011 Chortos-2 <chortos@inbox.lv>
 
 # TODO: copy the ansfile if not options.erase even if no validator is used
 
@@ -272,6 +272,32 @@
 		else:
 			return call_real(*args, **kwargs)
 
+# Emulate memory limits on platforms compatible with 4.3BSD but not XSI
+# I say 'emulate' because the OS will allow excessive memory usage
+# anyway; Upreckon will just treat the test case as not passed.
+# To do this, we not only require os.wait4 to be present but also
+# assume things about the implementation of subprocess.Popen.
+try:
+	def waitpid_emu(pid, options, _wait4=os.wait4):
+		global last_rusage
+		pid, status, last_rusage = _wait4(pid, options)
+		return pid, status
+	_waitpid = os.waitpid
+	os.waitpid = waitpid_emu
+	try:
+		defaults = Popen._internal_poll.__func__.__defaults__
+	except AttributeError:
+		# Python 2.5
+		defaults = Popen._internal_poll.im_func.func_defaults
+	i = defaults.index(_waitpid)
+	defaults = defaults[:i] + (waitpid_emu,) + defaults[i+1:]
+	try:
+		Popen._internal_poll.__func__.__defaults__ = defaults
+	except AttributeError:
+		pass
+	Popen._internal_poll.im_func.func_defaults = defaults
+except (AttributeError, ValueError):
+	pass
 
 __all__ = ('TestCase', 'load_problem', 'TestCaseNotPassed',
            'TimeLimitExceeded', 'CanceledByUser', 'WrongAnswer',
@@ -285,6 +311,7 @@
 
 class TestCaseNotPassed(Exception): __slots__ = ()
 class TimeLimitExceeded(TestCaseNotPassed): __slots__ = ()
+class MemoryLimitExceeded(TestCaseNotPassed): __slots__ = ()
 class CanceledByUser(TestCaseNotPassed): __slots__ = ()
 
 class WrongAnswer(TestCaseNotPassed):
@@ -539,6 +566,8 @@
 					pass
 		case.open_infile()
 		case.time_started = None
+		global last_rusage
+		last_rusage = None
 		if case.problem.config.stdio:
 			if options.erase and not case.validator or not case.problem.config.inname:
 				# TODO: re-use the same file name if possible
@@ -601,6 +630,8 @@
 										time.sleep(.001)
 						if config.globalconf.force_zero_exitcode and case.process.returncode or case.process.returncode < 0:
 							raise NonZeroExitCode(case.process.returncode)
+						if case.maxmemory and last_rusage and last_rusage.ru_maxrss > case.maxmemory * (1024 if sys.platform != 'darwin' else 1048576):
+							raise MemoryLimitExceeded
 						callback()
 						case.has_called_back = True
 						outfile.seek(0)
@@ -653,6 +684,8 @@
 							time.sleep(.001)
 			if config.globalconf.force_zero_exitcode and case.process.returncode or case.process.returncode < 0:
 				raise NonZeroExitCode(case.process.returncode)
+			if case.maxmemory and last_rusage and last_rusage.ru_maxrss > case.maxmemory * (1024 if sys.platform != 'darwin' else 1048576):
+				raise MemoryLimitExceeded
 			callback()
 			case.has_called_back = True
 			with open(case.problem.config.outname, 'rU') as output: