Repository: getnikola/nikola

Issue:
Improve performance of `nikola auto`
- [x] add support for `-n` to pass to `nikola build`
- [x] ~~show build output as it happens~~ doesn’t seem to work well
- [x] limit the number of reload requests per second
- [x] try to detect dead sockets (they send a `WARNING: asyncio: socket.send() raised exception.` now)


Discussion:

Resolving patch:
diff --git a/nikola/plugins/command/auto/__init__.py b/nikola/plugins/command/auto/__init__.py
--- a/nikola/plugins/command/auto/__init__.py
+++ b/nikola/plugins/command/auto/__init__.py
@@ -34,12 +34,13 @@
 import stat
 import subprocess
 import sys
-
+import typing
 import webbrowser
+
 import pkg_resources
 
 from nikola.plugin_categories import Command
-from nikola.utils import dns_sd, req_missing, get_theme_path
+from nikola.utils import dns_sd, req_missing, get_theme_path, makedirs
 
 try:
     import aiohttp
@@ -58,6 +59,8 @@
     Observer = None
 
 LRJS_PATH = os.path.join(os.path.dirname(__file__), 'livereload.js')
+REBUILDING_REFRESH_DELAY = 0.35
+IDLE_REFRESH_DELAY = 0.05
 
 if sys.platform == 'win32':
     asyncio.set_event_loop(asyncio.ProactorEventLoop())
@@ -71,6 +74,7 @@ class CommandAuto(Command):
     doc_purpose = "builds and serves a site; automatically detects site changes, rebuilds, and optionally refreshes a browser"
     dns_sd = None
     delta_last_rebuild = datetime.timedelta(milliseconds=100)
+    web_runner = None  # type: web.AppRunner
 
     cmd_options = [
         {
@@ -79,7 +83,7 @@ class CommandAuto(Command):
             'long': 'port',
             'default': 8000,
             'type': int,
-            'help': 'Port number)',
+            'help': 'Port number',
         },
         {
             'name': 'address',
@@ -112,13 +116,31 @@ class CommandAuto(Command):
             'type': bool,
             'help': 'Disable the server, automate rebuilds only'
         },
+        {
+            'name': 'process',
+            'short': 'n',
+            'long': 'process',
+            'default': 0,
+            'type': int,
+            'help': 'Number of subprocesses (nikola build argument)'
+        },
+        {
+            'name': 'parallel-type',
+            'short': 'P',
+            'long': 'parallel-type',
+            'default': 'process',
+            'type': str,
+            'help': "Parallelization mode ('process' or 'thread', nikola build argument)"
+        },
     ]
 
     def _execute(self, options, args):
         """Start the watcher."""
         self.sockets = []
         self.rebuild_queue = asyncio.Queue()
+        self.reload_queue = asyncio.Queue()
         self.last_rebuild = datetime.datetime.now()
+        self.is_rebuilding = False
 
         if aiohttp is None and Observer is None:
             req_missing(['aiohttp', 'watchdog'], 'use the "auto" command')
@@ -135,9 +157,9 @@ def _execute(self, options, args):
         if self.site.configuration_filename != 'conf.py':
             self.nikola_cmd.append('--conf=' + self.site.configuration_filename)
 
