Set large buffers on generators; DMOJ#295

Hackmas · Dec 8, 2017 · 32060ba · 32060ba
1 parent dd3455e
commit 32060ba
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 8 deletions.
diff --git a/dmoj/cptbox/sandbox.py b/dmoj/cptbox/sandbox.py
@@ -145,7 +145,8 @@ def __init__(self, debugger, _, args, executable=None, security=None, time=0, me
         self._executable = executable or _find_exe(args[0])
         self._args = args
         self._chdir = cwd
-        self._env = [utf8bytes('%s=%s' % i) for i in six.iteritems(env if env is not None else os.environ)]
+        self._env = [utf8bytes('%s=%s' % (arg, val))
+                     for arg, val in six.iteritems(env if env is not None else os.environ) if val is not None]
         self._time = time
         self._wall_time = time * 3 if wall_time is None else wall_time
         self._cpu_time = time + 5 if time else 0

diff --git a/dmoj/executors/mixins.py b/dmoj/executors/mixins.py
@@ -98,8 +98,7 @@ def get_security(self, launch_kwargs=None):
             def get_fs(self):
                 name = self.get_executor_name()
                 fs = BASE_FILESYSTEM + self.fs + env.get('extra_fs', {}).get(name, [])
-                if self.unbuffered:
-                    fs += [re.escape(self._file('setbufsize.so')) + '$']
+                fs += [re.escape(self._file('setbufsize.so')) + '$']
                 return fs
 
             def get_allowed_syscalls(self):
@@ -111,13 +110,19 @@ def get_address_grace(self):
             def get_env(self):
                 env = {'LANG': 'C'}
                 if self.unbuffered:
-                    unbuf = self._file('setbufsize.so')
-                    shutil.copyfile(setbufsize_path, unbuf)
-                    env['LD_PRELOAD'] = unbuf
                     env['CPTBOX_STDOUT_BUFFER_SIZE'] = 0
                 return env
 
             def launch(self, *args, **kwargs):
+                agent = self._file('setbufsize.so')
+                shutil.copyfile(setbufsize_path, agent)
+                env = {
+                    'LD_PRELOAD': agent,
+                    'CPTBOX_STDOUT_BUFFER_SIZE': kwargs.get('stdout_buffer_size'),
+                    'CPTBOX_STDERR_BUFFER_SIZE': kwargs.get('stderr_buffer_size'),
+                }
+                env.update(self.get_env())
+
                 return SecurePopen([utf8bytes(a) for a in self.get_cmdline() + list(args)],
                                    executable=utf8bytes(self.get_executable()),
                                    security=self.get_security(launch_kwargs=kwargs),
@@ -126,7 +131,7 @@ def launch(self, *args, **kwargs):
                                    time=kwargs.get('time'), memory=kwargs.get('memory'),
                                    wall_time=kwargs.get('wall_time'),
                                    stderr=(PIPE if kwargs.get('pipe_stderr', False) else None),
-                                   env=self.get_env(), cwd=utf8bytes(self._dir), nproc=self.get_nproc())
+                                   env=env, cwd=utf8bytes(self._dir), nproc=self.get_nproc())
 except ImportError:
     pass
 

diff --git a/dmoj/problem.py b/dmoj/problem.py
@@ -203,7 +203,10 @@ def _run_generator(self, gen, args=None):
         # e.g., an untrusted generator may be one generated via site-managed data by an
         # arbitrary user, who shouldn't be allowed to do arbitrary things on the host machine
         if use_sandbox:
-            proc = executor.launch(time=time_limit, memory=memory_limit, pipe_stderr=True)
+            # setting large buffers is really important, because otherwise stderr is unbuffered
+            # and the generator begins calling into cptbox Python code really frequently
+            proc = executor.launch(time=time_limit, memory=memory_limit, pipe_stderr=True,
+                                   stderr_buffer_size=65536, stdout_buffer_size=65536)
         else:
             proc = executor.launch_unsafe(*args, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                           stderr=subprocess.PIPE)