Implemented get_stats and clear_stats for Thrift RPC. (#249)

Signed-off-by: xuyinhao <[email protected]>

Author: xuyinhao

common/sai_client/sai_thrift_client/sai_thrift_client.py

@@ -208,6 +208,30 @@ def _operate_attributes(self, operation, attrs=(), oid=None, obj_type=None, key=
 
         return status, result
 
+    def _operate_stats(self, operation, attrs=(), oid=None, obj_type=None):
+        """
+        Get/Clear stats for given object
+        
+        operation (str): operation type, get or clear
+        """
+        if oid is not None:
+            oid = ThriftConverter.object_id(oid)
+
+        obj_type_name = self.get_object_type(oid, default=obj_type).name.lower()
+        object_oid = {f'{obj_type_name}_oid':oid}
+        sai_thrift_function = getattr(sai_adapter, f'sai_thrift_{operation}_{obj_type_name}_stats')
+        
+        result = {}
+
+        attr_kwargs = ThriftConverter.convert_counter_ids_to_thrift(attrs, obj_type_name)
+        result = sai_thrift_function(self.thrift_client, **object_oid, **attr_kwargs)
+        status = ThriftConverter.convert_to_sai_status_str(sai_adapter.status)
+
+        if status != 'SAI_STATUS_SUCCESS':
+            result = None
+
+        return status, result
+
     def cleanup(self):
         if self.thrift_transport:
             self.thrift_transport.close()
@@ -296,5 +320,27 @@ def bulk_set(self, obj_type, keys, attrs, do_assert=True):
                 self.set(obj_type + ":" + json.dumps(keys[i]), attr, do_assert)
         return "SAI_STATUS_SUCCESS", statuses
 
+    def get_stats(self, obj, attrs, do_assert=True):
+        obj_type, oid, _ = self.obj_to_items(obj)
+        status, result = self._operate_stats('get', attrs=attrs, oid=oid, obj_type=obj_type)
+
+        if do_assert:
+            assert status == 'SAI_STATUS_SUCCESS', f"get({obj}, {attrs}) --> {status}"
+        
+        result = [key for pair in result.items() for key in pair]
+        result = SaiData(json.dumps(result))
+
+        if do_assert:
+            return result
+        return status, result
+
+    def clear_stats(self, obj, attrs, do_assert=True):
+        obj_type, oid, _ = self.obj_to_items(obj)
+        status, _ = self._operate_stats('clear', attrs=attrs, oid=oid, obj_type=obj_type)
+
+        if do_assert:
+            assert status == 'SAI_STATUS_SUCCESS', f"clear({obj}, {attrs}) --> {status}"
+        return status
+    
     def get_object_key(self, obj_type=None):
         return dict()

common/sai_client/sai_thrift_client/sai_thrift_utils.py

@@ -6,6 +6,8 @@
 from sai_thrift.ttypes import *
 from sai_thrift import ttypes
 from sai_thrift.sai_headers import *
+from sai_thrift.sai_adapter import *
+import sai_thrift.sai_adapter as adapter
 from saichallenger.common.sai_data import SaiObjType, SaiStatus
 
 
@@ -19,6 +21,14 @@ def convert_attributes_to_thrift(attributes):
                 continue
             yield ThriftConverter.convert_attribute_name_to_thrift(name), ThriftConverter.convert_value_to_thrift(value, attr_name=name)
 
+    def convert_counter_ids_to_thrift(counter_ids, obj_type_name):
+        """
+        [ 'SAI_PORT_STAT_IF_IN_OCTETS', '', 'SAI_PORT_STAT_IF_IN_UCAST_PKTS', '' ] => { "counter_ids": [SAI_PORT_STAT_IF_IN_OCTETS, SAI_PORT_STAT_IF_IN_UCAST_PKTS]) }
+        [ 'SAI_QUEUE_STAT_PACKETS' ] => { "counter_ids": [SAI_QUEUE_STAT_PACKETS]) }
+        """
+        cnt_ids = ThriftConverter.convert_counter_ids_name_to_thrift(counter_ids, obj_type_name)
+        return {"counter_ids": cnt_ids}
+
     def convert_key_to_thrift(object_type, key = None):
         """
         Converts dictionary 'key' to the thrift key entry according to 'object_type':
@@ -59,6 +69,24 @@ def convert_attribute_name_to_thrift(attr):
         """
         return re.search('SAI_.*_ATTR_(.*)', attr).group(1).lower()
 
+    @staticmethod
+    def convert_counter_ids_name_to_thrift(counter_ids, obj_type_name):
+        """
+        [ 'SAI_PORT_STAT_IF_IN_OCTETS', '', 'SAI_PORT_STAT_IF_IN_UCAST_PKTS', '' ] =>  [SAI_PORT_STAT_IF_IN_OCTETS,SAI_PORT_STAT_IF_IN_UCAST_PKTS]
+        """
+        cnts_list = []
+        id_dict_name = "sai_get_{}_stats_counter_ids_dict".format(obj_type_name)
+        id_dict = getattr(adapter, id_dict_name)        
+        reverse_id_dict = {v: k for k, v in id_dict.items()}
+        for counter_id in counter_ids:
+            if counter_id == '':
+                continue
+            if counter_id in reverse_id_dict:
+                cnts_list.append(reverse_id_dict[counter_id])
+            else:
+                raise ValueError("Counter id {} is not supported for {}".format(counter_id, obj_type_name))
+        return cnts_list
+
     @staticmethod
     def convert_u8_to_thrift(u8_str):
         # Thrift does not support unsigned int notation.