first commit

scripts/run-guard.py

@@ -0,0 +1,159 @@
+#!/usr/bin/env python
+
+import os, sys, string, signal, resource, time, getopt, pickle
+
+class RunGuard:
+
+    def __init__(self):
+        self.timelimit = 1
+        self.memlimit = 65536 * 1024
+        self.timetime = 5
+        self.args = None
+        self.writeto = None
+        self.nproc = 1
+        self.ofile = 32
+        self.ldpreload = None
+        self.rundir = None
+
+        self.usepickle = False
+        self.memrss = 0
+        self.memdata = 0
+        self.memstack = 0
+        self.timeused = 0
+        self.exitcode = 0
+        self.sig = 0
+
+    def run(self):
+        self.parse_opts()
+
+        self.childpid = os.fork()
+        if self.childpid == 0:
+            self.execute()
+        else:
+            self.monitor()
+
+        self.write_result()
+
+    def parse_opts(self):
+        try:
+            pos = sys.argv.index('-x')
+            optlist, self.args = getopt.gnu_getopt(sys.argv[:pos], 'e:t:m:d:o:T:p')
+            self.args = sys.argv[pos:]
+        except ValueError:
+            optlist, self.args = getopt.gnu_getopt(sys.argv, 'e:t:m:d:o:T:p')
+        for o, v in optlist:
+            if o == '-e':
+                self.nproc += string.atoi(v)
+            if o == '-t':
+                self.timelimit = string.atoi(v)
+            if o == '-m':
+                self.memlimit = string.atoi(v) * 1024
+            if o == '-d':
+                self.rundir = v
+            if o == '-o':
+                self.writeto = v
+            if o == '-T':
+                self.timetime = string.atoi(v)
+            if o == '-p':
+                self.usepickle = True
+
+        v = os.getenv('GUARD_RLIMIT_OFILE')
+        if v: self.ofile = string.atoi(v)
+        self.ldpreload = os.getenv('GUARD_LD_PRELOAD')
+
+    def execute(self):
+        if self.rundir != None and os.path.isdir(self.rundir):
+            os.chdir(self.rundir)
+
+        if self.nproc:
+            resource.setrlimit(resource.RLIMIT_NPROC, (self.nproc, self.nproc))
+        if self.ofile:
+            resource.setrlimit(resource.RLIMIT_OFILE, (self.ofile, self.ofile))
+        if self.ldpreload:
+            os.putenv('LD_PRELOAD', self.ldpreload)
+
+        resource.setrlimit(resource.RLIMIT_CPU, (self.timelimit, self.timelimit))
+        resource.setrlimit(resource.RLIMIT_AS, (self.memlimit, self.memlimit))
+        os.execv(self.args[1], self.args[1:])
+        sys.exit(127) # exit with JGE
+
+    def monitor(self):
+        pid = 0
+        remaintime = self.timelimit * self.timetime
+        while remaintime > 0:
+            pid, status, ru = os.wait4(self.childpid, os.WNOHANG)
+            self._get_memused()
+            if pid > 0: break
+            time.sleep(0.05)
+            remaintime -= 0.05
+
+        while pid == 0:
+            try:
+                os.kill(self.childpid, signal.SIGKILL)
+            except OSError, e:
+                pass
+            pid, status, ru = os.wait4(self.childpid, os.WNOHANG)
+            time.sleep(0.1)
+
+        if os.WIFEXITED(status):
+            self.exitcode = os.WEXITSTATUS(status)
+
+        if os.WIFSIGNALED(status):
+            self.sig = os.WTERMSIG(status)
+
+        self.timeused = ru[0] + ru[1]
+
+    def _get_memused(self):
+        procdir = '/proc/%d' % self.childpid
+        if os.path.isdir(procdir):
+            cmdline = file(procdir + '/cmdline', 'r').readlines()
+
+            # do not get memory usage of this script after just fork
+            if len(cmdline) > 0 and \
+               string.strip(cmdline[0], '\0') != \
+               string.join(self.args[1:], '\0'):
+                return
+
+            procstatus = file(procdir + '/status', 'r')
+            rss = 0; data = 0; stack = 0
+            for line in procstatus:
+                n = line[0:6]
+                if n == 'VmRSS:':
+                    rss = string.atoi(line[7:-3])
+                if n == 'VmData':
+                    data = string.atoi(line[8:-3])
+                if n == 'VmStk:':
+                    stack = string.atoi(line[7:-3])
+            self.memrss = max(self.memrss, rss)
+            if self.memdata + self.memstack < data + stack:
+                self.memdata = data
+                self.memstack = stack
+        return
+
+    def write_result(self):
+        if self.writeto == None:
+            f = sys.stdout
+        else:
+            f = file(self.writeto, 'w')
+
+        if self.usepickle:
+            obj = { 'exitcode' : self.exitcode,
+                    'sig' : self.sig,
+                    'timeused' : self.timeused,
+                    'memrss' : self.memrss,
+                    'memdata' : self.memdata,
+                    'memstack' : self.memstack }
+            pickle.dump(obj, f)
+        else:
+            print >>f, "exitcode: %d" % self.exitcode
+            print >>f, "sig: %d" % self.sig
+            print >>f, "time: %.3f" % self.timeused
+            print >>f, "rss: %d" % self.memrss
+            print >>f, "data: %d" % self.memdata
+            print >>f, "stack: %d" % self.memstack
+
+        if self.writeto != None: f.close()
+
+if __name__ == '__main__':
+    os.umask(0002)
+    RunGuard().run()