Use non-blocking file/subprocess IO

2 files changed, 55 insertions, 23 deletions

diff --git a/texgen.py b/texgen.py
index 61d90f7..eda0e0b 100644
--- a/texgen.py
+++ b/texgen.py
@@ -1,10 +1,11 @@
 import asyncio
+import asyncio.subprocess as subprocess
 import contextlib
 import hashlib
 import os
 import shutil
-import typing
-import subprocess
+
+from aiofile import async_open
 
 
 @contextlib.contextmanager
@@ -40,26 +41,28 @@ class TexGenerator:
         with temp_dir(os.path.join(self._temp_path, os.urandom(24).hex())) as workdir:
             job_name = 'texput'
             tex_file_path = os.path.join(workdir, f'{job_name}.tex')
-            with open(tex_file_path, 'w', encoding='utf-8') as f:
-                f.write(tex_source)
-            with subprocess.Popen(
-                    [
-                        'xelatex',
-                        '-interaction=nonstopmode',
-                        '-halt-on-error',
-                        tex_file_path,
-                    ],
-                    cwd=workdir,
-                    stdin=subprocess.PIPE,
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.STDOUT,
-                    text=True,
-                    shell=True,
-            ) as proc:
-                stdout, stderr = proc.communicate(input='', timeout=self._task_timeout)
-                print('STDOUT', stdout)
-                print('STDERR', stderr)
-                if proc.returncode != 0:
-                    raise TexGenerationError(f'xelatex process exited with non-zero code {proc.returncode}')
+            async with async_open(tex_file_path, 'w', encoding='utf-8') as f:
+                await f.write(tex_source)
+            proc = await subprocess.create_subprocess_exec(
+                *[
+                    'xelatex',
+                    '-interaction=nonstopmode',
+                    '-halt-on-error',
+                    tex_file_path,
+                ],
+                cwd=workdir,
+                stdin=subprocess.PIPE,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+            try:
+                await asyncio.wait_for(proc.wait(), timeout=10)
+            except TimeoutError:
+                raise TexGenerationError('xelatex timed out')
+            stdout, stderr = [(await x.read()).decode() for x in (proc.stdout, proc.stderr)]
+            print('STDOUT', stdout)
+            print('STDERR', stderr)
+            if proc.returncode != 0:
+                raise TexGenerationError(f'xelatex process exited with non-zero code {proc.returncode}')
             os.rename(os.path.join(workdir, f'{job_name}.pdf'), cache_file_path)
         return cache_file_path
diff --git a/web.py b/web.py
index e2d7fdf..82126b7 100644
--- a/web.py
+++ b/web.py
@@ -1,7 +1,11 @@
+import asyncio
+import os
 import re
 
+import uvicorn
 from fastapi import FastAPI, Response
 from fastapi.responses import FileResponse
+from uvicorn.loops.auto import auto_loop_setup
 
 import htmlcache
 import texgen
@@ -32,3 +36,28 @@ async def get_utaten_lyric_tex(item_id: str):
         return Response(content=tex, media_type='application/x-tex')
     except htmlcache.TexSourceGenerationError as e:
         return Response(content=str(e), status_code=503)
+
+
[email protected]_event("startup")
+async def startup_event():
+    pass
+
+
+def setup_loop():
+    if os.name == 'nt':
+        # use ProactorEventLoop to support async subprocess on Windows
+        print('Driving event loop with IOCP.')
+        loop = asyncio.ProactorEventLoop()
+        asyncio.set_event_loop(loop)
+    auto_loop_setup()
+
+
+if __name__ == '__main__':
+    setup_loop()
+    uvicorn.run(
+        'web:app',
+        host='127.0.0.1',
+        port=8000,
+        log_level='info',
+        loop='none',  # use custom loop initializer
+    )