Add compatibility with StackContext

Author: m-czernek

salt/client/mixins.py

@@ -29,6 +29,12 @@
 import salt.utils.user
 import salt.utils.versions
 
+from salt import USE_VENDORED_TORNADO
+if USE_VENDORED_TORNADO:
+    from salt.ext.tornado.stack_context import StackContext
+else:
+    from contextlib import nullcontext as StackContext
+
 log = logging.getLogger(__name__)
 
 CLIENT_INTERNAL_KEYWORDS = frozenset(
@@ -378,26 +384,27 @@ def low(self, fun, low, print_event=True, full_return=False):
                 data["fun_args"] = list(args) + ([kwargs] if kwargs else [])
                 func_globals["__jid_event__"].fire_event(data, "new")
 
-                func = self.functions[fun]
-                try:
-                    data["return"] = func(*args, **kwargs)
-                except TypeError as exc:
-                    data[
-                        "return"
-                    ] = "\nPassed invalid arguments: {}\n\nUsage:\n{}".format(
-                        exc, func.__doc__
-                    )
-                try:
-                    data["success"] = self.context.get("retcode", 0) == 0
-
-                except AttributeError:
-                    # Assume a True result if no context attribute
-                    data["success"] = True
-                if isinstance(data["return"], dict) and "data" in data["return"]:
-                    # some functions can return boolean values
-                    data["success"] = salt.utils.state.check_result(
-                        data["return"]["data"]
-                    )
+                with StackContext(self.functions.context_dict.clone):
+                    func = self.functions[fun]
+                    try:
+                        data["return"] = func(*args, **kwargs)
+                    except TypeError as exc:
+                        data[
+                            "return"
+                        ] = "\nPassed invalid arguments: {}\n\nUsage:\n{}".format(
+                            exc, func.__doc__
+                        )
+                    try:
+                        data["success"] = self.context.get("retcode", 0) == 0
+
+                    except AttributeError:
+                        # Assume a True result if no context attribute
+                        data["success"] = True
+                    if isinstance(data["return"], dict) and "data" in data["return"]:
+                        # some functions can return boolean values
+                        data["success"] = salt.utils.state.check_result(
+                            data["return"]["data"]
+                        )
 
             except (Exception, SystemExit) as ex:  # pylint: disable=broad-except
                 if isinstance(ex, salt.exceptions.NotImplemented):

salt/metaproxy/deltaproxy.py

@@ -56,6 +56,12 @@
 from salt.utils.event import tagify
 from salt.utils.process import SignalHandlingProcess, default_signals
 
+from salt import USE_VENDORED_TORNADO
+if USE_VENDORED_TORNADO:
+    from salt.ext.tornado.stack_context import StackContext
+else:
+    from contextlib import nullcontext as StackContext
+
 log = logging.getLogger(__name__)
 
 
@@ -562,10 +568,11 @@ def target(cls, minion_instance, opts, data, connected):
         uid = salt.utils.user.get_uid(user=opts.get("user", None))
         minion_instance.proc_dir = salt.minion.get_proc_dir(opts["cachedir"], uid=uid)
 
-    if isinstance(data["fun"], tuple) or isinstance(data["fun"], list):
-        ProxyMinion._thread_multi_return(minion_instance, opts, data)
-    else:
-        ProxyMinion._thread_return(minion_instance, opts, data)
+    with StackContext(minion_instance.ctx):
+        if isinstance(data["fun"], tuple) or isinstance(data["fun"], list):
+            ProxyMinion._thread_multi_return(minion_instance, opts, data)
+        else:
+            ProxyMinion._thread_return(minion_instance, opts, data)
 
 
 def thread_return(cls, minion_instance, opts, data):

salt/metaproxy/proxy.py

@@ -52,6 +52,12 @@
 from salt.utils.event import tagify
 from salt.utils.process import SignalHandlingProcess, default_signals
 
+from salt import USE_VENDORED_TORNADO
+if USE_VENDORED_TORNADO:
+    from salt.ext.tornado.stack_context import StackContext
+else:
+    from contextlib import nullcontext as StackContext
+
 log = logging.getLogger(__name__)
 
 
@@ -380,10 +386,11 @@ def target(cls, minion_instance, opts, data, connected):
                 opts["cachedir"], uid=uid
             )
 
