trace simple websocket messages - python

Hello XRay Team,

im having a project in python which uses websockets library. now we want to use xray to get some insights and i have to instrument it. as a starting point, i decided to set up a sample project to get to know xray, but it seems i cannot make it work.

when a client connects, im getting an error:

connection handler failed
Traceback (most recent call last):
File "/home/ec2-user/environment/projects/xray/websocket.py", line 25, in echo
subsegment = xray_recorder.begin_subsegment('process_message')
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 313, in begin_subsegment
return self._begin_subsegment_helper(name, namespace)
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 287, in _begin_subsegment_helper
segment = self.current_segment()
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 271, in current_segment
entity = self.get_trace_entity()
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/recorder.py", line 411, in get_trace_entity
return self.context.get_trace_entity()
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/context.py", line 101, in get_trace_entity
return self.handle_context_missing()
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/context.py", line 125, in handle_context_missing
raise SegmentNotFoundException(MISSING_SEGMENT_MSG)
aws_xray_sdk.core.exceptions.exceptions.SegmentNotFoundException: cannot find the current segment/subsegment, please make sure you have a segment open

#!/usr/bin/env python

import asyncio
from websockets.server import serve

from aws_xray_sdk.core.async_context import AsyncContext
from aws_xray_sdk.core import xray_recorder

# Configure the X-Ray recorder
xray_recorder.configure(
    service='websocket', 
    context=AsyncContext(), 
    context_missing='RUNTIME_ERROR', 
    daemon_address='localhost:2000'
)

async def echo(websocket):
    segment = xray_recorder.begin_segment('echo')

    try:
        async for message in websocket:
            subsegment = xray_recorder.begin_subsegment('process_message')

            try:
                await websocket.send(message)
            except Exception as e:
                subsegment.add_exception(e)
            finally:
                xray_recorder.end_subsegment()

    except Exception as e:
        segment.add_exception(e)
    finally:
        xray_recorder.end_segment()

async def main():
    async with serve(echo, "localhost", 8765):
        await asyncio.Future()  # run forever

asyncio.run(main())

what am i doing wrong?

Reproduce el fallo de conexión de websocket.py, especialmente echo y las llamadas a begin_segment/begin_subsegment mostradas en el issue. Lee aws_xray_sdk.core.async_context.AsyncContext y los métodos de contexto del recorder relacionados con current_segment y get_trace_entity. La tarea está terminada cuando un cliente puede intercambiar mensajes websocket sin SegmentNotFoundException y los mensajes quedan trazados.