Improve MIBUpdater to re-connect DBConnector when re-init data.

src/ax_interface/mib.py

@@ -30,17 +30,29 @@ def __init__(self):
 
     async def start(self):
         # Run the update while we are allowed
+        redis_exception_happen = False
         while self.run_event.is_set():
             try:
                 # reinit internal structures
                 if self.update_counter > self.reinit_rate:
+                    # reconnect when redis exception happen
+                    if redis_exception_happen:
+                        self.reinit_connection()
+
                     self.reinit_data()
                     self.update_counter = 0
                 else:
                     self.update_counter += 1
 
                 # run the background update task
                 self.update_data()
+                redis_exception_happen = False
+            except RuntimeError:
+                # Any unexpected exception or error, log it and keep running
+                logger.exception("MIBUpdater.start() caught an unexpected exception during update_data()")
+                # When redis server restart, swsscommon will throw swsscommon.RedisError, redis connection need re-initialize in reinit_data()
+                # TODO: change to swsscommon.RedisError
+                redis_exception_happen = True
             except Exception:
                 # Any unexpected exception or error, log it and keep running
                 logger.exception("MIBUpdater.start() caught an unexpected exception during update_data()")
@@ -55,6 +67,12 @@ def reinit_data(self):
         """
         return
 
+    def reinit_connection(self):
+        """
+        Reinit redis connection task. Children may override this method.
+        """
+        return
+
     def update_data(self):
         """
         Background task. Children must override this method.

src/sonic_ax_impl/mibs/ietf/rfc1213.py

@@ -134,6 +134,9 @@ def __init__(self):
         self.nexthop_map = {}
         self.route_list = []
 
+    def reinit_connection(self):
+        Namespace.connect_all_dbs(self.db_conn, mibs.APPL_DB)
+
     def update_data(self):
         """
         Update redis (caches config)
@@ -216,6 +219,9 @@ def __init__(self):
 
         self.namespace_db_map = Namespace.get_namespace_db_map(self.db_conn)
 
+    def reinit_connection(self):
+        Namespace.connect_namespace_dbs(self.db_conn)
+
     def reinit_data(self):
         """
         Subclass update interface information

src/sonic_ax_impl/mibs/ietf/rfc2737.py

@@ -267,6 +267,9 @@ def create_physical_entity_updaters(self):
         """
         return [creator(self) for creator in PhysicalTableMIBUpdater.physical_entity_updater_types]
 
+    def reinit_connection(self):
+        Namespace.connect_all_dbs(self.statedb, mibs.STATE_DB)
+
     def reinit_data(self):
         """
         Re-initialize all data.

src/sonic_ax_impl/mibs/ietf/rfc2863.py

@@ -103,6 +103,9 @@ def __init__(self):
 
         self.namespace_db_map = Namespace.get_namespace_db_map(self.db_conn)
 
+    def reinit_connection(self):
+        Namespace.connect_namespace_dbs(self.db_conn)
+
     def reinit_data(self):
         """
         Subclass update interface information

src/sonic_ax_impl/mibs/ietf/rfc3433.py

@@ -374,6 +374,9 @@ def __init__(self):
         self.thermal_sensor = []
         self.broken_transceiver_info = []
 
+    def reinit_connection(self):
+        Namespace.connect_all_dbs(self.statedb, mibs.STATE_DB)
+
     def reinit_data(self):
         """
         Reinit data, clear cache
@@ -385,7 +388,6 @@ def reinit_data(self):
         self.ent_phy_sensor_precision_map = {}
         self.ent_phy_sensor_value_map = {}
         self.ent_phy_sensor_oper_state_map = {}
-
         transceiver_dom_encoded = Namespace.dbs_keys(self.statedb, mibs.STATE_DB, self.TRANSCEIVER_DOM_KEY_PATTERN)
         if transceiver_dom_encoded:
             self.transceiver_dom = [entry for entry in transceiver_dom_encoded]

src/sonic_ax_impl/mibs/ietf/rfc4292.py

@@ -17,6 +17,9 @@ def __init__(self):
         ## loopback ip string -> ip address object
         self.loips = {}
 
+    def reinit_connection(self):
+        Namespace.connect_all_dbs(self.db_conn, mibs.APPL_DB)
+
     def reinit_data(self):
         """
         Subclass update loopback information

src/sonic_ax_impl/mibs/ietf/rfc4363.py

@@ -41,6 +41,9 @@ def fdb_vlanmac(self, fdb):
             return None
         return (int(vlan_id),) + mac_decimals(fdb["mac"])
 
+    def reinit_connection(self):
+        Namespace.connect_namespace_dbs(self.db_conn)
+
     def reinit_data(self):
         """
         Subclass update interface information

src/sonic_ax_impl/mibs/vendor/cisco/ciscoPfcExtMIB.py

@@ -28,6 +28,9 @@ def __init__(self):
         self.if_range = []
         self.namespace_db_map = Namespace.get_namespace_db_map(self.db_conn)
 
+    def reinit_connection(self):
+        Namespace.connect_namespace_dbs(self.db_conn)
+
     def reinit_data(self):
         """
         Subclass update interface information

src/sonic_ax_impl/mibs/vendor/cisco/ciscoSwitchQosMIB.py

@@ -67,6 +67,9 @@ def __init__(self):
         self.port_index_namespace = {}
         self.namespace_db_map = Namespace.get_namespace_db_map(self.db_conn)
 
+    def reinit_connection(self):
+        Namespace.connect_namespace_dbs(self.db_conn)
+
     def reinit_data(self):
         """
         Subclass update interface information

tests/test_rfc1213.py

@@ -1,4 +1,6 @@
+import asyncio
 import os
+import sonic_ax_impl
 import sys
 from unittest import TestCase
 
@@ -10,7 +12,8 @@
 modules_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, os.path.join(modules_path, 'src'))
 
-from sonic_ax_impl.mibs.ietf.rfc1213 import NextHopUpdater
+from sonic_ax_impl.mibs.ietf.rfc1213 import NextHopUpdater, InterfacesUpdater
+
 
 class TestNextHopUpdater(TestCase):
 
@@ -43,3 +46,62 @@ def test_NextHopUpdater_route_no_next_hop(self):
             mocked_warning.assert_has_calls(expected)
 
         self.assertTrue(len(updater.route_list) == 0)
+
+
+class TestNextHopUpdaterRedisException(TestCase):
+    def __init__(self, name):
+        super().__init__(name)
+        self.throw_exception = True
+        self.updater = NextHopUpdater()
+
+    # setup mock method, throw exception when first time call it
+    def mock_dbs_keys(self, *args, **kwargs):
+        if self.throw_exception:
+            self.throw_exception = False
+            raise RuntimeError
+
+        self.updater.run_event.clear()
+        return None
+
+    @mock.patch('sonic_ax_impl.mibs.Namespace.dbs_get_all', mock.MagicMock(return_value=({"ifname": "Ethernet0,Ethernet4"})))
+    def test_NextHopUpdater_redis_exception(self):
+        with mock.patch('sonic_ax_impl.mibs.Namespace.dbs_keys', self.mock_dbs_keys):
+            with mock.patch('ax_interface.logger.exception') as mocked_exception:
+                self.updater.run_event.set()
+                self.updater.frequency = 1
+                self.updater.reinit_rate = 1
+                self.updater.update_counter = 1
+                loop = asyncio.get_event_loop()
+                loop.run_until_complete(self.updater.start())
+                loop.close()
+
+                # check warning
+                expected = [
+                    mock.call("MIBUpdater.start() caught an unexpected exception during update_data()")
+                ]
+                mocked_exception.assert_has_calls(expected)
+
+
+    @mock.patch('sonic_ax_impl.mibs.init_mgmt_interface_tables', mock.MagicMock(return_value=([{}, {}])))
+    def test_InterfacesUpdater_re_init_redis_exception(self):
+
+        def mock_get_sync_d_from_all_namespace(per_namespace_func, db_conn):
+            if per_namespace_func == sonic_ax_impl.mibs.init_sync_d_interface_tables:
+                return [{}, {}, {}, {}]
+
+            if per_namespace_func == sonic_ax_impl.mibs.init_sync_d_vlan_tables:
+                return [{}, {}, {}]
+
+            if per_namespace_func == sonic_ax_impl.mibs.init_sync_d_rif_tables:
+                return [{}, {}]
+
+            return [{}, {}, {}, {}, {}]
+
+        updater = InterfacesUpdater()
+        with mock.patch('sonic_ax_impl.mibs.Namespace.get_sync_d_from_all_namespace', mock_get_sync_d_from_all_namespace):
+            with mock.patch('sonic_ax_impl.mibs.Namespace.connect_namespace_dbs') as connect_namespace_dbs:
+                updater.reinit_connection()
+                updater.reinit_data()
+
+                # check re-init
+                connect_namespace_dbs.assert_called()

tests/test_rfc3433.py

@@ -25,4 +25,15 @@ def test_PhysicalSensorTableMIBUpdater_transceiver_info_key_missing(self):
             # check warning
             mocked_warn.assert_called()
 
-        self.assertTrue(len(updater.sub_ids) == 0)
+        self.assertTrue(len(updater.sub_ids) == 0)
+
+    @mock.patch('sonic_ax_impl.mibs.Namespace.dbs_keys', mock.MagicMock(return_value=(None)))
+    @mock.patch('swsscommon.swsscommon.SonicV2Connector.keys', mock.MagicMock(return_value=(None)))
+    def test_PhysicalSensorTableMIBUpdater_re_init_redis_exception(self):
+        updater = PhysicalSensorTableMIBUpdater()
+
+        with mock.patch('sonic_ax_impl.mibs.Namespace.connect_all_dbs') as connect_all_dbs:
+            updater.reinit_connection()
+
+            # check re-init
+            connect_all_dbs.assert_called()

tests/test_rfc4292.py

@@ -59,3 +59,14 @@ def test_RouteUpdater_route_no_iframe(self):
             mocked_warning.assert_has_calls(expected)
 
         self.assertTrue(len(updater.route_dest_list) == 0)
+
+
+    @mock.patch('sonic_ax_impl.mibs.Namespace.dbs_keys', mock.MagicMock(return_value=(None)))
+    def test_RouteUpdater_re_init_redis_exception(self):
+        updater = RouteUpdater()
+
+        with mock.patch('sonic_ax_impl.mibs.Namespace.connect_all_dbs') as connect_all_dbs:
+            updater.reinit_connection()
+
+            # check re-init
+            connect_all_dbs.assert_called()

tests/test_rfc4363.py

@@ -1,5 +1,6 @@
 import os
 import sys
+import sonic_ax_impl
 from unittest import TestCase
 
 if sys.version_info.major == 3:
@@ -26,3 +27,22 @@ def test_FdbUpdater_ent_bridge_port_id_attr_missing(self):
             mocked_warn.assert_called()
 
         self.assertTrue(len(updater.vlanmac_ifindex_list) == 0)
+
+
+    @mock.patch('sonic_ax_impl.mibs.Namespace.dbs_keys', mock.MagicMock(return_value=(None)))
+    @mock.patch('sonic_ax_impl.mibs.Namespace.dbs_get_bridge_port_map', mock.MagicMock(return_value=(None)))
+    def test_RouteUpdater_re_init_redis_exception(self):
+        updater = FdbUpdater()
+
+        def mock_get_sync_d_from_all_namespace(per_namespace_func, db_conn):
+            if per_namespace_func == sonic_ax_impl.mibs.init_sync_d_interface_tables:
+                return [{}, {}, {}, {}]
+
+            return [{}, {}, {}, {}, {}]
+
+        with mock.patch('sonic_ax_impl.mibs.Namespace.get_sync_d_from_all_namespace', mock_get_sync_d_from_all_namespace):
+            with mock.patch('sonic_ax_impl.mibs.Namespace.connect_namespace_dbs') as connect_namespace_dbs:
+                updater.reinit_connection()
+
+                # check re-init
+                connect_namespace_dbs.assert_called()