-    if isinstance(data["fun"], tuple) or isinstance(data["fun"], list):
-        ProxyMinion._thread_multi_return(minion_instance, opts, data)
-    else:
-        ProxyMinion._thread_return(minion_instance, opts, data)
+    with StackContext(minion_instance.ctx):
+        if isinstance(data["fun"], tuple) or isinstance(data["fun"], list):
+            ProxyMinion._thread_multi_return(minion_instance, opts, data)
+        else:
+            ProxyMinion._thread_return(minion_instance, opts, data)
 
 
 def thread_return(cls, minion_instance, opts, data):

salt/minion.py

@@ -101,6 +101,11 @@
 except ImportError:
     HAS_WIN_FUNCTIONS = False
 
+from salt import USE_VENDORED_TORNADO
+if USE_VENDORED_TORNADO:
+    from salt.ext.tornado.stack_context import ExceptionStackContext
+else:
+    from contextlib import nullcontext as ExceptionStackContext
 
 log = logging.getLogger(__name__)
 
@@ -1675,9 +1680,10 @@ def handle_timeout(*_):
 
                 timeout_handler = handle_timeout
 
-            # pylint: disable=unexpected-keyword-arg
-            self._send_req_async(load, timeout)
-            # pylint: enable=unexpected-keyword-arg
+            with ExceptionStackContext(timeout_handler):
+                # pylint: disable=unexpected-keyword-arg
+                self._send_req_async(load, timeout)
+                # pylint: enable=unexpected-keyword-arg
         return True
 
     @tornado.gen.coroutine
@@ -2251,11 +2257,12 @@ def timeout_handler(*_):
                 timeout_handler()
                 return ""
         else:
-            # pylint: disable=unexpected-keyword-arg
-            ret_val = self._send_req_async(
-                load, timeout=timeout
-            )
-            # pylint: enable=unexpected-keyword-arg
+            with ExceptionStackContext(timeout_handler):
+                # pylint: disable=unexpected-keyword-arg
+                ret_val = self._send_req_async(
+                    load, timeout=timeout
+                )
+                # pylint: enable=unexpected-keyword-arg
 
         log.trace("ret_val = %s", ret_val)  # pylint: disable=no-member
         return ret_val
@@ -2341,11 +2348,12 @@ def timeout_handler(*_):
                 timeout_handler()
                 return ""
         else:
-            # pylint: disable=unexpected-keyword-arg
-            ret_val = self._send_req_async(
-                load, timeout=timeout
-            )
-            # pylint: enable=unexpected-keyword-arg
+            with ExceptionStackContext(timeout_handler):
+                # pylint: disable=unexpected-keyword-arg
+                ret_val = self._send_req_async(
+                    load, timeout=timeout
+                )
+                # pylint: enable=unexpected-keyword-arg
 
         log.trace("ret_val = %s", ret_val)  # pylint: disable=no-member
         return ret_val
@@ -3270,18 +3278,19 @@ def timeout_handler(*args):
             log.warning("Unable to forward pub data: %s", args[1])
             return True
 
-        self.local.pub_async(
-            data["tgt"],
-            data["fun"],
-            data["arg"],
-            data["tgt_type"],
-            data["ret"],
-            data["jid"],
-            data["to"],
-            io_loop=self.io_loop,
-            callback=lambda _: None,
-            **kwargs
-        )
+        with ExceptionStackContext(timeout_handler):
+            self.local.pub_async(
+                data["tgt"],
+                data["fun"],
+                data["arg"],
+                data["tgt_type"],
+                data["ret"],
+                data["jid"],
+                data["to"],
+                io_loop=self.io_loop,
+                callback=lambda _: None,
+                **kwargs
+            )
 
     def _send_req_sync(self, load, timeout):
         if self.opts["minion_sign_messages"]:

salt/utils/thin.py

@@ -24,7 +24,7 @@
 
 import salt
 import salt.exceptions
-import tornado as tornado
+import tornado
 import salt.utils.files
 import salt.utils.hashutils
 import salt.utils.json