Add performance tracing

catkit2/testbed/__init__.py

@@ -5,12 +5,19 @@
  'CatkitLogHandler',
  'TestbedProxy',
  'ServiceProxy',
- 'Experiment'
+ 'Experiment',
+ 'TraceWriter',
+ 'trace_interval',
+ 'trace_instant',
+ 'trace_counter',
+ 'ZmqDistributor',
 ]
 
 from .testbed import *
 from .experiment import *
 from .service import *
 from .logging import *
+from .tracing import *
+from .distributor import *
 from .testbed_proxy import *
 from .service_proxy import *

catkit2/testbed/distributor.py

@@ -0,0 +1,86 @@
+import zmq
+import threading
+import traceback
+
+
+class ZmqDistributor:
+ '''Collects messages on a port and re-publish them on another.
+
+ This operates on a separate thread after it is started.
+
+ Parameters
+ ----------
+ context : zmq.Context
+ A previously-created ZMQ context. All sockets will be created on this context.
+ input_port : integer
+ The port number for the incoming log messages.
+ output_port : integer
+ The port number for the outgoing log messages.
+ callback : function
+ A callback to call with each message.
+ '''
+ def __init__(self, context, input_port, output_port, callback=None):
+ self.context = context
+ self.input_port = input_port
+ self.output_port = output_port
+ self.callback = callback
+
+ self.shutdown_flag = threading.Event()
+ self.thread = None
+
+ self.is_running = threading.Event()
+
+ def start(self):
+ '''Start the proxy thread.
+ '''
+ self.thread = threading.Thread(target=self._forwarder)
+ self.thread.start()
+
+ self.is_running.wait()
+
+ def stop(self):
+ '''Stop the proxy thread.
+
+ This function waits until the thread is actually stopped.
+ '''
+ self.shutdown_flag.set()
+
+ if self.thread:
+ self.thread.join()
+
+ self.is_running.clear()
+
+ def _forwarder(self):
+ '''Create sockets and republish all received log messages.
+
+ .. note::
+ This function should not be called directly. Use
+ :func:`~catkit2.testbed.ZmqDistributor.start()` to start the proxy.
+ '''
+ collector = self.context.socket(zmq.PULL)
+ collector.RCVTIMEO = 50
+ collector.bind(f'tcp://*:{self.input_port}')
+
+ publicist = self.context.socket(zmq.PUB)
+ publicist.bind(f'tcp://*:{self.output_port}')
+
+ self.is_running.set()
+
+ while not self.shutdown_flag.is_set():
+ try:
+ try:
+ log_message = collector.recv_multipart()
+ publicist.send_multipart(log_message)
+ except zmq.ZMQError as e:
+ if e.errno == zmq.EAGAIN:
+ # Timed out.
+ continue
+ else:
+ raise RuntimeError('Error during receive') from e
+
+ if self.callback:
+ self.callback(log_message)
+ except Exception:
+ # Something went wrong during handling of the log message.
+ # Let's ignore this error, but still print the exception.
+ print(traceback.format_exc())

catkit2/testbed/logging.py

@@ -3,7 +3,6 @@
 import zmq
 import json
 import contextlib
-import traceback
 from colorama import Fore, Back, Style
 
 from ..catkit_bindings import submit_log_entry, Severity
@@ -27,79 +26,6 @@ def emit(self, record):
 
  submit_log_entry(filename, line, function, severity, message)
 
-class LogDistributor:
- '''Collects log messages on a port and re-publish them on another.
-
- This operates on a separate thread after it is started.
-
- Parameters
- ----------
- context : zmq.Context
- A previously-created ZMQ context. All sockets will be created on this context.
- input_port : integer
- The port number for the incoming log messages.
- output_port : integer
- The port number for the outgoing log messages.
- '''
- def __init__(self, context, input_port, output_port):
- self.context = context
- self.input_port = input_port
- self.output_port = output_port
-
- self.shutdown_flag = threading.Event()
- self.thread = None
-
- def start(self):
- '''Start the proxy thread.
- '''
- self.thread = threading.Thread(target=self.forwarder)
- self.thread.start()
-
- def stop(self):
- '''Stop the proxy thread.
-
- This function waits until the thread is actually stopped.
- '''
- self.shutdown_flag.set()
-
- if self.thread:
- self.thread.join()
-
- def forwarder(self):
- '''Create sockets and republish all received log messages.
-
- .. note::
- This function should not be called directly. Use
- :func:`~catkit2.testbed.LoggingProxy.start` to start the proxy.
- '''
- collector = self.context.socket(zmq.PULL)
- collector.RCVTIMEO = 50
- collector.bind(f'tcp://*:{self.input_port}')
-
- publicist = self.context.socket(zmq.PUB)
- publicist.bind(f'tcp://*:{self.output_port}')
-
- while not self.shutdown_flag.is_set():
- try:
- try:
- log_message = collector.recv_multipart()
- publicist.send_multipart(log_message)
- except zmq.ZMQError as e:
- if e.errno == zmq.EAGAIN:
- # Timed out.
- continue
- else:
- raise RuntimeError('Error during receive') from e
-
- log_message = log_message[0].decode('utf-8')
- log_message = json.loads(log_message)
-
- print(f'[{log_message["service_id"]}] {log_message["message"]}')
- except Exception:
- # Something went wrong during handling of the log message.
- # Let's ignore this error, but still print the exception.
- print(traceback.format_exc())
-
 class LogObserver:
  def __init__(self, host, port):
  self.context = zmq.Context()

catkit2/testbed/testbed.py

@@ -13,6 +13,7 @@
 
 from ..catkit_bindings import LogForwarder, Server, ServiceState, DataStream, get_timestamp, is_alive_state, Client
 from .logging import *
+from .distributor import ZmqDistributor
 
 from ..proto import testbed_pb2 as testbed_proto
 from ..proto import service_pb2 as service_proto
@@ -167,11 +168,13 @@ def __init__(self, port, is_simulated, config):
  self.config = config
 
  self.services = {}
+ self.launched_processes = []
 
  self.log_distributor = None
  self.log_handler = None
  self.log_forwarder = None
- self.launched_processes = []
+
+ self.tracing_distributor = None
 
  self.log = logging.getLogger(__name__)
 
@@ -271,6 +274,9 @@ def run(self):
  self.start_log_distributor()
  self.setup_logging()
 
+ # Start tracing distributor.
+ self.start_tracing_distributor()
+
  heartbeat_thread = threading.Thread(target=self.do_heartbeats)
  heartbeat_thread.start()
 
@@ -316,6 +322,9 @@ def run(self):
  # Shut down the server.
  self.server.stop()
 
+ # Stop tracing distributor.
+ self.stop_tracing_distributor()
+
  # Stop the logging.
  self.destroy_logging()
  self.stop_log_distributor()
@@ -392,9 +401,15 @@ def destroy_logging(self):
  def start_log_distributor(self):
  '''Start the log distributor.
  '''
+ def callback(log_message):
+ log_message = log_message[0].decode('utf-8')
+ log_message = json.loads(log_message)
+
+ print(f'[{log_message["service_id"]}] {log_message["message"]}')
+
  self.logging_ingress_port, self.logging_egress_port = get_unused_port(num_ports=2)
 
- self.log_distributor = LogDistributor(self.context, self.logging_ingress_port, self.logging_egress_port)
+ self.log_distributor = ZmqDistributor(self.context, self.logging_ingress_port, self.logging_egress_port, callback)
  self.log_distributor.start()
 
  def stop_log_distributor(self):
@@ -404,6 +419,19 @@ def stop_log_distributor(self):
  self.log_distributor.stop()
  self.log_distributor = None
 
+ def start_tracing_distributor(self):
+ '''Start the tracing distributor.
+ '''
+ self.tracing_ingress_port, self.tracing_egress_port = get_unused_port(num_ports=2)
+
+ self.tracing_distributor = ZmqDistributor(self.context, self.tracing_ingress_port, self.tracing_egress_port)
+ self.tracing_distributor.start()
+
+ def stop_tracing_distributor(self):
+ if self.tracing_distributor:
+ self.tracing_distributor.stop()
+ self.tracing_distributor = None
+
  def on_start_service(self, data):
  request = testbed_proto.StartServiceRequest()
  request.ParseFromString(data)