fix: Docker pull hanging on Windows

Author: rboni-dk

dk-installer.py

@@ -7,6 +7,7 @@
 import datetime
 import functools
 import hashlib
+import io
 import ipaddress
 import json
 import logging
@@ -121,7 +122,6 @@ def collect_user_input(fields: list[str]) -> dict[str, str]:
                 if value := input(f"{CONSOLE.MARGIN}{field.capitalize()!s: >20}: "):
                     res[field] = value
     except KeyboardInterrupt:
-        print("")  # Moving the cursor back to the start
         raise AbortAction
     finally:
         CONSOLE.space()
@@ -162,44 +162,56 @@ def get_installer_version():
         return "N/A"
 
 
-class StreamIterator:
-    def __init__(self, proc, stream, file_path):
-        self.proc = proc
-        self.stream = stream
-        self.file_path = file_path
-        self.file = None
-        self.bytes_written = 0
-
-    def __iter__(self):
-        return self
-
-    def __next__(self):
-        for return_anyway in (False, True):
-            # We poll the process status before consuming the stream to make sure the StopIteration condition
-            # is not vulnerable to a race condition.
-            ret = self.proc.poll()
-            line = self.stream.readline()
-            if line:
-                if not self.file:
-                    self.file = open(self.file_path, "wb")
-                self.file.write(line)
-                self.bytes_written += len(line)
-                return line
-            if ret is not None and not line:
-                raise StopIteration
-            if not return_anyway:
-                time.sleep(0.1)
-        return line
+@contextlib.contextmanager
+def stream_iterator(
+    proc: subprocess.Popen, stdout_file_path: pathlib.Path, stderr_file_path: pathlib.Path, timeout: float = 1.0
+):
+    stdout_buffer = io.TextIOWrapper(io.BytesIO())
+    stderr_buffer = io.TextIOWrapper(io.BytesIO())
 
-    def __enter__(self):
-        return self
+    def _iter():
+        proc_exited = False
+        stdout_read_pos = 0
+        stderr_read_pos = 0
+        while not proc_exited:
+            try:
+                stdout_partial, stderr_partial = proc.communicate(timeout=timeout)
+            except subprocess.TimeoutExpired as exc:
+                stdout_partial = exc.output
+                stderr_partial = exc.stderr
+            else:
+                proc_exited = True
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        for _ in iter(self):
+            if stdout_partial is not None:
+                stdout_buffer.buffer.seek(0)
+                stdout_buffer.buffer.write(stdout_partial)
+            if stderr_partial is not None:
+                stderr_buffer.buffer.seek(0)
+                stderr_buffer.buffer.write(stderr_partial)
+
+            while True:
+                stdout_buffer.seek(stdout_read_pos)
+                stdout_line = stdout_buffer.readline()
+                stdout_read_pos = stdout_buffer.tell()
+                stderr_buffer.seek(stderr_read_pos)
+                stderr_line = stderr_buffer.readline()
+                stderr_read_pos = stderr_buffer.tell()
+
+                if not (stdout_line or stderr_line):
+                    break
+
+                yield stdout_line.strip(), stderr_line.strip()
+
+    iterator = _iter()
+    try:
+        yield iterator
+    finally:
+        for _ in iterator:
             pass
-        if self.file:
-            self.file.close()
-        return False
+        if stdout_buffer.buffer.tell():
+            stdout_file_path.write_bytes(stdout_buffer.buffer.getvalue())
+        if stderr_buffer.buffer.tell():
+            stderr_file_path.write_bytes(stderr_buffer.buffer.getvalue())
 
 
 #
@@ -616,7 +628,7 @@ def run_cmd_retries(self, *cmd, timeout, retries, raise_on_non_zero=True, env=No
         cmd_fail_exception = None
         while retries > 0:
             try:
-                with self.start_cmd(*cmd, raise_on_non_zero=raise_on_non_zero, env=env, **popen_args) as (proc, *_):
+                with self.start_cmd(*cmd, raise_on_non_zero=raise_on_non_zero, env=env, **popen_args) as (proc, _):
                     try:
                         proc.wait(timeout=timeout)
                     except subprocess.TimeoutExpired as e:
@@ -648,28 +660,29 @@ def run_cmd(
         env=None,
         **popen_args,
     ):
-        with self.start_cmd(*cmd, raise_on_non_zero=raise_on_non_zero, env=env, **popen_args) as (proc, stdout, stderr):
+        with self.start_cmd(*cmd, raise_on_non_zero=raise_on_non_zero, env=env, **popen_args) as (proc, stream_iter):
             if input:
                 proc.stdin.write(input)
-            proc.stdin.close()
 
             if echo:
-                for line in stdout:
-                    if line:
-                        CONSOLE.msg(line.decode().strip())
-            elif capture_text:
-                return b"".join(stdout).decode()
-            elif capture_json:
-                try:
-                    return json.loads(b"".join(stdout).decode())
-                except json.JSONDecodeError:
-                    LOG.warning("Error decoding JSON from stdout")
-                    return {}
+                for stdout_line, _ in stream_iter:
+                    if stdout_line:
+                        CONSOLE.msg(stdout_line.strip())
+            elif capture_text or capture_json:
+                text = "".join([stdout_line for stdout_line, _ in stream_iter])
+                if capture_text:
+                    return text
+                else:
+                    try:
+                        return json.loads(text)
+                    except json.JSONDecodeError:
+                        LOG.warning("Error decoding JSON from stdout")
+                        return
             elif capture_json_lines:
                 json_lines = []
-                for idx, output_line in enumerate(stdout):
+                for idx, (stdout_line, _) in enumerate(stream_iter):
                     try:
-                        json_lines.append(json.loads(output_line.decode()))
+                        json_lines.append(json.loads(stdout_line))
                     except json.JSONDecodeError:
                         LOG.warning(f"Error decoding JSON from stdout line #{idx}")
                 return json_lines
@@ -693,6 +706,8 @@ def start_cmd(self, *cmd, raise_on_non_zero=True, env=None, **popen_args):
                 stderr=subprocess.PIPE,
                 stdin=subprocess.PIPE,
                 env=env,
+                # Always using text=False to match the BytesIO used for buffering the output
+                text=False,
                 **popen_args,
             )
         except FileNotFoundError as e:
@@ -701,35 +716,34 @@ def start_cmd(self, *cmd, raise_on_non_zero=True, env=None, **popen_args):
 
         slug_cmd = re.sub(r"[^a-zA-Z]+", "-", cmd_str)[:100].strip("-")
 
-        def get_stream_iterator(stream_name):
-            file_name = f"{self._cmd_idx:04d}-{stream_name}-{slug_cmd}.txt"
-            file_path = self.session_folder.joinpath(file_name)
-            return StreamIterator(proc, getattr(proc, stream_name), file_path)
+        stdout_path, stderr_path = [
+            self.session_folder.joinpath(f"{self._cmd_idx:04d}-{stream_name}-{slug_cmd}.txt")
+            for stream_name in ("stdout", "stderr")
+        ]
 
         try:
-            with (
-                get_stream_iterator("stdout") as stdout_iter,
-                get_stream_iterator("stderr") as stderr_iter,
-            ):
-                try:
-                    yield proc, stdout_iter, stderr_iter
-                finally:
-                    proc.wait()
+            with stream_iterator(proc, stdout_path, stderr_path) as stream_iter:
+                yield proc, stream_iter
+
             if raise_on_non_zero and proc.returncode != 0:
                 raise CommandFailed
         # We capture and raise CommandFailed to allow the client code to raise an empty CommandFailed exception
         # but still get a contextualized exception at the end
         except CommandFailed as e:
             raise CommandFailed(self._cmd_idx, cmd_str, proc.returncode) from e.__cause__
+        # Cleaning the current line in case the command is interrupted
+        except KeyboardInterrupt:
+            print()
+            raise
         finally:
             elapsed = time.time() - started
             LOG.info(
-                "Command [%04d] returned [%d] in [%.3f] seconds. [%d] bytes in STDOUT, [%d] bytes in STDERR",
+                "Command [%04d] returned [%s] in [%.3f] seconds. [%d] bytes in STDOUT, [%d] bytes in STDERR",
                 self._cmd_idx,
                 proc.returncode,
                 elapsed,
-                stdout_iter.bytes_written,
-                stderr_iter.bytes_written,
+                stdout_path.stat().st_size if stdout_path.exists() else 0,
+                stderr_path.stat().st_size if stderr_path.exists() else 0,
             )
 
 
