Skip to content

Commit 967432b

Browse files
sl0thentr0pysl0thentr0py
committed
feat(socket): Support span streaming
1 parent 3142b16 commit 967432b

2 files changed

Lines changed: 196 additions & 81 deletions

File tree

‎sentry_sdk/integrations/socket.py‎

Lines changed: 68 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,9 @@
22

33
import sentry_sdk
44
from sentry_sdk._types import MYPY
5-
from sentry_sdk.consts import OP
5+
from sentry_sdk.consts import OP, SPANDATA
66
from sentry_sdk.integrations import Integration
7+
from sentry_sdk.tracing_utils import has_span_streaming_enabled
78

89
if MYPY:
910
from socket import AddressFamily, SocketKind
@@ -50,22 +51,39 @@ def create_connection(
5051
timeout: "Optional[float]" = socket._GLOBAL_DEFAULT_TIMEOUT, # type: ignore
5152
source_address: "Optional[Tuple[Union[bytearray, bytes, str], int]]" = None,
5253
) -> "socket.socket":
53-
integration = sentry_sdk.get_client().get_integration(SocketIntegration)
54+
client = sentry_sdk.get_client()
55+
integration = client.get_integration(SocketIntegration)
5456
if integration is None:
5557
return real_create_connection(address, timeout, source_address)
5658

57-
with sentry_sdk.start_span(
58-
op=OP.SOCKET_CONNECTION,
59-
name=_get_span_description(address[0], address[1]),
60-
origin=SocketIntegration.origin,
61-
) as span:
62-
span.set_data("address", address)
63-
span.set_data("timeout", timeout)
64-
span.set_data("source_address", source_address)
65-
66-
return real_create_connection(
67-
address=address, timeout=timeout, source_address=source_address
68-
)
59+
if has_span_streaming_enabled(client.options):
60+
with sentry_sdk.traces.start_span(
61+
name=_get_span_description(address[0], address[1]),
62+
attributes={
63+
"sentry.op": OP.SOCKET_CONNECTION,
64+
"sentry.origin": SocketIntegration.origin,
65+
},
66+
) as span:
67+
if address[0] is not None:
68+
span.set_attribute(SPANDATA.SERVER_ADDRESS, address[0])
69+
span.set_attribute(SPANDATA.SERVER_PORT, address[1])
70+
71+
return real_create_connection(
72+
address=address, timeout=timeout, source_address=source_address
73+
)
74+
else:
75+
with sentry_sdk.start_span(
76+
op=OP.SOCKET_CONNECTION,
77+
name=_get_span_description(address[0], address[1]),
78+
origin=SocketIntegration.origin,
79+
) as span:
80+
span.set_data("address", address)
81+
span.set_data("timeout", timeout)
82+
span.set_data("source_address", source_address)
83+
84+
return real_create_connection(
85+
address=address, timeout=timeout, source_address=source_address
86+
)
6987

7088
socket.create_connection = create_connection # type: ignore
7189

@@ -81,18 +99,44 @@ def getaddrinfo(
8199
proto: int = 0,
82100
flags: int = 0,
83101
) -> "List[Tuple[AddressFamily, SocketKind, int, str, Union[Tuple[str, int], Tuple[str, int, int, int], Tuple[int, bytes]]]]":
84-
integration = sentry_sdk.get_client().get_integration(SocketIntegration)
102+
client = sentry_sdk.get_client()
103+
integration = client.get_integration(SocketIntegration)
85104
if integration is None:
86105
return real_getaddrinfo(host, port, family, type, proto, flags)
87106

88-
with sentry_sdk.start_span(
89-
op=OP.SOCKET_DNS,
90-
name=_get_span_description(host, port),
91-
origin=SocketIntegration.origin,
92-
) as span:
93-
span.set_data("host", host)
94-
span.set_data("port", port)
95-
96-
return real_getaddrinfo(host, port, family, type, proto, flags)
107+
if has_span_streaming_enabled(client.options):
108+
with sentry_sdk.traces.start_span(
109+
name=_get_span_description(host, port),
110+
attributes={
111+
"sentry.op": OP.SOCKET_DNS,
112+
"sentry.origin": SocketIntegration.origin,
113+
},
114+
) as span:
115+
if isinstance(host, str):
116+
span.set_attribute(SPANDATA.SERVER_ADDRESS, host)
117+
elif isinstance(host, bytes):
118+
span.set_attribute(
119+
SPANDATA.SERVER_ADDRESS, host.decode(errors="replace")
120+
)
121+
if isinstance(port, int):
122+
span.set_attribute(SPANDATA.SERVER_PORT, port)
123+
elif isinstance(port, str):
124+
span.set_attribute(SPANDATA.SERVER_PORT, port)
125+
elif isinstance(port, bytes):
126+
span.set_attribute(
127+
SPANDATA.SERVER_PORT, port.decode(errors="replace")
128+
)
129+
130+
return real_getaddrinfo(host, port, family, type, proto, flags)
131+
else:
132+
with sentry_sdk.start_span(
133+
op=OP.SOCKET_DNS,
134+
name=_get_span_description(host, port),
135+
origin=SocketIntegration.origin,
136+
) as span:
137+
span.set_data("host", host)
138+
span.set_data("port", port)
139+
140+
return real_getaddrinfo(host, port, family, type, proto, flags)
97141

98142
socket.getaddrinfo = getaddrinfo

‎tests/integrations/socket/test_socket.py‎

Lines changed: 128 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -1,81 +1,152 @@
11
import socket
22

3+
import pytest
4+
5+
import sentry_sdk
36
from sentry_sdk import start_transaction
47
from sentry_sdk.integrations.socket import SocketIntegration
58
from tests.conftest import ApproxDict, create_mock_http_server
69

710
PORT = create_mock_http_server()
811

912

10-
def test_getaddrinfo_trace(sentry_init, capture_events):
11-
sentry_init(integrations=[SocketIntegration()], traces_sample_rate=1.0)
12-
events = capture_events()
13-
14-
with start_transaction():
15-
socket.getaddrinfo("localhost", PORT)
16-
17-
(event,) = events
18-
(span,) = event["spans"]
19-
20-
assert span["op"] == "socket.dns"
21-
assert span["description"] == f"localhost:{PORT}" # noqa: E231
22-
assert span["data"] == ApproxDict(
23-
{
24-
"host": "localhost",
25-
"port": PORT,
26-
}
13+
@pytest.mark.parametrize("span_streaming", [True, False])
14+
def test_getaddrinfo_trace(sentry_init, capture_events, capture_items, span_streaming):
15+
sentry_init(
16+
integrations=[SocketIntegration()],
17+
traces_sample_rate=1.0,
18+
_experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
2719
)
2820

29-
30-
def test_create_connection_trace(sentry_init, capture_events):
21+
if span_streaming:
22+
items = capture_items("span")
23+
with sentry_sdk.traces.start_span(name="root"):
24+
socket.getaddrinfo("localhost", PORT)
25+
sentry_sdk.flush()
26+
27+
spans = [item.payload for item in items if item.type == "span"]
28+
dns_span, _root = spans
29+
30+
assert dns_span["attributes"]["sentry.op"] == "socket.dns"
31+
assert dns_span["attributes"]["sentry.origin"] == "auto.socket.socket"
32+
assert dns_span["name"] == f"localhost:{PORT}" # noqa: E231
33+
assert dns_span["attributes"]["server.address"] == "localhost"
34+
assert dns_span["attributes"]["server.port"] == PORT
35+
else:
36+
events = capture_events()
37+
38+
with start_transaction():
39+
socket.getaddrinfo("localhost", PORT)
40+
41+
(event,) = events
42+
(span,) = event["spans"]
43+
44+
assert span["op"] == "socket.dns"
45+
assert span["description"] == f"localhost:{PORT}" # noqa: E231
46+
assert span["data"] == ApproxDict(
47+
{
48+
"host": "localhost",
49+
"port": PORT,
50+
}
51+
)
52+
53+
54+
@pytest.mark.parametrize("span_streaming", [True, False])
55+
def test_create_connection_trace(
56+
sentry_init, capture_events, capture_items, span_streaming
57+
):
3158
timeout = 10
3259

33-
sentry_init(integrations=[SocketIntegration()], traces_sample_rate=1.0)
34-
events = capture_events()
35-
36-
with start_transaction():
37-
socket.create_connection(("localhost", PORT), timeout, None)
38-
39-
(event,) = events
40-
(connect_span, dns_span) = event["spans"]
41-
# as getaddrinfo gets called in create_connection it should also contain a dns span
42-
43-
assert connect_span["op"] == "socket.connection"
44-
assert connect_span["description"] == f"localhost:{PORT}" # noqa: E231
45-
assert connect_span["data"] == ApproxDict(
46-
{
47-
"address": ["localhost", PORT],
48-
"timeout": timeout,
49-
"source_address": None,
50-
}
51-
)
52-
53-
assert dns_span["op"] == "socket.dns"
54-
assert dns_span["description"] == f"localhost:{PORT}" # noqa: E231
55-
assert dns_span["data"] == ApproxDict(
56-
{
57-
"host": "localhost",
58-
"port": PORT,
59-
}
60+
sentry_init(
61+
integrations=[SocketIntegration()],
62+
traces_sample_rate=1.0,
63+
_experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
6064
)
6165

62-
63-
def test_span_origin(sentry_init, capture_events):
66+
if span_streaming:
67+
items = capture_items("span")
68+
with sentry_sdk.traces.start_span(name="root"):
69+
socket.create_connection(("localhost", PORT), timeout, None)
70+
sentry_sdk.flush()
71+
72+
spans = [item.payload for item in items if item.type == "span"]
73+
# as getaddrinfo gets called in create_connection it should also contain a dns span
74+
# spans finish in order: dns (inner) ends first, connect ends, then root
75+
dns_span, connect_span, _root = spans
76+
77+
assert connect_span["attributes"]["sentry.op"] == "socket.connection"
78+
assert connect_span["name"] == f"localhost:{PORT}" # noqa: E231
79+
assert connect_span["attributes"]["server.address"] == "localhost"
80+
assert connect_span["attributes"]["server.port"] == PORT
81+
82+
assert dns_span["attributes"]["sentry.op"] == "socket.dns"
83+
assert dns_span["name"] == f"localhost:{PORT}" # noqa: E231
84+
assert dns_span["attributes"]["server.address"] == "localhost"
85+
assert dns_span["attributes"]["server.port"] == PORT
86+
else:
87+
events = capture_events()
88+
89+
with start_transaction():
90+
socket.create_connection(("localhost", PORT), timeout, None)
91+
92+
(event,) = events
93+
(connect_span, dns_span) = event["spans"]
94+
# as getaddrinfo gets called in create_connection it should also contain a dns span
95+
96+
assert connect_span["op"] == "socket.connection"
97+
assert connect_span["description"] == f"localhost:{PORT}" # noqa: E231
98+
assert connect_span["data"] == ApproxDict(
99+
{
100+
"address": ["localhost", PORT],
101+
"timeout": timeout,
102+
"source_address": None,
103+
}
104+
)
105+
106+
assert dns_span["op"] == "socket.dns"
107+
assert dns_span["description"] == f"localhost:{PORT}" # noqa: E231
108+
assert dns_span["data"] == ApproxDict(
109+
{
110+
"host": "localhost",
111+
"port": PORT,
112+
}
113+
)
114+
115+
116+
@pytest.mark.parametrize("span_streaming", [True, False])
117+
def test_span_origin(sentry_init, capture_events, capture_items, span_streaming):
64118
sentry_init(
65119
integrations=[SocketIntegration()],
66120
traces_sample_rate=1.0,
121+
_experiments={"trace_lifecycle": "stream" if span_streaming else "static"},
67122
)
68-
events = capture_events()
69123

70-
with start_transaction(name="foo"):
71-
socket.create_connection(("localhost", PORT), 1, None)
124+
if span_streaming:
125+
items = capture_items("span")
126+
with sentry_sdk.traces.start_span(name="foo"):
127+
socket.create_connection(("localhost", PORT), 1, None)
128+
sentry_sdk.flush()
129+
130+
spans = [item.payload for item in items if item.type == "span"]
131+
dns_span, connect_span, _root = spans
132+
133+
assert connect_span["attributes"]["sentry.op"] == "socket.connection"
134+
assert connect_span["attributes"]["sentry.origin"] == "auto.socket.socket"
135+
136+
assert dns_span["attributes"]["sentry.op"] == "socket.dns"
137+
assert dns_span["attributes"]["sentry.origin"] == "auto.socket.socket"
138+
else:
139+
events = capture_events()
140+
141+
with start_transaction(name="foo"):
142+
socket.create_connection(("localhost", PORT), 1, None)
72143

73-
(event,) = events
144+
(event,) = events
74145

75-
assert event["contexts"]["trace"]["origin"] == "manual"
146+
assert event["contexts"]["trace"]["origin"] == "manual"
76147

77-
assert event["spans"][0]["op"] == "socket.connection"
78-
assert event["spans"][0]["origin"] == "auto.socket.socket"
148+
assert event["spans"][0]["op"] == "socket.connection"
149+
assert event["spans"][0]["origin"] == "auto.socket.socket"
79150

80-
assert event["spans"][1]["op"] == "socket.dns"
81-
assert event["spans"][1]["origin"] == "auto.socket.socket"
151+
assert event["spans"][1]["op"] == "socket.dns"
152+
assert event["spans"][1]["origin"] == "auto.socket.socket"

0 commit comments

Comments
 (0)