From 622fa7ced0aad18e28184bbb36e50b7dd4c1f9a5 Mon Sep 17 00:00:00 2001
From: Simon Walker <s.r.walker101@googlemail.com>
Date: Mon, 3 Jun 2024 12:34:09 +0100
Subject: [PATCH] ls-restart: Wait for health check and ready scripts

---
 home/bin/ls-restart | 67 ++++++++++++++++++++++++++++++++++++---------
 1 file changed, 54 insertions(+), 13 deletions(-)

diff --git a/home/bin/ls-restart b/home/bin/ls-restart
index 565be90d..4d68eb92 100755
--- a/home/bin/ls-restart
+++ b/home/bin/ls-restart
@@ -1,10 +1,18 @@
 #!/usr/bin/env python
 
 import sys
+import logging
+import json
+import time
 import os
 import subprocess as sp
+from typing import Any
+import urllib.request
 import argparse
 
+logging.basicConfig(level=logging.WARNING)
+logger = logging.getLogger("ls-restart")
+
 
 def in_tmux() -> bool:
     return "TMUX" in os.environ
@@ -40,35 +48,67 @@ class DockerClient:
         cmd = ["docker", "kill", "--signal", "USR1", container_name]
         sp.check_call(cmd, stdout=sp.PIPE)
         if wait:
-            cls.wait_for_container(container_name)
+            logger.info("Waiting for new container to be healthy")
+            HealthCheck().wait()
+            
+class HealthCheck:
+    def wait(self):
+        while True:
+            if self.localstack_health():
+                break
+            time.sleep(0.5)
 
+    def localstack_health(self) -> bool:
+        return self._health_endpoint() and self._ready_init_scripts_endpoint()
 
-    @classmethod
-    def wait_for_container(cls, container_name: str):
-        process = sp.Popen(["docker", "logs", "-f", container_name, "--since", "0m"], stdout=sp.PIPE, stderr=sp.STDOUT)
-        if process.stdout is None:
-            raise RuntimeError("Programming error: no stdout")
+    def _health_endpoint(self) -> bool:
+        url = "http://localhost.localstack.cloud:4566/_localstack/health"
+        logger.debug("checking health of '%s'", url)
+        try:
+            self._json_request(url)
+        except Exception:
+            return False
 
-        while True:
-            line = process.stdout.readline()
-            if not line:
-                break
-            line = line.decode("utf-8").strip()
-            if line == "Ready.":
-                break
+        return True
+
+    def _ready_init_scripts_endpoint(self) -> bool:
+        url = "http://localhost.localstack.cloud:4566/_localstack/init"
+        logger.debug("checking health of '%s'", url)
+        try:
+            res = self._json_request(url)
+        except Exception:
+            return False
+
+        return res.get("completed", {}).get("READY", False)
 
+    def _json_request(self, url: str) -> Any:
+        res = urllib.request.urlopen(url)
+        if res.status != 200:
+            raise ValueError(f"Bad status: {res.status}")
+
+        data = res.read()
+        return json.loads(data.decode("utf-8"))
 
 def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("--no-wait", action="store_true", default=False)
+    parser.add_argument("-v", "--verbose", action="count", default=0)
     opts, args = parser.parse_known_args()
 
+    if opts.verbose == 1:
+        logger.setLevel(logging.INFO)
+    elif opts.verbose > 1:
+        logger.setLevel(logging.DEBUG)
+
     container_name = "localstack-main"
     if args:
         container_name = args[0]
 
+    logger.info("Working with container '%s'", container_name)
+
     running_containers = DockerClient.list_containers()
     if container_name in running_containers:
+        logger.info("Restarting container '%s'", container_name)
         DockerClient.restart_container(container_name, wait=not opts.no_wait)
         return 0
 
@@ -77,6 +117,7 @@ def main() -> int:
         print("No container name chosen or found, exiting", file=sys.stderr)
         return 1
 
+    logger.info("Restarting container '%s'", chosen_container_name)
     DockerClient.restart_container(chosen_container_name, wait=not opts.no_wait)
     return 0