feature: address threading race conditions concerns

Author: DanielMicrosoft

src/azure-cli-core/azure/cli/core/__init__.py

@@ -202,13 +202,14 @@ def _configure_style(self):
 
 
 class ModuleLoadResult:  # pylint: disable=too-few-public-methods
-    def __init__(self, module_name, command_table, group_table, elapsed_time, error=None, traceback_str=None):
+    def __init__(self, module_name, command_table, group_table, elapsed_time, error=None, traceback_str=None, command_loader=None):
         self.module_name = module_name
         self.command_table = command_table
         self.group_table = group_table
         self.elapsed_time = elapsed_time
         self.error = error
         self.traceback_str = traceback_str
+        self.command_loader = command_loader
 
 
 class MainCommandsLoader(CLICommandsLoader):
@@ -343,7 +344,7 @@ def _filter_modname(extensions):
                         # from an extension requires this map to be up-to-date.
                         # self._mod_to_ext_map[ext_mod] = ext_name
                         start_time = timeit.default_timer()
-                        extension_command_table, extension_group_table = \
+                        extension_command_table, extension_group_table, _ = \
                             _load_extension_command_loader(self, args, ext_mod)
                         import_extension_breaking_changes(ext_mod)
 
@@ -581,10 +582,10 @@ def _load_single_module(self, mod, args):
         import traceback
         try:
             start_time = timeit.default_timer()
-            module_command_table, module_group_table = _load_module_command_loader(self, args, mod)
+            module_command_table, module_group_table, command_loader = _load_module_command_loader(self, args, mod)
             import_module_breaking_changes(mod)
             elapsed_time = timeit.default_timer() - start_time
-            return ModuleLoadResult(mod, module_command_table, module_group_table, elapsed_time)
+            return ModuleLoadResult(mod, module_command_table, module_group_table, elapsed_time, command_loader=command_loader)
         except Exception as ex:  # pylint: disable=broad-except
             tb_str = traceback.format_exc()
             return ModuleLoadResult(mod, {}, {}, 0, ex, tb_str)
@@ -602,6 +603,12 @@ def _handle_module_load_error(self, result):
 
     def _process_successful_load(self, result):
         """Process successfully loaded module results."""
+        if result.command_loader:
+            self.loaders.append(result.command_loader)
+
+            for cmd in result.command_table:
+                self.cmd_to_loader_map[cmd] = [result.command_loader]
+
         for cmd in result.command_table.values():
             cmd.command_source = result.module_name
 

src/azure-cli-core/azure/cli/core/commands/__init__.py

@@ -1131,22 +1131,17 @@ def _load_command_loader(loader, args, name, prefix):
             logger.debug("Module '%s' is missing `get_command_loader` entry.", name)
 
     command_table = {}
+    command_loader = None
 
     if loader_cls:
         command_loader = loader_cls(cli_ctx=loader.cli_ctx)
-        loader.loaders.append(command_loader)  # This will be used by interactive
         if command_loader.supported_resource_type():
             command_table = command_loader.load_command_table(args)
-            if command_table:
-                for cmd in list(command_table.keys()):
-                    # TODO: If desired to for extension to patch module, this can be uncommented
-                    # if loader.cmd_to_loader_map.get(cmd):
-                    #    loader.cmd_to_loader_map[cmd].append(command_loader)
-                    # else:
-                    loader.cmd_to_loader_map[cmd] = [command_loader]
     else:
         logger.debug("Module '%s' is missing `COMMAND_LOADER_CLS` entry.", name)
-    return command_table, command_loader.command_group_table
+
+    group_table = command_loader.command_group_table if command_loader else {}
+    return command_table, group_table, command_loader
 
 
 def _load_extension_command_loader(loader, args, ext):