Update NetworkTracker to track time and allow setting of the PID to t…

…rack

src/nwb_benchmarks/core/_network_tracker.py

@@ -12,22 +12,21 @@
 
 
 @contextlib.contextmanager
-def network_activity_tracker(tshark_path: Union[pathlib.Path, None] = None):
-    """Context manager for tracking network activity and statistics for the code executed in the context"""
+def network_activity_tracker(tshark_path: Union[pathlib.Path, None] = None, pid: int = None):
+    """
+    Context manager for tracking network activity and statistics for the code executed in the context
+
+    :param tshark_path: Path to the tshark CLI command to use for tracking network traffic
+    :param pid: The id of the process to compute the network statistics for. If set to None, then the
+                 PID of the current process will be used.
+    """
     network_tracker = NetworkTracker()
 
     try:
         network_tracker.start_network_capture(tshark_path=tshark_path)
-        time.sleep(0.3)
-
-        t0 = time.time()
         yield network_tracker
     finally:
-        network_tracker.stop_network_capture()
-
-        t1 = time.time()
-        network_total_time = t1 - t0
-        network_tracker.network_statistics["network_total_time_in_seconds"] = network_total_time
+        network_tracker.stop_network_capture(pid=pid)
 
 
 class NetworkTracker:
@@ -52,11 +51,14 @@ def __init__(self):
         self.pid_packets = None
         self.network_statistics = None
         self.asv_network_statistics = None
+        self.__start_capture_time = None
 
     def start_network_capture(self, tshark_path: Union[pathlib.Path, None] = None):
         """
         Start capturing the connections on this machine as well as all network packets
 
+        :param tshark_path: Path to the tshark CLI command to use for tracking network traffic
+
         Side effects: This functions sets the following instance variables:
         * self.connections_thread
         * self.network_profile
@@ -69,10 +71,16 @@ def start_network_capture(self, tshark_path: Union[pathlib.Path, None] = None):
         self.network_profiler = NetworkProfiler()
         self.network_profiler.start_capture(tshark_path=tshark_path)
 
-    def stop_network_capture(self):
+        # start the main timer
+        self.__start_capture_time = time.time()
+
+    def stop_network_capture(self, pid: int = None):
         """
         Stop capturing network packets and connections.
 
+        :param pid: The id of the process to compute the network statistics for. If set to None, then the
+                    PID of the current process (i.e., os.getpid()) will be used.
+
         Note: This function will fail if `start_network_capture` was not called first.
 
         Side effects: This functions sets the following instance variables:
@@ -81,15 +89,23 @@ def stop_network_capture(self):
         * self.network_statistics
         * self.asv_network_statistics
         """
+        # stop capturing the network
         self.network_profiler.stop_capture()
         self.connections_thread.stop()
 
-        # get the connections for the PID of this process
-        self.pid_connections = self.connections_thread.get_connections_for_pid(os.getpid())
+        # compute the total time
+        stop_capture_time = time.time()
+        network_total_time = stop_capture_time - self.__start_capture_time
+
+        # get the connections for the PID of this process or the PID set by the user
+        if pid is None:
+            pid = os.getpid()
+        self.pid_connections = self.connections_thread.get_connections_for_pid(pid)
         # Parse packets and filter out all the packets for this process pid by matching with the pid_connections
         self.pid_packets = self.network_profiler.get_packets_for_connections(self.pid_connections)
         # Compute all the network statistics
         self.network_statistics = NetworkStatistics.get_statistics(packets=self.pid_packets)
+        self.network_statistics["network_total_time_in_seconds"] = network_total_time
 
         # Very special structure required by ASV
         # 'samples' is the value tracked in our results