code review - few cleanups

Signed-off-by: imarom <imarom@cisco.com>

5 files changed, 89 insertions, 66 deletions

diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py
index 654ceaf6..571334ee 100755
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_client.py
@@ -22,13 +22,11 @@ from collections import namedtuple
 from yaml import YAMLError
 import time
 import datetime
-import threading
 import re
 import random
 import json
 import traceback
 import os.path
-import argparse
 
 ############################     logger     #############################
 ############################                #############################
@@ -3007,32 +3005,42 @@ class STLClient(object):
     @__api_check(True)
     def start_capture (self, tx_ports, rx_ports, limit = 1000, mode = 'fixed'):
         """
-            Starts a capture to PCAP on port(s)
+            Starts a low rate packet capturing on the server
 
             :parameters:
                 tx_ports       - on which ports to capture TX
                 rx_ports       - on which ports to capture RX
                 limit          - limit how many packets will be written
+                                 memory requierment is O(9K * limit)
                 
-                mode           - 'fixed': when full, future packets will be
+                mode           - 'fixed': when full, newer packets will be
                                   dropped
-                                  'cyclic: when full, oldest packets will be
+                                  
+                                  'cyclic: when full, older packets will be
                                   dropped
                                   
             :returns:
-                returns a dictionary containing
+                returns a dictionary:
                 {'id: <new_id>, 'ts': <starting timestamp>}
                 
+                where 'id' is the new capture ID for future commands
+                and 'ts' is that server monotonic timestamp when
+                the capture was created
+                
             :raises:
                 + :exe:'STLError'
 
         """
-        
+        # TODO: remove this when TX is implemented
+        if tx_ports:
+            raise STLError('TX port capturing is not yet implemented')
+            
         # check arguments
         tx_ports = self._validate_port_list(tx_ports, allow_empty = True)
         rx_ports = self._validate_port_list(rx_ports, allow_empty = True)
         merge_ports = set(tx_ports + rx_ports)
         
+        # make sure at least one port to capture
         if not merge_ports:
             raise STLError("start_capture - must get at least one port to capture")
             
@@ -3044,13 +3052,13 @@ class STLClient(object):
             raise STLError("'mode' must be either 'fixed' or 'cyclic'")
         
         # verify service mode
-        non_service_ports =  list_difference(set(tx_ports + rx_ports), self.get_service_enabled_ports())
+        non_service_ports =  list_difference(merge_ports, self.get_service_enabled_ports())
         if non_service_ports:
-            raise STLError("Port(s) {0} are not under service mode. PCAP capturing requires all ports to be in service mode".format(non_service_ports))
+            raise STLError("Port(s) {0} are not under service mode. packet capturing requires all ports to be in service mode".format(non_service_ports))
         
             
         # actual job
-        self.logger.pre_cmd("Starting PCAP capturing up to {0} packets".format(limit))
+        self.logger.pre_cmd("Starting packet capturing up to {0} packets".format(limit))
         rc = self._transmit("capture", params = {'command': 'start', 'limit': limit, 'mode': mode, 'tx': tx_ports, 'rx': rx_ports})
         self.logger.post_cmd(rc)
 
@@ -3059,50 +3067,18 @@ class STLClient(object):
 
         return {'id': rc.data()['capture_id'], 'ts': rc.data()['start_ts']}
 
-        
-                
-    def __fetch_capture_packets (self, capture_id, output_filename, pkt_count):
-        self.logger.pre_cmd("Writing {0} packets to '{1}'".format(pkt_count, output_filename))
 
-        # create a PCAP file
-        writer = RawPcapWriter(output_filename, linktype = 1)
-        writer._write_header(None)
-        
-        # fetch
-        pending = pkt_count
-        rc = RC_OK()
-        while pending > 0:
-            rc = self._transmit("capture", params = {'command': 'fetch', 'capture_id': capture_id, 'pkt_limit': 50})
-            if not rc:
-                self.logger.post_cmd(rc)
-                raise STLError(rc)
-        
-            pkts      = rc.data()['pkts']
-            pending   = rc.data()['pending']
-            start_ts  = rc.data()['start_ts']
-            
-            for pkt in pkts:
-                ts = pkt['ts'] - start_ts
-                ts_sec  = int(ts)
-                ts_usec = int( (ts - ts_sec) * 1e6 )
-                
-                pkt_bin = base64.b64decode(pkt['binary'])
-                writer._write_packet(pkt_bin, sec = ts_sec, usec = ts_usec)
-                
-            
-            
-        
-        self.logger.post_cmd(rc)
-        
         
     @__api_check(True)
     def stop_capture (self, capture_id, output_filename = None):
         """
-            Stops an active capture
+            Stops an active capture and optionally save it to a PCAP file
 
             :parameters:
                 capture_id        - an active capture ID to stop
                 output_filename   - output filename to save capture
+                                    if None all captured packets 
+                                    will be discarded
 
             :raises:
                 + :exe:'STLError'
@@ -3116,7 +3092,7 @@ class STLClient(object):
         
         # stop
         
-        self.logger.pre_cmd("Stopping PCAP capture {0}".format(capture_id))
+        self.logger.pre_cmd("Stopping packet capture {0}".format(capture_id))
         rc = self._transmit("capture", params = {'command': 'stop', 'capture_id': capture_id})
         self.logger.post_cmd(rc)
         if not rc:
@@ -3137,6 +3113,47 @@ class STLClient(object):
             raise STLError(rc)
         
 
+            
+    # fetch packets from the server and save them to a file
+    def __fetch_capture_packets (self, capture_id, output_filename, pkt_count):
+        self.logger.pre_cmd("Writing {0} packets to '{1}'".format(pkt_count, output_filename))
+
+        # create a PCAP file
+        writer = RawPcapWriter(output_filename, linktype = 1)
+        writer._write_header(None)
+
+        pending = pkt_count
+        rc = RC_OK()
+        
+        # fetch with iteratios - each iteration up to 50 packets
+        while pending > 0:
+            rc = self._transmit("capture", params = {'command': 'fetch', 'capture_id': capture_id, 'pkt_limit': 50})
+            if not rc:
+                self.logger.post_cmd(rc)
+                raise STLError(rc)
+
+            # make sure we are getting some progress
+            assert(rc.data()['pending'] < pending)
+            
+            pkts      = rc.data()['pkts']
+            pending   = rc.data()['pending']
+            start_ts  = rc.data()['start_ts']
+            
+            # write packets
+            for pkt in pkts:
+                # split the server timestamp relative to the capture start time
+                ts_sec, ts_usec = sec_split_usec(pkt['ts'] - start_ts)
+                
+                pkt_bin = base64.b64decode(pkt['binary'])
+                writer._write_packet(pkt_bin, sec = ts_sec, usec = ts_usec)
+
+
+
+
+        self.logger.post_cmd(rc)
+
+            
+            
     @__api_check(True)
     def get_capture_status (self):
         """
@@ -3145,14 +3162,13 @@ class STLClient(object):
             info about the capture
 
         """
-
         rc = self._transmit("capture", params = {'command': 'status'})
-
         if not rc:
             raise STLError(rc)
 
         return rc.data()
 
+        
     @__api_check(True)
     def remove_all_captures (self):
         """
@@ -3160,7 +3176,7 @@ class STLClient(object):
         """
         captures = self.get_capture_status()
         
-        self.logger.pre_cmd("Removing all PCAP captures from server")
+        self.logger.pre_cmd("Removing all packet captures from server")
         
         for c in captures:
             # remove
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_jsonrpc_client.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_jsonrpc_client.py
index db216532..405f76be 100644
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_jsonrpc_client.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_jsonrpc_client.py
@@ -10,6 +10,7 @@ import struct
 from .trex_stl_types import *
 from .utils.common import random_id_gen
 from .utils.zipmsg import ZippedMsg
+from threading import Lock
 
 class bcolors:
     BLUE = '\033[94m'
@@ -72,6 +73,8 @@ class JsonRpcClient(object):
         self.id_gen = random_id_gen()
         self.zipper = ZippedMsg()
 
+        self.lock = Lock()
+        
     def get_connection_details (self):
         rc = {}
         rc['server'] = self.server
@@ -137,6 +140,12 @@ class JsonRpcClient(object):
 
    
     def send_msg (self, msg, retry = 0):
+        # REQ/RESP pattern in ZMQ requires no interrupts during the send
+        with self.lock:
+            return self.__send_msg(msg, retry)
+        
+        
+    def __send_msg (self, msg, retry = 0):
         # print before
         if self.logger.check_verbose(self.logger.VERBOSE_HIGH):
             self.verbose_msg("Sending Request To Server:\n\n" + self.pretty_json(msg) + "\n")
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_port.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_port.py
index 31d752af..1ef3a8ff 100644
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_port.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/trex_stl_port.py
@@ -492,22 +492,6 @@ class Port(object):
         return self.ok()
 
     
-    @owned
-    def start_capture (self, pcap_filename, mode, limit):
-
-        if mode != 'tx' and not self.is_service_mode_on():
-            return self.err('port service mode must be enabled for performing RX capturing. Please enable service mode')
-            
-        params = {"handler":        self.handler,
-                  "port_id":        self.port_id,
-                  "mode":           mode,
-                  "limit":          limit}
-
-        rc = self.transmit("start_capture", params)
-        if rc.bad():
-            return self.err(rc.err())
-
-        return self.ok()
      
     @writeable
     def set_l2_mode (self, dst_mac):
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py
index c386451b..72d3fa9f 100644
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/common.py
@@ -71,6 +71,11 @@ def list_difference (l1, l2):
 def is_sub_list (l1, l2):
     return set(l1) <= set(l2)
 
+# splits a timestamp in seconds to sec/usec
+def sec_split_usec (ts):
+    return int(ts), int( (ts - int(ts)) * 1e6 )
+    
+    
 # a simple passive timer
 class PassiveTimer(object):
 
diff --git a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py
index 8d3aedbe..53db533c 100755
--- a/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py
+++ b/scripts/automation/trex_control_plane/stl/trex_stl_lib/utils/parsing_opts.py
@@ -689,19 +689,23 @@ class _MergeAction(argparse._AppendAction):
 
 class CCmdArgParser(argparse.ArgumentParser):
 
-    def __init__(self, stateless_client = None, x = None, *args, **kwargs):
+    def __init__(self, stateless_client = None, *args, **kwargs):
         super(CCmdArgParser, self).__init__(*args, **kwargs)
         self.stateless_client = stateless_client
         self.cmd_name = kwargs.get('prog')
         self.register('action', 'merge', _MergeAction)
 
 
+        
     def add_arg_list (self, *args):
         populate_parser(self, *args)
 
+        
+    # a simple hook for add subparsers to add stateless client
     def add_subparsers(self, *args, **kwargs):
         sub = super(CCmdArgParser, self).add_subparsers(*args, **kwargs)
 
+        # save pointer to the original add parser method
         add_parser = sub.add_parser
         stateless_client = self.stateless_client
 
@@ -710,13 +714,17 @@ class CCmdArgParser(argparse.ArgumentParser):
             parser.stateless_client = stateless_client
             return parser
 
+        # override with the hook
         sub.add_parser = add_parser_hook
+        
         return sub
 
+        
     # hook this to the logger
     def _print_message(self, message, file=None):
         self.stateless_client.logger.log(message)
 
+        
     def error(self, message):
         self.print_usage()
         self._print_message(('%s: error: %s\n') % (self.prog, message))
@@ -783,6 +791,7 @@ class CCmdArgParser(argparse.ArgumentParser):
             # recover from system exit scenarios, such as "help", or bad arguments.
             return RC_ERR("'{0}' - {1}".format(self.cmd_name, "no action"))
 
+
     def formatted_error (self, msg):
         self.print_usage()
         self._print_message(('%s: error: %s\n') % (self.prog, msg))