From 4c59b13409e22471bc0a96947a91d3b6b38b90f3 Mon Sep 17 00:00:00 2001
From: Alex Kavanagh <alex.kavanagh@canonical.com>
Date: Fri, 17 May 2019 16:00:00 +0100
Subject: [PATCH] Fix execute command missing output with pauses in output text

For the container execute command, if the command executed has pauses in
the output text, the websocket will send binary messages.  These were
interpreted as the stream being completed, and pylxd closed the
websocket, thus losing the rest of the output.

Delving into the code indicated some very strange behaviour across the
threads, so as part of solving the problem, this has been cleaned up as
well.

Closes: #362
Signed-off-by: Alex Kavanagh <alex.kavanagh@canonical.com>
---
 pylxd/models/container.py | 38 +++++++++++++++++++++++++++++++-------
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/pylxd/models/container.py b/pylxd/models/container.py
index cbd8206f..4c4f91e3 100644
--- a/pylxd/models/container.py
+++ b/pylxd/models/container.py
@@ -353,7 +353,7 @@ def unfreeze(self, timeout=30, force=True, wait=False):
     def execute(
             self, commands, environment={}, encoding=None, decode=True,
             stdin_payload=None, stdin_encoding="utf-8",
-            stdout_handler=None, stderr_handler=None
+            stdout_handler=None, stderr_handler=None,
     ):
         """Execute a command on the container.
 
@@ -430,8 +430,17 @@ def execute(
                     break
                 time.sleep(.5)  # pragma: no cover
 
-            while len(manager.websockets.values()) > 0:
-                time.sleep(.1)  # pragma: no cover
+            try:
+                stdin.close()
+            except BrokenPipeError:
+                pass
+
+            stdout.finish_soon()
+            stderr.finish_soon()
+            manager.close_all()
+
+            while not stdout.finished or not stderr.finished:
+                time.sleep(.1)  # progma: no cover
 
             manager.stop()
             manager.join()
@@ -618,6 +627,9 @@ def __init__(self, manager, *args, **kwargs):
         self.encoding = kwargs.pop('encoding', None)
         self.handler = kwargs.pop('handler', None)
         self.message_encoding = None
+        self.finish_off = False
+        self.finished = False
+        self.last_message_empty = False
         super(_CommandWebsocketClient, self).__init__(*args, **kwargs)
 
     def handshake_ok(self):
@@ -626,15 +638,27 @@ def handshake_ok(self):
 
     def received_message(self, message):
         if message.data is None or len(message.data) == 0:
-            self.manager.remove(self)
-            return
+            self.last_message_empty = True
+            if self.finish_off:
+                self.finished = True
+                return
+        else:
+            self.last_message_empty = False
         if message.encoding and self.message_encoding is None:
             self.message_encoding = message.encoding
         if self.handler:
             self.handler(self._maybe_decode(message.data))
         self.buffer.append(message.data)
-        if isinstance(message, BinaryMessage):
-            self.manager.remove(self)
+        if self.finish_off and isinstance(message, BinaryMessage):
+            self.finished = True
+
+    def closed(self, code, reason=None):
+        self.finished = True
+
+    def finish_soon(self):
+        self.finish_off = True
+        if self.last_message_empty:
+            self.finished = True
 
     def _maybe_decode(self, buffer):
         if self.decode and buffer is not None: