From 1ccb63c45ed8a841d48c4815882b9e839f862724 Mon Sep 17 00:00:00 2001 From: Neel Shah Date: Tue, 19 May 2026 16:04:50 +0200 Subject: [PATCH 1/2] feat(socket): Support span streaming --- sentry_sdk/integrations/socket.py | 92 ++++++++--- tests/integrations/socket/test_socket.py | 185 ++++++++++++++++------- 2 files changed, 196 insertions(+), 81 deletions(-) diff --git a/sentry_sdk/integrations/socket.py b/sentry_sdk/integrations/socket.py index 02d7075e00..2e4464730d 100644 --- a/sentry_sdk/integrations/socket.py +++ b/sentry_sdk/integrations/socket.py @@ -2,8 +2,9 @@ import sentry_sdk from sentry_sdk._types import MYPY -from sentry_sdk.consts import OP +from sentry_sdk.consts import OP, SPANDATA from sentry_sdk.integrations import Integration +from sentry_sdk.tracing_utils import has_span_streaming_enabled if MYPY: from socket import AddressFamily, SocketKind @@ -50,22 +51,39 @@ def create_connection( timeout: "Optional[float]" = socket._GLOBAL_DEFAULT_TIMEOUT, # type: ignore source_address: "Optional[Tuple[Union[bytearray, bytes, str], int]]" = None, ) -> "socket.socket": - integration = sentry_sdk.get_client().get_integration(SocketIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(SocketIntegration) if integration is None: return real_create_connection(address, timeout, source_address) - with sentry_sdk.start_span( - op=OP.SOCKET_CONNECTION, - name=_get_span_description(address[0], address[1]), - origin=SocketIntegration.origin, - ) as span: - span.set_data("address", address) - span.set_data("timeout", timeout) - span.set_data("source_address", source_address) - - return real_create_connection( - address=address, timeout=timeout, source_address=source_address - ) + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=_get_span_description(address[0], address[1]), + attributes={ + "sentry.op": OP.SOCKET_CONNECTION, + "sentry.origin": SocketIntegration.origin, + }, + ) as span: + if address[0] is not None: + span.set_attribute(SPANDATA.SERVER_ADDRESS, address[0]) + span.set_attribute(SPANDATA.SERVER_PORT, address[1]) + + return real_create_connection( + address=address, timeout=timeout, source_address=source_address + ) + else: + with sentry_sdk.start_span( + op=OP.SOCKET_CONNECTION, + name=_get_span_description(address[0], address[1]), + origin=SocketIntegration.origin, + ) as span: + span.set_data("address", address) + span.set_data("timeout", timeout) + span.set_data("source_address", source_address) + + return real_create_connection( + address=address, timeout=timeout, source_address=source_address + ) socket.create_connection = create_connection # type: ignore @@ -81,18 +99,44 @@ def getaddrinfo( proto: int = 0, flags: int = 0, ) -> "List[Tuple[AddressFamily, SocketKind, int, str, Union[Tuple[str, int], Tuple[str, int, int, int], Tuple[int, bytes]]]]": - integration = sentry_sdk.get_client().get_integration(SocketIntegration) + client = sentry_sdk.get_client() + integration = client.get_integration(SocketIntegration) if integration is None: return real_getaddrinfo(host, port, family, type, proto, flags) - with sentry_sdk.start_span( - op=OP.SOCKET_DNS, - name=_get_span_description(host, port), - origin=SocketIntegration.origin, - ) as span: - span.set_data("host", host) - span.set_data("port", port) - - return real_getaddrinfo(host, port, family, type, proto, flags) + if has_span_streaming_enabled(client.options): + with sentry_sdk.traces.start_span( + name=_get_span_description(host, port), + attributes={ + "sentry.op": OP.SOCKET_DNS, + "sentry.origin": SocketIntegration.origin, + }, + ) as span: + if isinstance(host, str): + span.set_attribute(SPANDATA.SERVER_ADDRESS, host) + elif isinstance(host, bytes): + span.set_attribute( + SPANDATA.SERVER_ADDRESS, host.decode(errors="replace") + ) + if isinstance(port, int): + span.set_attribute(SPANDATA.SERVER_PORT, port) + elif isinstance(port, str): + span.set_attribute(SPANDATA.SERVER_PORT, port) + elif isinstance(port, bytes): + span.set_attribute( + SPANDATA.SERVER_PORT, port.decode(errors="replace") + ) + + return real_getaddrinfo(host, port, family, type, proto, flags) + else: + with sentry_sdk.start_span( + op=OP.SOCKET_DNS, + name=_get_span_description(host, port), + origin=SocketIntegration.origin, + ) as span: + span.set_data("host", host) + span.set_data("port", port) + + return real_getaddrinfo(host, port, family, type, proto, flags) socket.getaddrinfo = getaddrinfo diff --git a/tests/integrations/socket/test_socket.py b/tests/integrations/socket/test_socket.py index cc109e0968..158f85a4b7 100644 --- a/tests/integrations/socket/test_socket.py +++ b/tests/integrations/socket/test_socket.py @@ -1,5 +1,8 @@ import socket +import pytest + +import sentry_sdk from sentry_sdk import start_transaction from sentry_sdk.integrations.socket import SocketIntegration from tests.conftest import ApproxDict, create_mock_http_server @@ -7,75 +10,143 @@ PORT = create_mock_http_server() -def test_getaddrinfo_trace(sentry_init, capture_events): - sentry_init(integrations=[SocketIntegration()], traces_sample_rate=1.0) - events = capture_events() - - with start_transaction(): - socket.getaddrinfo("localhost", PORT) - - (event,) = events - (span,) = event["spans"] - - assert span["op"] == "socket.dns" - assert span["description"] == f"localhost:{PORT}" # noqa: E231 - assert span["data"] == ApproxDict( - { - "host": "localhost", - "port": PORT, - } +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_getaddrinfo_trace(sentry_init, capture_events, capture_items, span_streaming): + sentry_init( + integrations=[SocketIntegration()], + traces_sample_rate=1.0, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - -def test_create_connection_trace(sentry_init, capture_events): + if span_streaming: + items = capture_items("span") + with sentry_sdk.traces.start_span(name="root"): + socket.getaddrinfo("localhost", PORT) + sentry_sdk.flush() + + spans = [item.payload for item in items if item.type == "span"] + dns_span, _root = spans + + assert dns_span["attributes"]["sentry.op"] == "socket.dns" + assert dns_span["attributes"]["sentry.origin"] == "auto.socket.socket" + assert dns_span["name"] == f"localhost:{PORT}" # noqa: E231 + assert dns_span["attributes"]["server.address"] == "localhost" + assert dns_span["attributes"]["server.port"] == PORT + else: + events = capture_events() + + with start_transaction(): + socket.getaddrinfo("localhost", PORT) + + (event,) = events + (span,) = event["spans"] + + assert span["op"] == "socket.dns" + assert span["description"] == f"localhost:{PORT}" # noqa: E231 + assert span["data"] == ApproxDict( + { + "host": "localhost", + "port": PORT, + } + ) + + +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_create_connection_trace( + sentry_init, capture_events, capture_items, span_streaming +): timeout = 10 - sentry_init(integrations=[SocketIntegration()], traces_sample_rate=1.0) - events = capture_events() - - with start_transaction(): - socket.create_connection(("localhost", PORT), timeout, None) - - (event,) = events - (connect_span, dns_span) = event["spans"] - # as getaddrinfo gets called in create_connection it should also contain a dns span - - assert connect_span["op"] == "socket.connection" - assert connect_span["description"] == f"localhost:{PORT}" # noqa: E231 - assert connect_span["data"] == ApproxDict( - { - "address": ["localhost", PORT], - "timeout": timeout, - "source_address": None, - } - ) - - assert dns_span["op"] == "socket.dns" - assert dns_span["description"] == f"localhost:{PORT}" # noqa: E231 - assert dns_span["data"] == ApproxDict( - { - "host": "localhost", - "port": PORT, - } + sentry_init( + integrations=[SocketIntegration()], + traces_sample_rate=1.0, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - -def test_span_origin(sentry_init, capture_events): + if span_streaming: + items = capture_items("span") + with sentry_sdk.traces.start_span(name="root"): + socket.create_connection(("localhost", PORT), timeout, None) + sentry_sdk.flush() + + spans = [item.payload for item in items if item.type == "span"] + # as getaddrinfo gets called in create_connection it should also contain a dns span + # spans finish in order: dns (inner) ends first, connect ends, then root + dns_span, connect_span, _root = spans + + assert connect_span["attributes"]["sentry.op"] == "socket.connection" + assert connect_span["name"] == f"localhost:{PORT}" # noqa: E231 + assert connect_span["attributes"]["server.address"] == "localhost" + assert connect_span["attributes"]["server.port"] == PORT + + assert dns_span["attributes"]["sentry.op"] == "socket.dns" + assert dns_span["name"] == f"localhost:{PORT}" # noqa: E231 + assert dns_span["attributes"]["server.address"] == "localhost" + assert dns_span["attributes"]["server.port"] == PORT + else: + events = capture_events() + + with start_transaction(): + socket.create_connection(("localhost", PORT), timeout, None) + + (event,) = events + (connect_span, dns_span) = event["spans"] + # as getaddrinfo gets called in create_connection it should also contain a dns span + + assert connect_span["op"] == "socket.connection" + assert connect_span["description"] == f"localhost:{PORT}" # noqa: E231 + assert connect_span["data"] == ApproxDict( + { + "address": ["localhost", PORT], + "timeout": timeout, + "source_address": None, + } + ) + + assert dns_span["op"] == "socket.dns" + assert dns_span["description"] == f"localhost:{PORT}" # noqa: E231 + assert dns_span["data"] == ApproxDict( + { + "host": "localhost", + "port": PORT, + } + ) + + +@pytest.mark.parametrize("span_streaming", [True, False]) +def test_span_origin(sentry_init, capture_events, capture_items, span_streaming): sentry_init( integrations=[SocketIntegration()], traces_sample_rate=1.0, + _experiments={"trace_lifecycle": "stream" if span_streaming else "static"}, ) - events = capture_events() - with start_transaction(name="foo"): - socket.create_connection(("localhost", PORT), 1, None) + if span_streaming: + items = capture_items("span") + with sentry_sdk.traces.start_span(name="foo"): + socket.create_connection(("localhost", PORT), 1, None) + sentry_sdk.flush() + + spans = [item.payload for item in items if item.type == "span"] + dns_span, connect_span, _root = spans + + assert connect_span["attributes"]["sentry.op"] == "socket.connection" + assert connect_span["attributes"]["sentry.origin"] == "auto.socket.socket" + + assert dns_span["attributes"]["sentry.op"] == "socket.dns" + assert dns_span["attributes"]["sentry.origin"] == "auto.socket.socket" + else: + events = capture_events() + + with start_transaction(name="foo"): + socket.create_connection(("localhost", PORT), 1, None) - (event,) = events + (event,) = events - assert event["contexts"]["trace"]["origin"] == "manual" + assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["op"] == "socket.connection" - assert event["spans"][0]["origin"] == "auto.socket.socket" + assert event["spans"][0]["op"] == "socket.connection" + assert event["spans"][0]["origin"] == "auto.socket.socket" - assert event["spans"][1]["op"] == "socket.dns" - assert event["spans"][1]["origin"] == "auto.socket.socket" + assert event["spans"][1]["op"] == "socket.dns" + assert event["spans"][1]["origin"] == "auto.socket.socket" From ff88f0443051df0033c29140f96bc77821bc3a33 Mon Sep 17 00:00:00 2001 From: Neel Shah Date: Tue, 19 May 2026 16:46:50 +0200 Subject: [PATCH 2/2] port int --- sentry_sdk/integrations/socket.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/sentry_sdk/integrations/socket.py b/sentry_sdk/integrations/socket.py index 2e4464730d..775170fb9f 100644 --- a/sentry_sdk/integrations/socket.py +++ b/sentry_sdk/integrations/socket.py @@ -118,14 +118,14 @@ def getaddrinfo( span.set_attribute( SPANDATA.SERVER_ADDRESS, host.decode(errors="replace") ) + if isinstance(port, int): span.set_attribute(SPANDATA.SERVER_PORT, port) - elif isinstance(port, str): - span.set_attribute(SPANDATA.SERVER_PORT, port) - elif isinstance(port, bytes): - span.set_attribute( - SPANDATA.SERVER_PORT, port.decode(errors="replace") - ) + elif port is not None: + try: + span.set_attribute(SPANDATA.SERVER_PORT, int(port)) + except (ValueError, TypeError): + pass return real_getaddrinfo(host, port, family, type, proto, flags) else: