From 7a84972770ac9cce8507fb7ebeeb63d9cbb9e979 Mon Sep 17 00:00:00 2001
From: Pascal Vizeli <pascal.vizeli@syshack.ch>
Date: Tue, 17 Oct 2017 16:04:37 +0200
Subject: [PATCH] Better close/loop handling (#221)

* Better close/loop handling

* Update bootstrap.py

* Update __main__.py

* Update core.py

* Update __main__.py

* Update __main__.py

* Update supervisor.py

* Update supervisor.py

* Update const.py

* fix lint
---
 hassio/__main__.py        | 26 ++++++++++++++------------
 hassio/bootstrap.py       |  8 ++++----
 hassio/const.py           |  2 --
 hassio/core.py            |  7 +++----
 hassio/dock/supervisor.py |  3 +--
 5 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/hassio/__main__.py b/hassio/__main__.py
index 44b9dcd7e..5ab7639c5 100644
--- a/hassio/__main__.py
+++ b/hassio/__main__.py
@@ -13,11 +13,12 @@ _LOGGER = logging.getLogger(__name__)
 # pylint: disable=invalid-name
 if __name__ == "__main__":
     bootstrap.initialize_logging()
+    loop = asyncio.get_event_loop()
 
     if not bootstrap.check_environment():
-        exit(1)
+        sys.exit(1)
 
-    loop = asyncio.get_event_loop()
+    # init executor pool
     executor = ThreadPoolExecutor(thread_name_prefix="SyncWorker")
     loop.set_default_executor(executor)
 
@@ -27,19 +28,20 @@ if __name__ == "__main__":
 
     bootstrap.migrate_system_env(config)
 
-    _LOGGER.info("Run Hassio setup")
+    _LOGGER.info("Setup HassIO")
     loop.run_until_complete(hassio.setup())
 
-    _LOGGER.info("Start Hassio")
     loop.call_soon_threadsafe(loop.create_task, hassio.start())
-    loop.call_soon_threadsafe(bootstrap.reg_signal, loop, hassio)
+    loop.call_soon_threadsafe(bootstrap.reg_signal, loop)
 
-    _LOGGER.info("Run Hassio loop")
-    loop.run_forever()
-
-    _LOGGER.info("Cleanup system")
-    executor.shutdown(wait=False)
-    loop.close()
+    try:
+        _LOGGER.info("Run HassIO")
+        loop.run_forever()
+    finally:
+        _LOGGER.info("Stopping HassIO")
+        loop.run_until_complete(hassio.stop())
+        executor.shutdown(wait=False)
+        loop.close()
 
     _LOGGER.info("Close Hassio")
-    sys.exit(hassio.exit_code)
+    sys.exit(0)
diff --git a/hassio/bootstrap.py b/hassio/bootstrap.py
index daa09371a..4a9468b63 100644
--- a/hassio/bootstrap.py
+++ b/hassio/bootstrap.py
@@ -123,22 +123,22 @@ def check_environment():
     return True
 
 
-def reg_signal(loop, hassio):
+def reg_signal(loop):
     """Register SIGTERM, SIGKILL to stop system."""
     try:
         loop.add_signal_handler(
-            signal.SIGTERM, lambda: loop.create_task(hassio.stop()))
+            signal.SIGTERM, lambda: loop.call_soon(loop.stop))
     except (ValueError, RuntimeError):
         _LOGGER.warning("Could not bind to SIGTERM")
 
     try:
         loop.add_signal_handler(
-            signal.SIGHUP, lambda: loop.create_task(hassio.stop()))
+            signal.SIGHUP, lambda: loop.call_soon(loop.stop))
     except (ValueError, RuntimeError):
         _LOGGER.warning("Could not bind to SIGHUP")
 
     try:
         loop.add_signal_handler(
-            signal.SIGINT, lambda: loop.create_task(hassio.stop()))
+            signal.SIGINT, lambda: loop.call_soon(loop.stop))
     except (ValueError, RuntimeError):
         _LOGGER.warning("Could not bind to SIGINT")
diff --git a/hassio/const.py b/hassio/const.py
index cfa8a9a1e..891e8df13 100644
--- a/hassio/const.py
+++ b/hassio/const.py
@@ -20,8 +20,6 @@ RUN_WATCHDOG_HOMEASSISTANT_DOCKER = 15
 RUN_WATCHDOG_HOMEASSISTANT_API = 300
 RUN_CLEANUP_API_SESSIONS = 900
 
-RESTART_EXIT_CODE = 100
-
 FILE_HASSIO_ADDONS = Path(HASSIO_DATA, "addons.json")
 FILE_HASSIO_CONFIG = Path(HASSIO_DATA, "config.json")
 FILE_HASSIO_HOMEASSISTANT = Path(HASSIO_DATA, "homeassistant.json")
diff --git a/hassio/core.py b/hassio/core.py
index be7d3cdcb..6150b0e1b 100644
--- a/hassio/core.py
+++ b/hassio/core.py
@@ -177,7 +177,7 @@ class HassIO(object):
             if self.homeassistant.version == 'landingpage':
                 self.loop.create_task(self.homeassistant.install())
 
-    async def stop(self, exit_code=0):
+    async def stop(self):
         """Stop a running orchestration."""
         # don't process scheduler anymore
         self.scheduler.suspend = True
@@ -185,7 +185,6 @@ class HassIO(object):
         # process stop tasks
         self.websession.close()
         self.homeassistant.websession.close()
-        await asyncio.wait([self.api.stop(), self.dns.stop()], loop=self.loop)
 
-        self.exit_code = exit_code
-        self.loop.stop()
+        # process async stop tasks
+        await asyncio.wait([self.api.stop(), self.dns.stop()], loop=self.loop)
diff --git a/hassio/dock/supervisor.py b/hassio/dock/supervisor.py
index 96e9fc2b4..5e4591b73 100644
--- a/hassio/dock/supervisor.py
+++ b/hassio/dock/supervisor.py
@@ -6,7 +6,6 @@ import docker
 
 from .interface import DockerInterface
 from .util import docker_process
-from ..const import RESTART_EXIT_CODE
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -52,7 +51,7 @@ class DockerSupervisor(DockerInterface):
         _LOGGER.info("Update supervisor docker to %s:%s", self.image, tag)
 
         if await self.loop.run_in_executor(None, self._install, tag):
-            self.loop.create_task(self.stop_callback(RESTART_EXIT_CODE))
+            self.loop.call_later(2, self.loop.stop)
             return True
 
         return False