@@ -1068,7 +1082,7 @@ def pre_execute(self, action, args):
             capture_json=True,
             raise_on_non_zero=False,
         )
-        if "Name" in env_json:
+        if env_json and "Name" in env_json:
             CONSOLE.msg(
                 "Found a minikube profile with the same name. If a previous attempt to run this installer failed,"
             )
@@ -1455,27 +1469,26 @@ def get_parser(self, sub_parsers):
     def execute(self, args):
         CONSOLE.title("Expose Observability ports")
 
-        try:
-            with self.start_cmd(
-                "minikube",
-                "kubectl",
-                "--profile",
-                args.profile,
-                "--",
-                "--namespace",
-                args.namespace,
-                "--address",
-                "0.0.0.0",
-                "port-forward",
-                "service/observability-ui",
-                f"{args.port}:http",
-                raise_on_non_zero=False,
-            ) as (proc, stdout, stderr):
-                for output in stdout:
-                    if output:
-                        break
-
-                if proc.poll() is None:
+        with self.start_cmd(
+            "minikube",
+            "kubectl",
+            "--profile",
+            args.profile,
+            "--",
+            "--namespace",
+            args.namespace,
+            "--address",
+            "0.0.0.0",
+            "port-forward",
+            "service/observability-ui",
+            f"{args.port}:http",
+            raise_on_non_zero=False,
+        ) as (proc, stream_iter):
+            succeeded = False
+            for stdout_line, stderr_line in stream_iter:
+                if stderr_line:
+                    CONSOLE.msg(stderr_line)
+                if stdout_line and proc.poll() is None:
                     url = f"http://localhost:{args.port}"
                     for service, label in SERVICES_LABELS.items():
                         CONSOLE.msg(f"{label:>20}: {SERVICES_URLS[service].format(url)}")
@@ -1495,33 +1508,26 @@ def execute(self, args):
                             file.write(json.dumps(json_config))
                     except Exception:
                         LOG.exception(f"Unable to update {DEMO_CONFIG_FILE} file with exposed port")
-                else:
-                    for output in stderr:
-                        if output:
-                            CONSOLE.msg(output.decode().strip())
-                    raise CommandFailed
 
                 try:
                     while proc.poll() is None:
                         time.sleep(10)
                 except KeyboardInterrupt:
-                    # The empty print forces the terminal cursor to move to the first column
+                    proc.terminate()
                     print()
+                    succeeded = True
 
-                proc.terminate()
-
+        if succeeded:
             CONSOLE.msg("The services are no longer exposed.")
-
-        except Exception as e:
-            LOG.exception("Something went wrong exposing the services ports")
+        else:
             CONSOLE.space()
             CONSOLE.msg("The platform could not have its ports exposed.")
             CONSOLE.msg(
                 f"Verify if the platform is running and installer has permission to listen at the port {args.port}."
             )
             CONSOLE.space()
             CONSOLE.msg(f"If port {args.port} is in use, use the command option --port to specify an alternate value.")
-            raise AbortAction from e
+            raise AbortAction
 
 
 class ObsDeleteAction(Action):
@@ -1558,7 +1564,7 @@ class DemoContainerAction(Action):
     requirements = [REQ_DOCKER, REQ_DOCKER_DAEMON]
 
     def run_dk_demo_container(self, command: str):
-        with self.start_cmd(
+        self.run_cmd(
             "docker",
             "run",
             "--rm",
@@ -1572,14 +1578,8 @@ def run_dk_demo_container(self, command: str):
             "host.docker.internal:host-gateway",
             DEMO_IMAGE,
             command,
-        ) as (proc, stdout, stderr):
-            try:
-                for line in stdout:
-                    if line:
-                        CONSOLE.msg(line.decode().strip())
-            except KeyboardInterrupt:
-                print("")
-                proc.terminate()
+            echo=True,
+        )
 
 
 class ObsRunDemoAction(DemoContainerAction):
@@ -1593,6 +1593,7 @@ def execute(self, args):
             CONSOLE.title("Demo FAILED")
             CONSOLE.space()
             CONSOLE.msg(f"To retry the demo, first run `python3 {INSTALLER_NAME} {args.prod} delete-demo`")
+            raise AbortAction
         else:
             CONSOLE.title("Demo SUCCEEDED")