tinfoil: Ensure sockets don't leak even when exceptions occur

We're seeing leaking open socket connections when errors occur and
tinfoil is in use. Improve the exception handling so the sockets
are closed even if exceptions occur, allowing more robust behaviour
when things go wrong.

Signed-off-by: Richard Purdie <richard.purdie@linuxfoundation.org>

1 files changed, 16 insertions, 12 deletions

diff --git a/lib/bb/tinfoil.py b/lib/bb/tinfoil.py
index 5755e5a34..2fb1bb7d2 100644
--- a/lib/bb/tinfoil.py
+++ b/lib/bb/tinfoil.py
@@ -806,18 +806,22 @@ class Tinfoil:
         prepare() has been called, or use a with... block when you create
         the tinfoil object which will ensure that it gets called.
         """
-        if self.server_connection:
-            self.run_command('clientComplete')
-            _server_connections.remove(self.server_connection)
-            bb.event.ui_queue = []
-            self.server_connection.terminate()
-            self.server_connection = None
-
-        # Restore logging handlers to how it looked when we started
-        if self.oldhandlers:
-            for handler in self.logger.handlers:
-                if handler not in self.oldhandlers:
-                    self.logger.handlers.remove(handler)
+        try:
+            if self.server_connection:
+                try:
+                    self.run_command('clientComplete')
+                finally:
+                    _server_connections.remove(self.server_connection)
+                    bb.event.ui_queue = []
+                    self.server_connection.terminate()
+                    self.server_connection = None
+
+        finally:
+            # Restore logging handlers to how it looked when we started
+            if self.oldhandlers:
+                for handler in self.logger.handlers:
+                    if handler not in self.oldhandlers:
+                        self.logger.handlers.remove(handler)
 
     def _reconvert_type(self, obj, origtypename):
         """