Fix: use add_trace_field() over add_context() for tracer (#21)

This way all events are marked with the command they operate under,
instead of only the root event. This makes sampling a lot easier,
as we know the command we are doing it for.

Author: TrueBrain

openttd_protocol/tracer.py

@@ -10,7 +10,7 @@ def wrapper(func):
 
     tracer = beeline.tracer
     untraced = beeline.untraced
-    add_context = beeline.add_context
+    add_trace_field = beeline.add_trace_field
 
 except ImportError:
     # Honeycomb Beeline package is not installed. Mock the tracer functions.
@@ -34,5 +34,5 @@ def __exit__(self, type, value, traceback):
     def untraced(func):
         return func
 
-    def add_context(payload):
+    def add_trace_field(key, value):
         pass

openttd_protocol/wire/tcp.py

@@ -181,8 +181,6 @@ async def _process_queue(self):
                 log.info("Dropping invalid packet from %s:%d: %r", self.source.ip, self.source.port, err)
                 raise SocketClosed
 
-            tracer.add_context({"command": f"receive.{packet_type.name}"})
-
             with tracer.tracer("tcp.packet-handler"):
                 await getattr(self._callback, f"receive_{packet_type.name}")(self.source, **kwargs)
 
@@ -194,19 +192,21 @@ def receive_packet(self, source, data):
             raise PacketInvalidSize(len(data) + 2, length)
 
         # Check if type is in range
-        type, data = read_uint8(data)
-        if type >= self.PACKET_END:
-            raise PacketInvalidType(type)
+        packet_type, data = read_uint8(data)
+        if packet_type >= self.PACKET_END:
+            raise PacketInvalidType(packet_type)
 
         # Check if we expect this packet
-        type = self.PacketType(type)
-        func = getattr(self, f"receive_{type.name}", None)
+        packet_type = self.PacketType(packet_type)
+        func = getattr(self, f"receive_{packet_type.name}", None)
         if func is None:
-            raise PacketInvalidType(type)
+            raise PacketInvalidType(packet_type)
+
+        tracer.add_trace_field("command", f"receive.{packet_type.name}")
 
         # Process this packet
         kwargs = func(source, data)
-        return type, kwargs
+        return packet_type, kwargs
 
     @tracer.traced("tcp")
     async def send_packet(self, data):