-        # Run an initial build so we are up-to-date (synchronously)
-        self.logger.info("Rebuilding the site...")
-        subprocess.call(self.nikola_cmd)
+        if options and options.get('process'):
+            self.nikola_cmd += ['--process={}'.format(options['process']),
+                                '--parallel-type={}'.format(options['parallel-type'])]
 
         port = options and options.get('port')
         self.snippet = '''<script>document.write('<script src="http://'
@@ -166,6 +188,9 @@ def _execute(self, options, args):
         watched.add(pkg_resources.resource_filename('nikola', ''))
 
         out_folder = self.site.config['OUTPUT_FOLDER']
+        if not os.path.exists(out_folder):
+            makedirs(out_folder)
+
         if options and options.get('browser'):
             browser = True
         else:
@@ -178,14 +203,6 @@ def _execute(self, options, args):
 
         host = options['address'].strip('[').strip(']') or dhost
 
-        # Set up asyncio server
-        webapp = web.Application()
-        webapp.router.add_get('/livereload.js', self.serve_livereload_js)
-        webapp.router.add_get('/robots.txt', self.serve_robots_txt)
-        webapp.router.add_route('*', '/livereload', self.websocket_handler)
-        resource = IndexHtmlStaticResource(True, self.snippet, '', out_folder)
-        webapp.router.register_resource(resource)
-
         # Prepare asyncio event loop
         # Required for subprocessing to work
         loop = asyncio.get_event_loop()
@@ -197,8 +214,10 @@ def _execute(self, options, args):
         self.has_server = not options['no-server']
 
         if self.has_server:
-            handler = webapp.make_handler()
-            srv = loop.run_until_complete(loop.create_server(handler, host, port))
+            loop.run_until_complete(self.set_up_server(host, port, out_folder))
+
+        # Run an initial build so we are up-to-date. The server is running, but we are not watching yet.
+        loop.run_until_complete(self.run_initial_rebuild())
 
         self.wd_observer = Observer()
         # Watch output folders and trigger reloads
@@ -208,12 +227,12 @@ def _execute(self, options, args):
         # Watch input folders and trigger rebuilds
         for p in watched:
             if os.path.exists(p):
-                self.wd_observer.schedule(NikolaEventHandler(self.run_nikola_build, loop), p, recursive=True)
+                self.wd_observer.schedule(NikolaEventHandler(self.queue_rebuild, loop), p, recursive=True)
 
         # Watch config file (a bit of a hack, but we need a directory)
         _conf_fn = os.path.abspath(self.site.configuration_filename or 'conf.py')
         _conf_dn = os.path.dirname(_conf_fn)
-        self.wd_observer.schedule(ConfigEventHandler(_conf_fn, self.run_nikola_build, loop), _conf_dn, recursive=False)
+        self.wd_observer.schedule(ConfigEventHandler(_conf_fn, self.queue_rebuild, loop), _conf_dn, recursive=False)
         self.wd_observer.start()
 
         win_sleeper = None
@@ -239,8 +258,6 @@ def _execute(self, options, args):
             loop.close()
             return
 
-        host, port = srv.sockets[0].getsockname()
-
         if options['ipv6'] or '::' in host:
             server_url = "http://[{0}]:{1}/".format(host, port)
         else:
@@ -257,7 +274,8 @@ def _execute(self, options, args):
         # Run the event loop forever and handle shutdowns.
         try:
             # Run rebuild queue
-            queue_future = asyncio.ensure_future(self.run_rebuild_queue())
+            rebuild_queue_fut = asyncio.ensure_future(self.run_rebuild_queue())
+            reload_queue_fut = asyncio.ensure_future(self.run_reload_queue())
 
             self.dns_sd = dns_sd(port, (options['ipv6'] or '::' in host))
             loop.run_forever()
@@ -269,17 +287,35 @@ def _execute(self, options, args):
                 win_sleeper.cancel()
             if self.dns_sd:
                 self.dns_sd.Reset()
-            queue_future.cancel()
-            srv.close()
-            loop.run_until_complete(srv.wait_closed())
-            loop.run_until_complete(webapp.shutdown())
-            loop.run_until_complete(handler.shutdown(5.0))
-            loop.run_until_complete(webapp.cleanup())
+            rebuild_queue_fut.cancel()
+            reload_queue_fut.cancel()
+            loop.run_until_complete(self.web_runner.cleanup())
             self.wd_observer.stop()
             self.wd_observer.join()
         loop.close()
 
-    async def run_nikola_build(self, event):
+    async def set_up_server(self, host: str, port: int, out_folder: str) -> None:
+        """Set up aiohttp server and start it."""
+        webapp = web.Application()
+        webapp.router.add_get('/livereload.js', self.serve_livereload_js)
+        webapp.router.add_get('/robots.txt', self.serve_robots_txt)
+        webapp.router.add_route('*', '/livereload', self.websocket_handler)
+        resource = IndexHtmlStaticResource(True, self.snippet, '', out_folder)
+        webapp.router.register_resource(resource)
+        webapp.on_shutdown.append(self.remove_websockets)
+
+        self.web_runner = web.AppRunner(webapp)
+        await self.web_runner.setup()
+        website = web.TCPSite(self.web_runner, host, port)
+        await website.start()
+
+    async def run_initial_rebuild(self) -> None:
+        """Run an initial rebuild."""
+        await self._rebuild_site()
+        # If there are any clients, have them reload the root.
+        await self._send_reload_command(self.site.config['INDEX_FILE'])
+
+    async def queue_rebuild(self, event) -> None:
         """Rebuild the site."""
         # Move events have a dest_path, some editors like gedit use a
         # move on larger save operations for write protection
@@ -299,34 +335,61 @@ async def run_nikola_build(self, event):
         self.logger.debug('Queuing rebuild from {0}'.format(event_path))
         await self.rebuild_queue.put((datetime.datetime.now(), event_path))
 
-    async def run_rebuild_queue(self):
+    async def run_rebuild_queue(self) -> None:
         """Run rebuilds from a queue (Nikola can only build in a single instance)."""
         while True:
             date, event_path = await self.rebuild_queue.get()
             if date < (self.last_rebuild + self.delta_last_rebuild):
                 self.logger.debug("Skipping rebuild from {0} (within delta)".format(event_path))
                 continue
-            self.last_rebuild = datetime.datetime.now()
+            await self._rebuild_site(event_path)
+
+    async def _rebuild_site(self, event_path: typing.Optional[str] = None) -> None:
+        """Rebuild the site."""
+        self.is_rebuilding = True
+        self.last_rebuild = datetime.datetime.now()
+        if event_path:
             self.logger.info('REBUILDING SITE (from {0})'.format(event_path))
-            p = await asyncio.create_subprocess_exec(*self.nikola_cmd, stderr=subprocess.PIPE)
-            exit_code = await p.wait()
-            error = await p.stderr.read()
-            errord = error.decode('utf-8')
-
-            if exit_code != 0:
-                self.logger.error(errord)
-                await self.send_to_websockets({'command': 'alert', 'message': errord})
+        else:
+            self.logger.info('REBUILDING SITE')
+
+        p = await asyncio.create_subprocess_exec(*self.nikola_cmd, stderr=subprocess.PIPE)
+        exit_code = await p.wait()
+        out = (await p.stderr.read()).decode('utf-8')
+
+        if exit_code != 0:
+            self.logger.error("Rebuild failed\n" + out)
+            await self.send_to_websockets({'command': 'alert', 'message': out})
+        else:
+            self.logger.info("Rebuild successful\n" + out)
+
+        self.is_rebuilding = False
+
+    async def run_reload_queue(self) -> None:
+        """Send reloads from a queue to limit CPU usage."""
+        while True:
+            p = await self.reload_queue.get()
+            self.logger.info('REFRESHING: {0}'.format(p))
+            await self._send_reload_command(p)
+            if self.is_rebuilding:
+                await asyncio.sleep(REBUILDING_REFRESH_DELAY)
             else:
-                self.logger.info("Rebuild successful\n" + errord)
+                await asyncio.sleep(IDLE_REFRESH_DELAY)
+
+    async def _send_reload_command(self, path: str) -> None:
+        """Send a reload command."""
+        await self.send_to_websockets({'command': 'reload', 'path': path, 'liveCSS': True})
 
-    async def reload_page(self, event):
+    async def reload_page(self, event) -> None:
         """Reload the page."""
         # Move events have a dest_path, some editors like gedit use a
         # move on larger save operations for write protection
-        event_path = event.dest_path if hasattr(event, 'dest_path') else event.src_path
+        if event:
+            event_path = event.dest_path if hasattr(event, 'dest_path') else event.src_path
+        else:
+            event_path = self.site.config['OUTPUT_FOLDER']
         p = os.path.relpath(event_path, os.path.abspath(self.site.config['OUTPUT_FOLDER'])).replace(os.sep, '/')
-        self.logger.info('REFRESHING: {0}'.format(p))
-        await self.send_to_websockets({'command': 'reload', 'path': p, 'liveCSS': True})
+        await self.reload_queue.put(p)
 
     async def serve_livereload_js(self, request):
         """Handle requests to /livereload.js and serve the JS file."""
@@ -359,7 +422,7 @@ async def websocket_handler(self, request):
                     await ws.send_json(response)
                 elif message['command'] != 'info':
                     self.logger.warning("Unknown command in message: {0}".format(message))
-            elif msg.type == aiohttp.WSMsgType.CLOSED:
+            elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSING):
                 break
             elif msg.type == aiohttp.WSMsgType.CLOSE:
                 self.logger.debug("Closing WebSocket")
@@ -376,7 +439,13 @@ async def websocket_handler(self, request):
 
         return ws
 
-    async def send_to_websockets(self, message):
+    async def remove_websockets(self, app) -> None:
+        """Remove all websockets."""
+        for ws in self.sockets:
+            await ws.close()
+        self.sockets.clear()
+
+    async def send_to_websockets(self, message: dict) -> None:
         """Send a message to all open WebSockets."""
         to_delete = []
         for ws in self.sockets:
@@ -386,6 +455,9 @@ async def send_to_websockets(self, message):
 
             try:
                 await ws.send_json(message)
+                if ws._close_code:
+                    await ws.close()
+                    to_delete.append(ws)
             except RuntimeError as e:
                 if 'closed' in e.args[0]:
                     self.logger.warning("WebSocket {0} closed uncleanly".format(ws))
@@ -397,7 +469,7 @@ async def send_to_websockets(self, message):
             self.sockets.remove(ws)
 
 
-async def windows_ctrlc_workaround():
+async def windows_ctrlc_workaround() -> None:
     """Work around bpo-23057."""
     # https://bugs.python.org/issue23057
     while True:
@@ -416,13 +488,12 @@ def __init__(self, modify_html=True, snippet="</head>", *args, **kwargs):
         self.snippet = snippet
         super().__init__(*args, **kwargs)
 
-    async def _handle(self, request):
+    async def _handle(self, request: web.Request) -> web.Response:
         """Handle incoming requests (pass to handle_file)."""
         filename = request.match_info['filename']
-        ret = await self.handle_file(request, filename)
-        return ret
+        return await self.handle_file(request, filename)
 
-    async def handle_file(self, request, filename, from_index=None):
+    async def handle_file(self, request: web.Request, filename: str, from_index=None) -> web.Response:
         """Handle file requests."""
         try:
             filepath = self._directory.joinpath(filename).resolve()
@@ -471,7 +542,7 @@ async def handle_file(self, request, filename, from_index=None):
 
         return ret
 
-    def transform_html(self, text):
+    def transform_html(self, text: str) -> str:
         """Apply some transforms to HTML content."""
         # Inject livereload.js
         text = text.replace('</head>', self.snippet, 1)


Test patch: