[BFN] Canceling PSU platform API calls on SIGTERM

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

@@ -6,6 +6,7 @@
     import json
     from collections import OrderedDict
     from sonic_py_common import device_info
+    from platform_utils import limit_execution_time
 
 except ImportError as e:
     raise ImportError(str(e) + "- required module not found")
@@ -24,6 +25,7 @@ def get_bios_version():
     except subprocess.CalledProcessError as e:
         raise RuntimeError("Failed to get BIOS version")
 
+@limit_execution_time(1)
 def get_bmc_version():
     """
     Retrieves the firmware version of the BMC

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

@@ -3,11 +3,19 @@
 try:
     import os
     import subprocess
+    import signal
+    from functools import wraps
 
 except ImportError as e:
     raise ImportError(str(e) + "- required module not found")
 
 def file_create(path, mode=None):
+    """
+    Ensure that file is created with the appropriate permissions
+    Args:
+        path: full path of a file
+        mode: file permission in octal representation  
+    """
     def run_cmd(cmd):
         if os.geteuid() != 0:
             cmd.insert(0, 'sudo')
@@ -18,5 +26,55 @@ def run_cmd(cmd):
         run_cmd(['mkdir', '-p', file_path])
     if not os.path.isfile(path):
         run_cmd(['touch', path])
-    if (mode is not None):    
+    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
+
+def limit_execution_time(execution_time_secs: int):
+    """
+    Wrapper for a function whose execution time must be limited
+    Args:
+        execution_time_secs: maximum execution time in seconds,
+        after which the function execution will be stopped
+    """
+    def wrapper(func):
+        @wraps(func)
+        def execution_func(*args, **kwargs):
+            def handler(sig, frame):
+                if sigalrm_handler:
+                    sigalrm_handler(sig, frame)
+                raise Exception("Canceling {}() execution...".format(func.__name__))
+
+            sigalrm_handler = signal.getsignal(signal.SIGALRM)
+            signal.signal(signal.SIGALRM, handler)
+            signal.alarm(execution_time_secs)
+            result = None
+            try:
+                result = func(*args, **kwargs)
+            finally:
+                signal.alarm(0)
+
+            return result
+        return execution_func
+    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():
         """