trace simple websocket messages - python

[daTobiGit]

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?

[daTobiGit]

same same for aiohttp?

import asyncio
from aiohttp import web
from aws_xray_sdk.ext.aiohttp.middleware import middleware
from aws_xray_sdk.core import xray_recorder
from aws_xray_sdk.core.async_context import AsyncContext
from aws_xray_sdk.core import patch_all

patch_all()

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

# @xray_recorder.capture_async('## websocket_handler')
async def websocket_handler(request):

    ws = web.WebSocketResponse()
    await ws.prepare(request)
    segment = xray_recorder.begin_segment('websocket_connection')
    print(f"segment id:{segment.trace_id}")

    try:
        async for msg in ws:
            subsegment = xray_recorder.begin_segment('process_message')
            
            if msg.type == web.WSMsgType.TEXT:
                await ws.send_str(msg.data)
            elif msg.type == web.WSMsgType.ERROR:
                print(f'WebSocket connection closed with exception {ws.exception()}')
        
            xray_recorder.end_segment()

    except Exception as e:
        print(e)
        segment.add_exception(e)
    
    xray_recorder.end_segment()
    return ws

async def init_app():
    app = web.Application(middlewares=[middleware])
    app.add_routes([web.get('/ws', websocket_handler)])
    return app

if __name__ == '__main__':
    app = init_app()
    web.run_app(app, port=8765)

[daTobiGit]

doesnt even work with asyncio only:

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

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

@xray_recorder.capture_async('## sub1')
async def sub1():
    print("sub1")
    # time.sleep(1)

@xray_recorder.capture_async('## main')
async def main():
    print("main+")
    await sub1()
    print("main-")

if __name__ == '__main__':
    xray_recorder.begin_segment('start')
    asyncio.run(main())
    print("end")

[daTobiGit]

This is the trace:

Traceback (most recent call last):
File "/home/ec2-user/environment/projects/async/async.py", line 26, in
asyncio.run(main())
File "/usr/lib64/python3.9/asyncio/runners.py", line 44, in run
return loop.run_until_complete(main)
File "/usr/lib64/python3.9/asyncio/base_events.py", line 647, in run_until_complete
return future.result()
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/async_recorder.py", line 29, in call
return await self.recorder.record_subsegment_async(
File "/home/ec2-user/.local/lib/python3.9/site-packages/aws_xray_sdk/core/async_recorder.py", line 75, in record_subsegment_async
subsegment = self.begin_subsegment(name, namespace)
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

Cloud9 with AmazonLinux2023

[srprash]

Hi @daTobiGit .
I think it is due to the async nature of these libraries that the parent segment is not present on the same task as the subsegment being created, and I think this is a known bug in the X-Ray SDK.

Since you are just starting out with using X-Ray, I would recommend you to check out the AWS Distro for OpenTelemetry. You can use the ADOT Python to automatically instrument the libraries in your Python application. The current version of the ADOT Python supports the following library instrumentations: https://github.com/open-telemetry/opentelemetry-python-contrib/blob/release/v1.22.x-0.43bx/instrumentation/README.md

[daTobiGit]

Hi @srprash,

thank you for your answer. i will try the other solution, but is there already a bugfix on the roadmap?