[BFN] Canceling PSU platform API calls on SIGTERM

Signed-off-by: Andriy Kokhan <[email protected]>

platform/barefoot/sonic-platform-modules-bfn-montara/sonic_platform/platform_utils.py

@@ -3,6 +3,8 @@
 try:
     import os
     import subprocess
+    import signal
+    from functools import wraps
 
 except ImportError as e:
     raise ImportError(str(e) + "- required module not found")
@@ -20,3 +22,24 @@ def run_cmd(cmd):
         run_cmd(['touch', path])
     if (mode is not None):    
         run_cmd(['chmod', mode, path])
+
+def cancel_on_sigterm(func):
+    """
+    Wrapper for a function which has to be cancel on SIGTERM
+    """
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        def handler(sig, frame):
+            if sigterm_handler:
+                sigterm_handler(sig, frame)
+            raise Exception("Canceling {}() execution...".format(func.__name__))
+
+        sigterm_handler = signal.getsignal(signal.SIGTERM)
+        signal.signal(signal.SIGTERM, handler)
+        result = None
+        try:
+            result = func(*args, **kwargs)
+        finally:
+            signal.signal(signal.SIGTERM, sigterm_handler)
+        return result
+    return wrapper

platform/barefoot/sonic-platform-modules-bfn-montara/sonic_platform/psu.py

@@ -4,18 +4,26 @@
     import os
     import sys
     import time
+    import signal
+    import syslog
 
     sys.path.append(os.path.dirname(__file__))
 
     from .platform_thrift_client import thrift_try
 
     from sonic_platform_base.psu_base import PsuBase
+    from platform_utils import cancel_on_sigterm
+
 except ImportError as e:
     raise ImportError (str(e) + "- required module not found")
 
 class Psu(PsuBase):
     """Platform-specific PSU class"""
 
+    sigterm = False
+    sigterm_default_handler = None
+    cls_inited = False
+
     def __init__(self, index):
         PsuBase.__init__(self)
         self.__index = index
@@ -24,6 +32,21 @@ def __init__(self, index):
         # STUB IMPLEMENTATION
         self.color = ""
 
+        syslog.syslog(syslog.LOG_INFO, "Created PSU #{} instance".format(self.__index))
+        if not Psu.cls_inited:
+            Psu.sigterm_default_handler = signal.getsignal(signal.SIGTERM)
+            signal.signal(signal.SIGTERM, Psu.signal_handler)
+            if Psu.sigterm_default_handler:
+                syslog.syslog(syslog.LOG_INFO, "Default SIGTERM handler overridden!!")
+            Psu.cls_inited = True
+
+    @classmethod
+    def signal_handler(cls, sig, frame):
+        if cls.sigterm_default_handler:
+            cls.sigterm_default_handler(sig, frame)
+        syslog.syslog(syslog.LOG_INFO, "Canceling PSU platform API calls...")
+        cls.sigterm = True
+
     '''
     Units of returned info object values:
         vin - V
@@ -33,20 +56,23 @@ def __init__(self, index):
         fspeed - RPM
     '''
     def __info_get(self):
+        @cancel_on_sigterm
         def psu_info_get(client):
             return client.pltfm_mgr.pltfm_mgr_pwr_supply_info_get(self.__index)
 
         # Update cache once per 2 seconds
-        if self.__ts + 2 < time.time():
+        if self.__ts + 2 < time.time() and not Psu.sigterm:
             self.__info = None
             try:
                 self.__info = thrift_try(psu_info_get, attempts=1)
+            except Exception as e:
+                if "Canceling" in str(e):
+                    syslog.syslog(syslog.LOG_INFO, "{}".format(e))
             finally:
                 self.__ts = time.time()
                 return self.__info
         return self.__info
 
-
     @staticmethod
     def get_num_psus():
         """