Merge pull request #736 from netenglabs/pacer-pt2

Support pacing commands/logins

Author: ddutt

suzieq/poller/controller/source/base_source.py

@@ -49,6 +49,7 @@ def __init__(self, input_data, validate: bool = True) -> None:
             'jump_host_key_file',
             'ignore_known_hosts',
             'slow_host',
+            'per_cmd_auth',
         ]
 
         self._load(input_data)
@@ -204,6 +205,8 @@ def set_device(self, inventory: Dict[str, Dict]):
         transport = None
         ignore_known_hosts = False
         slow_host = False
+        per_cmd_auth = True
+        retries_on_auth_fail = 0
         port = None
         devtype = None
 
@@ -223,6 +226,8 @@ def set_device(self, inventory: Dict[str, Dict]):
                 transport = transport.value
             ignore_known_hosts = self._device.get('ignore-known-hosts', False)
             slow_host = self._device.get('slow-host', False)
+            per_cmd_auth = self._device.get('per-cmd-auth', True)
+            retries_on_auth_fail = self._device.get('retries-on-auth-fail', 0)
             port = self._device.get('port')
             devtype = self._device.get('devtype')
 
@@ -237,13 +242,19 @@ def set_device(self, inventory: Dict[str, Dict]):
                 'port': node.get('port') or port or 22,
                 'devtype': node.get('devtype') or devtype,
                 'slow_host': node.get('slow_host', '') or slow_host,
+                'per_cmd_auth': ((node.get('per_cmd_auth', '') != '')
+                                 or per_cmd_auth),
+                'retries_on_auth_fail': ((node.get('retries_on_auth_fail',
+                                                   -1) != -1) or
+                                         retries_on_auth_fail)
             })
 
     def _validate_device(self):
         if self._device:
             dev_fields = ['name', 'jump-host', 'jump-host-key-file',
                           'ignore-known-hosts', 'transport', 'port',
-                          'slow-host', 'devtype']
+                          'slow-host', 'per_cmd_auth', 'retries_on_auth_fail',
+                          'devtype']
             inv_fields = [x for x in self._device if x not in dev_fields]
             if inv_fields:
                 raise InventorySourceError(

suzieq/poller/controller/utils/inventory_models.py

@@ -30,10 +30,14 @@ class DeviceModel(BaseModel):
     ignore_known_hosts: Optional[bool] = Field(
         alias='ignore-known-hosts', default=False)
     slow_host: Optional[bool] = Field(alias='slow-host', default=False)
+    per_cmd_auth: Optional[bool] = Field(alias='per-cmd-auth', default=True)
+    retries_on_auth_fail: Optional[int] = Field(alias='retries-on-auth-fail',
+                                                default=1)
     transport: Optional[PollerTransport]
     port: Optional[str]
     devtype: Optional[str]
 
+    # pylint: disable=no-self-argument
     @validator('jump_host')
     def jump_host_validator(cls, value):
         '''Validate jump host format'''

suzieq/poller/worker/inventory/inventory.py

@@ -33,6 +33,7 @@ def __init__(self, add_task_fn: Callable, **kwargs) -> None:
         self.add_task_fn = add_task_fn
         self._max_outstanding_cmd = 0
         self._cmd_semaphore = None
+        self._cmd_pacer_mutex = None
 
         self.connect_timeout = kwargs.pop('connect_timeout', 15)
         self.ssh_config_file = kwargs.pop('ssh_config_file', None)
@@ -75,6 +76,7 @@ async def build_inventory(self) -> Dict[str, Node]:
 
         if self._max_outstanding_cmd:
             self._cmd_semaphore = asyncio.Semaphore(self._max_outstanding_cmd)
+            self._cmd_pacer_mutex = asyncio.Lock()
 
         # Initialize the nodes in the inventory
         self._nodes = await self._init_nodes(inventory_list)
@@ -123,19 +125,30 @@ async def _init_nodes(self, inventory_list:
 
         for host in inventory_list:
             new_node = Node()
-            init_tasks += [new_node.initialize(
-                **host,
-                cmd_sem=self._cmd_semaphore,
-                connect_timeout=self.connect_timeout,
-                ssh_config_file=self.ssh_config_file
-            )]
+            if self._max_outstanding_cmd > 0:
+                init_tasks += [new_node.initialize(
+                    **host,
+                    cmd_sem=self._cmd_semaphore,
+                    cmd_mutex=self._cmd_pacer_mutex,
+                    cmd_pacer_sleep=float(1/self._max_outstanding_cmd),
+                    connect_timeout=self.connect_timeout,
+                    ssh_config_file=self.ssh_config_file
+                )]
+            else:
+                init_tasks += [new_node.initialize(
+                    **host,
+                    cmd_sem=self._cmd_semaphore,
+                    cmd_mutex=self._cmd_pacer_mutex,
+                    connect_timeout=self.connect_timeout,
+                    ssh_config_file=self.ssh_config_file
+                )]
 
         for n in asyncio.as_completed(init_tasks):
             try:
                 newnode = await n
             except Exception as e:  # pylint: disable=broad-except
                 logger.error(
-                    f'Encountered error {e} in initializing node')
+                    f'Encountered error {str(e)} in initializing node')
                 continue
             if newnode.devtype == "unsupported":
                 logger.error(

suzieq/poller/worker/nodes/node.py

@@ -109,9 +109,14 @@ async def initialize(self, **kwargs) -> TNode:
         self._current_exception = None
         self.api_key = None
         self._stdin = self._stdout = self._long_proc = None
-        self._retry = True      # set to False if authentication fails
+        self._max_retries_on_auth_fail = (kwargs.get('retries_on_auth_fail')
+                                          or 0) + 1
+        self._retry = self._max_retries_on_auth_fail
         self._discovery_lock = asyncio.Lock()
         self._cmd_sem = kwargs.get('cmd_sem', None)
+        self._cmd_mutex = kwargs.get('cmd_mutex', None)
+        self._cmd_pacer_sleep = kwargs.get('cmd_pacer_sleep', None)
+        self.per_cmd_auth = kwargs.get('per_cmd_auth', True)
 
         self.address = kwargs["address"]
         self.hostname = kwargs["address"]  # default till we get hostname
@@ -242,31 +247,33 @@ def is_connected(self):
         return self._conn is not None
 
     @asynccontextmanager
-    async def limit_pipeline(self):
+    async def cmd_pacer(self, use_sem: bool = True):
         '''Context Manager to implement throttling of commands.
 
         In many networks, backend authentication servers such as TACACS which
         handle authentication of logins and even command execution, cannot
         large volumes of authentication requests. Thanks to our use of
         asyncio, we can easily sends hundreds of connection requests to such
         servers, which effectively turns into authentication failures. To
-        handle this, we add a user-specified maximum of simultaneous
-        commands/logins at any given time. This code implements that locking
-        context.
+        handle this, we add a user-specified maximum of rate of cmds/sec
+        that the authentication can handle, and we pace it out. This code
+        implements that pacer.
+
+        Some networks communicate with a backend authentication server only
+        on login while others contact it for authorization of a command as
+        well. Its to handle this difference that we pass use_sem. Users set
+        the per_cmd_auth to True if authorization is used. The caller of this
+        function sets the use_sem apppropriately depending on when the context
+        is invoked.
+
+        Args:
+          use_sem(bool): True if you want to use the pacer
         '''
-        if self._cmd_sem:
-            self.logger.debug(
-                f'{self.transport}://{self.hostname}:{self.port}: Get lock')
-            await self._cmd_sem.acquire()
-            self.logger.debug(
-                f'{self.transport}://{self.hostname}:{self.port}: Got lock')
-            try:
+        if self._cmd_sem and use_sem:
+            async with self._cmd_sem:
+                async with self._cmd_mutex:
+                    await asyncio.sleep(self._cmd_pacer_sleep)
                 yield
-            finally:
-                self.logger.debug(
-                    f'{self.transport}://{self.hostname}:{self.port}: '
-                    'Free lock')
-                self._cmd_sem.release()
         else:
             yield
 
@@ -643,7 +650,7 @@ async def _init_ssh(self, init_dev_data=True, use_lock=True) -> None:
                         self.ssh_ready.release()
                     return
 
-            async with self.limit_pipeline():
+            async with self.cmd_pacer():
                 try:
                     if self._tunnel:
                         self._conn = await self._tunnel.connect_ssh(
@@ -659,6 +666,8 @@ async def _init_ssh(self, init_dev_data=True, use_lock=True) -> None:
                     self.logger.info(
                         f"Connected to {self.address}:{self.port} at "
                         f"{time.time()}")
+                    # Reset authentication fail attempt on success
+                    self._retry = self._max_retries_on_auth_fail
                 except Exception as e:  # pylint: disable=broad-except
                     if isinstance(e, asyncssh.HostKeyNotVerifiable):
                         self.logger.error(
@@ -672,7 +681,7 @@ async def _init_ssh(self, init_dev_data=True, use_lock=True) -> None:
                             f'Authentication failed to {self.address}. '
                             'Not retrying to avoid locking out user. Please '
                             'restart poller with proper authentication')
-                        self._retry = False
+                        self._retry -= 1
                     else:
                         self.logger.error('Unable to connect to '
                                           f'{self.address}:{self.port}, {e}')
@@ -790,7 +799,7 @@ async def _ssh_gather(self, service_callback: Callable,
             cb_token.node_token = self.bootupTimestamp
 
         timeout = timeout or self.cmd_timeout
-        async with self.limit_pipeline():
+        async with self.cmd_pacer(self.per_cmd_auth):
             for cmd in cmd_list:
                 try:
                     output = await asyncio.wait_for(self._conn.run(cmd),
@@ -1162,7 +1171,7 @@ async def _rest_gather(self, service_callback, cmd_list, cb_token,
         output = []
         status = 200  # status OK
 
-        async with self.limit_pipeline():
+        async with self.cmd_pacer(self.per_cmd_auth):
             try:
                 async with aiohttp.ClientSession(
                         auth=auth, conn_timeout=self.connect_timeout,
@@ -1309,7 +1318,7 @@ async def _init_rest(self):
         url = "https://{0}:{1}/nclu/v1/rpc".format(self.address, self.port)
         headers = {"Content-Type": "application/json"}
 
-        async with self.limit_pipeline():
+        async with self.cmd_pacer(self.per_cmd_auth):
             try:
                 async with aiohttp.ClientSession(
                         auth=auth, timeout=self.cmd_timeout,
@@ -1334,7 +1343,7 @@ async def _rest_gather(self, service_callback, cmd_list, cb_token,
         url = "https://{0}:{1}/nclu/v1/rpc".format(self.address, self.port)
         headers = {"Content-Type": "application/json"}
 
-        async with self.limit_pipeline():
+        async with self.cmd_pacer(self.per_cmd_auth):
             try:
                 async with aiohttp.ClientSession(
                         auth=auth,
@@ -1524,7 +1533,7 @@ async def _init_ssh(self, init_dev_data=True,
         if self.is_connected and not self._stdin:
             self.logger.info(
                 f'Trying to create Persistent SSH for {self.hostname}')
-            async with self.limit_pipeline():
+            async with self.cmd_pacer(self.per_cmd_auth):
                 try:
                     self._stdin, self._stdout, self._stderr = \
                         await self._conn.open_session(term_type='xterm')
@@ -1537,11 +1546,15 @@ async def _init_ssh(self, init_dev_data=True,
                             await self._close_connection()
                             self._conn = None
                             self._stdin = None
-                            self._retry = False  # No retry if escalation fails
+                            self._retry -= 1
                             if use_lock:
                                 self.ssh_ready.release()
                             return
+                    # Reset number of retries on successful auth
+                    self._retry = self._max_retries_on_auth_fail
                 except Exception as e:
+                    if isinstance(e, asyncssh.misc.PermissionDenied):
+                        self._retry -= 1
                     self.current_exception = e
                     self.logger.error('Unable to create persistent SSH session'
                                       f' for {self.hostname} due to {str(e)}')
@@ -1657,7 +1670,7 @@ async def _ssh_gather(self, service_callback, cmd_list, cb_token, oformat,
             return
 
         timeout = timeout or self.cmd_timeout
-        async with self.limit_pipeline():
+        async with self.cmd_pacer(self.per_cmd_auth):
             for cmd in cmd_list:
                 try:
                     if self.slow_host:
@@ -1885,7 +1898,7 @@ async def _fetch_init_dev_data(self):
         try:
             res = []
             # temporary hack to detect device info using ssh
-            async with self.limit_pipeline():
+            async with self.cmd_pacer():
                 async with asyncssh.connect(
                         self.address, port=22, username=self.username,
                         password=self.password, known_hosts=None) as conn:
@@ -1922,7 +1935,7 @@ async def get_api_key(self):
         url = f"https://{self.address}:{self.port}/api/?type=keygen&user=" \
             f"{self.username}&password={self.password}"
 
-        async with self.limit_pipeline():
+        async with self.cmd_pacer(self.per_cmd_auth):
             async with self._session.get(url, timeout=self.connect_timeout) \
                     as response:
                 status, xml = response.status, await response.text()
@@ -1974,7 +1987,7 @@ def _extract_nos_version(self, data: str) -> None:
     async def _init_rest(self):
         # In case of PANOS, getting here means REST is up
         if not self._session:
-            async with self.limit_pipeline():
+            async with self.cmd_pacer(self.per_cmd_auth):
                 try:
                     self._session = aiohttp.ClientSession(
                         conn_timeout=self.connect_timeout,
@@ -2010,7 +2023,7 @@ async def _rest_gather(self, service_callback, cmd_list, cb_token,
             await service_callback(result, cb_token)
             return
 
-        async with self.limit_pipeline():
+        async with self.cmd_pacer(self.per_cmd_auth):
             try:
                 for cmd in cmd_list:
                     url_cmd = f"{url}?type=op&cmd={cmd}&key={self.api_key}"

tests/integration/sqcmds/common-samples/deprecated.yml

@@ -68,23 +68,23 @@ tests:
   marks: network show deprecated
   output: "WARNING: 'network show' is deprecated. Use 'namespace show' instead.\n\
     \     namespace  deviceCnt  ...  hasMlag                       lastUpdate\n0 \
-    \    dual-bgp         14  ...     True 2021-12-10 18:37:11.921000+00:00\n1   \
-    \ dual-evpn         14  ...     True 2021-12-10 18:34:58.729000+00:00\n2     \
-    \     eos         14  ...     True 2021-12-10 18:33:51.163000+00:00\n3       \
-    \ junos         12  ...    False 2022-02-05 14:03:21.736000+00:00\n4        mixed\
-    \          8  ...    False 2021-12-10 18:33:53.280000+00:00\n5         nxos  \
-    \       14  ...     True 2021-12-10 18:33:51.216000+00:00\n6    ospf-ibgp    \
-    \     14  ...     True 2022-01-10 16:01:05.167000+00:00\n7  ospf-single      \
-    \   14  ...    False 2021-12-10 18:36:06.022000+00:00\n8        panos        \
-    \ 14  ...     True 2021-12-14 16:21:20.095000+00:00\n9          vmx          5\
-    \  ...    False 2021-12-10 18:33:50.623000+00:00\n\n[10 rows x 9 columns]\n"
+    \    dual-bgp         14  ...     True 2022-05-15 04:08:24.070000+00:00\n1   \
+    \ dual-evpn         14  ...     True 2022-05-15 04:06:10.720000+00:00\n2     \
+    \     eos         14  ...     True 2022-05-15 04:05:01.207000+00:00\n3       \
+    \ junos         12  ...    False 2022-05-15 04:05:01.227000+00:00\n4        mixed\
+    \          8  ...    False 2022-05-15 04:05:00.852000+00:00\n5         nxos  \
+    \       14  ...     True 2022-05-15 13:25:17.544000+00:00\n6    ospf-ibgp    \
+    \     14  ...     True 2022-05-15 04:05:00.802000+00:00\n7  ospf-single      \
+    \   14  ...    False 2022-05-15 04:07:17.124000+00:00\n8        panos        \
+    \ 14  ...     True 2022-05-15 04:05:01.957000+00:00\n9          vmx          5\
+    \  ...    False 2022-05-15 04:04:59.544000+00:00\n\n[10 rows x 9 columns]\n"
 - command: network summarize
   data-directory: tests/data/parquet/
   format: text
   marks: network summarize deprecated
   output: "WARNING: 'network summarize' is deprecated. Use 'namespace summarize' instead.\n\
     \                         summary\nnamespacesCnt                 10\nservicePerNsStat\
-    \  [13, 18, 17.5]\nnsWithMlagCnt                  6\nnsWithBgpCnt            \
+    \  [13, 18, 17.0]\nnsWithMlagCnt                  6\nnsWithBgpCnt            \
     \       8\nnsWithOspfCnt                  7\nnsWithVxlanCnt                 6\n\
     nsWithErrsvcCnt                0\n"
 - command: network top --what=deviceCnt
@@ -93,11 +93,11 @@ tests:
   marks: network top deprecated
   output: "WARNING: 'network top' is deprecated. Use 'namespace top' instead.\n  \
     \   namespace  deviceCnt  ...  hasMlag                       lastUpdate\n0   \
-    \     panos         14  ...     True 2021-12-14 16:21:20.095000+00:00\n1  ospf-single\
-    \         14  ...    False 2021-12-10 18:36:06.022000+00:00\n2    ospf-ibgp  \
-    \       14  ...     True 2022-01-10 16:01:05.167000+00:00\n3         nxos    \
-    \     14  ...     True 2021-12-10 18:33:51.216000+00:00\n4          eos      \
-    \   14  ...     True 2021-12-10 18:33:51.163000+00:00\n\n[5 rows x 9 columns]\n"
+    \     panos         14  ...     True 2022-05-15 04:05:01.957000+00:00\n1  ospf-single\
+    \         14  ...    False 2022-05-15 04:07:17.124000+00:00\n2    ospf-ibgp  \
+    \       14  ...     True 2022-05-15 04:05:00.802000+00:00\n3         nxos    \
+    \     14  ...     True 2022-05-15 13:25:17.544000+00:00\n4          eos      \
+    \   14  ...     True 2022-05-15 04:05:01.207000+00:00\n\n[5 rows x 9 columns]\n"
 - command: network unique
   data-directory: tests/data/parquet/
   format: text

tests/integration/test_rest.py

@@ -397,8 +397,6 @@ def get_args_to_match(sqobj: SqObject, verbs: List[str]) -> List[str]:
 def get_supported_verbs(sqobj: SqObject) -> List[str]:
     if sqobj.table == 'routes':
         supported_verbs = [e.value for e in RouteVerbs]
-    elif sqobj.table == 'network':
-        supported_verbs = [e.value for e in NetworkVerbs]
     else:
         if sqobj._valid_assert_args:
             # the assertion is supported for this class
@@ -408,6 +406,9 @@ def get_supported_verbs(sqobj: SqObject) -> List[str]:
         if sqobj.table == 'tables':
             supported_verbs.append('describe')
 
+    if sqobj.table == 'network':
+        supported_verbs += [e.value for e in NetworkVerbs]
+
     return supported_verbs
 
 
@@ -685,6 +686,9 @@ def test_routes_sqobj_consistency():
                 assert_missing_args(top_args, query_params, 'top', table)
             query_params = query_params.difference(top_args)
 
+            if table == 'network' and 'find' not in route.verbs:
+                # do not check deprecated commands. They are already checked
+                continue
             args_to_match = get_args_to_match(sqobj, route.verbs)
             args_to_match = {a for a in args_to_match
                              if a not in common_args.union(top_args)}