From 9d5d70564f569fac9a9212082382f602dc106f77 Mon Sep 17 00:00:00 2001
From: Douglas RAILLARD <douglas.raillard@arm.com>
Date: Wed, 15 Jan 2020 17:16:47 +0000
Subject: [PATCH] target: Use tls_property() to manage a thread-local
 connection

This frees the connection to have to handle threading issues, since each thread
using the Target will have its own connection. The connection will be garbage
collected when the thread using it dies, avoiding connection leaks.
---
 devlib/target.py | 30 ++++++++++++++----------------
 1 file changed, 14 insertions(+), 16 deletions(-)

diff --git a/devlib/target.py b/devlib/target.py
index e3765a3..2b8296c 100644
--- a/devlib/target.py
+++ b/devlib/target.py
@@ -53,7 +53,7 @@ from devlib.utils.android import AdbConnection, AndroidProperties, LogcatMonitor
 from devlib.utils.misc import memoized, isiterable, convert_new_lines
 from devlib.utils.misc import commonprefix, merge_lists
 from devlib.utils.misc import ABI_MAP, get_cpu_name, ranges_to_list
-from devlib.utils.misc import batch_contextmanager
+from devlib.utils.misc import batch_contextmanager, tls_property
 from devlib.utils.types import integer, boolean, bitmask, identifier, caseless_string, bytes_regex
 
 
@@ -214,22 +214,19 @@ class Target(object):
         cmd = "cat /proc/self/smaps | {0} grep KernelPageSize | {0} head -n 1 | {0} awk '{{ print $2 }}'"
         return int(self.execute(cmd.format(self.busybox)))
 
-    @property
-    def conn(self):
-        if self._connections:
-            tid = id(threading.current_thread())
-            if tid not in self._connections:
-                self._connections[tid] = self.get_connection()
-            return self._connections[tid]
-        else:
-            return None
-
     @property
     def shutils(self):
         if self._shutils is None:
             self._setup_shutils()
         return self._shutils
 
+    @tls_property
+    def _conn(self):
+        return self.get_connection()
+
+    # Add a basic property that does not require calling to get the value
+    conn = _conn.basic_property
+
     def __init__(self,
                  connection_settings=None,
                  platform=None,
@@ -242,6 +239,7 @@ class Target(object):
                  conn_cls=None,
                  is_container=False
                  ):
+
         self._is_rooted = None
         self.connection_settings = connection_settings or {}
         # Set self.platform: either it's given directly (by platform argument)
@@ -271,7 +269,6 @@ class Target(object):
         self._installed_binaries = {}
         self._installed_modules = {}
         self._cache = {}
-        self._connections = {}
         self._shutils = None
         self._file_transfer_cache = None
         self.busybox = None
@@ -290,8 +287,9 @@ class Target(object):
 
     def connect(self, timeout=None, check_boot_completed=True):
         self.platform.init_target_connection(self)
-        tid = id(threading.current_thread())
-        self._connections[tid] = self.get_connection(timeout=timeout)
+        # Forcefully set the thread-local value for the connection, with the
+        # timeout we want
+        self.conn = self.get_connection(timeout=timeout)
         if check_boot_completed:
             self.wait_boot_complete(timeout)
         self._resolve_paths()
@@ -304,9 +302,9 @@ class Target(object):
             self._install_module(get_module('bl'))
 
     def disconnect(self):
-        for conn in self._connections.values():
+        connections = self._conn.get_all_values()
+        for conn in connections:
             conn.close()
-        self._connections = {}
 
     def get_connection(self, timeout=None):
         if self.conn_cls is None: