Skip to content

Commit 8e0c49e

Browse files
jy-tanjy-tan
committed
Fix calling libraries and add notes
1 parent ce3b3c9 commit 8e0c49e

6 files changed

Lines changed: 335 additions & 170 deletions

File tree

‎drift/instrumentation/aiohttp/instrumentation.py‎

Lines changed: 33 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ def __init__(self, message: str, method: str, url: str):
4444
SpanStatus,
4545
StatusCode,
4646
TuskDriftMode,
47+
calling_library_context,
4748
)
4849
from ..base import InstrumentationBase
4950
from ..http import HttpSpanData, HttpTransformEngine
@@ -164,32 +165,41 @@ async def patched_request(
164165
if sdk.mode == TuskDriftMode.DISABLED:
165166
return await original_request(client_self, method, str_or_url, **kwargs)
166167

167-
async def original_call():
168-
return await original_request(client_self, method, str_or_url, **kwargs)
168+
# Set calling_library_context to suppress socket instrumentation warnings
169+
# for internal socket calls (e.g., aiohappyeyeballs connection management)
170+
context_token = calling_library_context.set("aiohttp")
171+
try:
169172

170-
# REPLAY mode
171-
if sdk.mode == TuskDriftMode.REPLAY:
172-
return await handle_replay_mode(
173-
replay_mode_handler=lambda: instrumentation_self._handle_replay_request(
174-
sdk, module, method, url_str, **kwargs
173+
async def original_call():
174+
return await original_request(client_self, method, str_or_url, **kwargs)
175+
176+
# REPLAY mode
177+
if sdk.mode == TuskDriftMode.REPLAY:
178+
return await handle_replay_mode(
179+
replay_mode_handler=lambda: instrumentation_self._handle_replay_request(
180+
sdk, module, method, url_str, **kwargs
181+
),
182+
no_op_request_handler=lambda: instrumentation_self._get_default_response(
183+
module, method, url_str
184+
),
185+
is_server_request=False,
186+
)
187+
188+
# RECORD mode
189+
return await handle_record_mode(
190+
original_function_call=original_call,
191+
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_request(
192+
client_self,
193+
method,
194+
str_or_url,
195+
is_pre_app_start,
196+
original_request,
197+
**kwargs,
175198
),
176-
no_op_request_handler=lambda: instrumentation_self._get_default_response(module, method, url_str),
177-
is_server_request=False,
199+
span_kind=OTelSpanKind.CLIENT,
178200
)
179-
180-
# RECORD mode
181-
return await handle_record_mode(
182-
original_function_call=original_call,
183-
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_request(
184-
client_self,
185-
method,
186-
str_or_url,
187-
is_pre_app_start,
188-
original_request,
189-
**kwargs,
190-
),
191-
span_kind=OTelSpanKind.CLIENT,
192-
)
201+
finally:
202+
calling_library_context.reset(context_token)
193203

194204
# Apply patch
195205
module.ClientSession._request = patched_request

‎drift/instrumentation/aiohttp/notes.md‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# aiohttp Instrumentation Notes
2+
3+
## 1. Socket Instrumentation Warnings from `aiohappyeyeballs`
4+
5+
### Background
6+
7+
aiohttp uses `aiohappyeyeballs` internally for connection management. This library implements the "Happy Eyeballs" algorithm (RFC 8305) for fast and reliable connection establishment by racing IPv4 and IPv6 connections.
8+
9+
When making HTTP requests, aiohttp's connection flow is:
10+
11+
1. Application calls `session.get()` or similar method
12+
2. This calls `ClientSession._request()` (which we patch)
13+
3. aiohttp internally uses `aiohappyeyeballs` to establish the TCP connection
14+
4. `aiohappyeyeballs` uses low-level asyncio APIs like `loop.sock_connect()`
15+
5. These schedule callbacks via the asyncio event loop
16+
17+
### Known Limitation
18+
19+
**Problem:** In REPLAY mode, socket instrumentation warnings are still logged for aiohttp requests, even though the HTTP layer is properly instrumented.
20+
21+
**Root Cause:** Python's `contextvars` propagate through `async/await` chains, but NOT through all asyncio internal callback mechanisms. When `aiohappyeyeballs` schedules socket operations via `loop.sock_connect()`, these run as event loop callbacks that don't inherit the `calling_library_context` we set in the patched `_request()` method.
22+
23+
The stack trace shows the context boundary:
24+
25+
```text
26+
File "asyncio/events.py", line 88, in _run
27+
self._context.run(self._callback, *self._args)
28+
File "aiohappyeyeballs/_staggered.py", line 132, in run_one_coro
29+
result = await coro_fn()
30+
File "aiohappyeyeballs/impl.py", line 208, in _connect_sock
31+
await loop.sock_connect(sock, address)
32+
```
33+
34+
The `_context.run()` call creates a new context for the callback, which doesn't include our `calling_library_context`.
35+
36+
### Why This Happens Only with aiohttp
37+
38+
Other HTTP client libraries like httpx don't have this issue because:
39+
40+
| Library | Connection Layer | Context Propagation |
41+
|---------|-----------------|---------------------|
42+
| **httpx** | Uses `httpcore` which maintains context through async/await | ✅ Works correctly |
43+
| **requests** | Synchronous, no callbacks | ✅ Works correctly |
44+
| **urllib3** | Synchronous, no callbacks | ✅ Works correctly |
45+
| **aiohttp** | Uses `aiohappyeyeballs` with asyncio callbacks | ⚠️ Context lost in callbacks |
46+
47+
### Impact
48+
49+
| Aspect | Impact |
50+
|--------|--------|
51+
| **Functionality** | ✅ No impact - instrumentation works correctly |
52+
| **E2E Tests** | ✅ All tests pass |
53+
| **Logs** | ⚠️ Spurious "Unpatched dependency alert" warnings in REPLAY mode |
54+
| **CLI Alerts** | ⚠️ False positive alerts sent to CLI |
55+
56+
The warnings are cosmetic - they don't indicate actual unpatched dependencies since we successfully intercept at the higher `ClientSession._request()` level.
57+
58+
### Mitigation
59+
60+
The e2e test framework already handles this by skipping socket instrumentation warning checks for aiohttp:
61+
62+
```python
63+
# Note: The socket instrumentation will detect and warn about these calls, but this
64+
# is expected behavior. The warnings indicate that there are low-level socket calls
65+
# happening that are not directly instrumented by our aiohttp instrumentation.
66+
# This is normal for aiohappyeyeballs and doesn't indicate an instrumentation bug
67+
# since we successfully intercept at the higher ClientSession._request() level.
68+
```
69+
70+
### Code Location
71+
72+
The `calling_library_context` is set in `_patch_client_session()` in [`instrumentation.py`](./instrumentation.py):
73+
74+
```python
75+
async def patched_request(client_self, method: str, str_or_url, **kwargs):
76+
# ...
77+
# Set calling_library_context to suppress socket instrumentation warnings
78+
# for internal socket calls (e.g., aiohappyeyeballs connection management)
79+
context_token = calling_library_context.set("aiohttp")
80+
try:
81+
# ... handle REPLAY or RECORD mode ...
82+
finally:
83+
calling_library_context.reset(context_token)
84+
```
85+
86+
### Future Considerations
87+
88+
**Potential approaches to eliminate the warnings:**
89+
90+
1. **Stack trace inspection**: Modify socket instrumentation to check the stack trace for known libraries like `aiohappyeyeballs` and suppress warnings. However, this adds overhead to every socket call and is fragile if library internals change.
91+
92+
2. **asyncio context propagation**: Python's `contextvars` module has `copy_context()` which could theoretically be used, but we don't control how `aiohappyeyeballs` or asyncio schedules callbacks.
93+
94+
3. **Library-specific skip list**: Add a configuration to socket instrumentation to skip warnings for known libraries. This is the most practical approach but requires maintenance as new libraries are added.
95+
96+
**Current approach:** Accept the warnings for aiohttp as expected behavior. The e2e test framework already handles this correctly, and the warnings don't affect functionality. Users who see these warnings in their logs can be informed they're expected for aiohttp.
97+
98+
### Test Coverage
99+
100+
The e2e tests verify that:
101+
102+
- All aiohttp request types work correctly in RECORD and REPLAY modes
103+
- The socket instrumentation warning check is skipped for aiohttp
104+
- Traces are captured correctly despite the warnings

‎drift/instrumentation/httpx/instrumentation.py‎

Lines changed: 69 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ def __init__(self, message: str, method: str, url: str):
4545
SpanStatus,
4646
StatusCode,
4747
TuskDriftMode,
48+
calling_library_context,
4849
)
4950
from ..base import InstrumentationBase
5051
from ..http import HttpSpanData, HttpTransformEngine
@@ -285,33 +286,42 @@ def patched_send(
285286
if sdk.mode == TuskDriftMode.DISABLED:
286287
return original_send(client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects)
287288

288-
def original_call():
289-
return original_send(client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects)
289+
# Set calling_library_context to suppress socket instrumentation warnings
290+
# for internal socket calls made by httpx or its dependencies
291+
context_token = calling_library_context.set("httpx")
292+
try:
293+
294+
def original_call():
295+
return original_send(
296+
client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects
297+
)
298+
299+
# REPLAY mode: Use handle_replay_mode for proper background request handling
300+
if sdk.mode == TuskDriftMode.REPLAY:
301+
return handle_replay_mode(
302+
replay_mode_handler=lambda: instrumentation_self._handle_replay_send_sync(
303+
sdk, module, request, auth=auth, client=client_self
304+
),
305+
no_op_request_handler=lambda: instrumentation_self._get_default_response(module, url_str),
306+
is_server_request=False,
307+
)
290308

291-
# REPLAY mode: Use handle_replay_mode for proper background request handling
292-
if sdk.mode == TuskDriftMode.REPLAY:
293-
return handle_replay_mode(
294-
replay_mode_handler=lambda: instrumentation_self._handle_replay_send_sync(
295-
sdk, module, request, auth=auth, client=client_self
309+
# RECORD mode: Use handle_record_mode for proper is_pre_app_start handling
310+
return handle_record_mode(
311+
original_function_call=original_call,
312+
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send_sync(
313+
client_self,
314+
request,
315+
stream,
316+
is_pre_app_start,
317+
original_send,
318+
auth=auth,
319+
follow_redirects=follow_redirects,
296320
),
297-
no_op_request_handler=lambda: instrumentation_self._get_default_response(module, url_str),
298-
is_server_request=False,
321+
span_kind=OTelSpanKind.CLIENT,
299322
)
300-
301-
# RECORD mode: Use handle_record_mode for proper is_pre_app_start handling
302-
return handle_record_mode(
303-
original_function_call=original_call,
304-
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send_sync(
305-
client_self,
306-
request,
307-
stream,
308-
is_pre_app_start,
309-
original_send,
310-
auth=auth,
311-
follow_redirects=follow_redirects,
312-
),
313-
span_kind=OTelSpanKind.CLIENT,
314-
)
323+
finally:
324+
calling_library_context.reset(context_token)
315325

316326
# Apply patch
317327
module.Client.send = patched_send
@@ -467,37 +477,44 @@ async def patched_send(
467477
client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects
468478
)
469479

470-
async def original_call():
471-
return await original_send(
472-
client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects
473-
)
480+
# Set calling_library_context to suppress socket instrumentation warnings
481+
# for internal socket calls made by httpx or its dependencies
482+
context_token = calling_library_context.set("httpx")
483+
try:
474484

475-
# REPLAY mode: Use handle_replay_mode for proper background request handling
476-
# handle_replay_mode returns coroutine which we await
477-
if sdk.mode == TuskDriftMode.REPLAY:
478-
return await handle_replay_mode(
479-
replay_mode_handler=lambda: instrumentation_self._handle_replay_send_async(
480-
sdk, module, request, auth=auth, client=client_self
485+
async def original_call():
486+
return await original_send(
487+
client_self, request, stream=stream, auth=auth, follow_redirects=follow_redirects
488+
)
489+
490+
# REPLAY mode: Use handle_replay_mode for proper background request handling
491+
# handle_replay_mode returns coroutine which we await
492+
if sdk.mode == TuskDriftMode.REPLAY:
493+
return await handle_replay_mode(
494+
replay_mode_handler=lambda: instrumentation_self._handle_replay_send_async(
495+
sdk, module, request, auth=auth, client=client_self
496+
),
497+
no_op_request_handler=lambda: instrumentation_self._get_default_response(module, url_str),
498+
is_server_request=False,
499+
)
500+
501+
# RECORD mode: Use handle_record_mode for proper is_pre_app_start handling
502+
# handle_record_mode returns coroutine which we await
503+
return await handle_record_mode(
504+
original_function_call=original_call,
505+
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send_async(
506+
client_self,
507+
request,
508+
stream,
509+
is_pre_app_start,
510+
original_send,
511+
auth=auth,
512+
follow_redirects=follow_redirects,
481513
),
482-
no_op_request_handler=lambda: instrumentation_self._get_default_response(module, url_str),
483-
is_server_request=False,
514+
span_kind=OTelSpanKind.CLIENT,
484515
)
485-
486-
# RECORD mode: Use handle_record_mode for proper is_pre_app_start handling
487-
# handle_record_mode returns coroutine which we await
488-
return await handle_record_mode(
489-
original_function_call=original_call,
490-
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send_async(
491-
client_self,
492-
request,
493-
stream,
494-
is_pre_app_start,
495-
original_send,
496-
auth=auth,
497-
follow_redirects=follow_redirects,
498-
),
499-
span_kind=OTelSpanKind.CLIENT,
500-
)
516+
finally:
517+
calling_library_context.reset(context_token)
501518

502519
# Apply patch
503520
module.AsyncClient.send = patched_send

‎drift/instrumentation/requests/instrumentation.py‎

Lines changed: 27 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ def __init__(self, message: str, method: str, url: str):
4444
SpanStatus,
4545
StatusCode,
4646
TuskDriftMode,
47+
calling_library_context,
4748
)
4849
from ..base import InstrumentationBase
4950
from ..http import HttpSpanData, HttpTransformEngine
@@ -122,28 +123,34 @@ def patched_send(session_self, request, **kwargs):
122123
if sdk.mode == TuskDriftMode.DISABLED:
123124
return original_send(session_self, request, **kwargs)
124125

125-
# Extract URL for default response handler
126-
url = request.url
127-
128-
def original_call():
129-
return original_send(session_self, request, **kwargs)
126+
# Set calling_library_context to suppress socket instrumentation warnings
127+
# for internal socket calls made by requests or its dependencies (urllib3)
128+
context_token = calling_library_context.set("requests")
129+
try:
130+
# Extract URL for default response handler
131+
url = request.url
132+
133+
def original_call():
134+
return original_send(session_self, request, **kwargs)
135+
136+
# REPLAY mode: Use handle_replay_mode for proper background request handling
137+
if sdk.mode == TuskDriftMode.REPLAY:
138+
return handle_replay_mode(
139+
replay_mode_handler=lambda: instrumentation_self._handle_replay_send(sdk, request, **kwargs),
140+
no_op_request_handler=lambda: instrumentation_self._get_default_response(url),
141+
is_server_request=False,
142+
)
130143

131-
# REPLAY mode: Use handle_replay_mode for proper background request handling
132-
if sdk.mode == TuskDriftMode.REPLAY:
133-
return handle_replay_mode(
134-
replay_mode_handler=lambda: instrumentation_self._handle_replay_send(sdk, request, **kwargs),
135-
no_op_request_handler=lambda: instrumentation_self._get_default_response(url),
136-
is_server_request=False,
144+
# RECORD mode: Use handle_record_mode for proper is_pre_app_start handling
145+
return handle_record_mode(
146+
original_function_call=original_call,
147+
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send(
148+
session_self, request, is_pre_app_start, original_send, **kwargs
149+
),
150+
span_kind=OTelSpanKind.CLIENT,
137151
)
138-
139-
# RECORD mode: Use handle_record_mode for proper is_pre_app_start handling
140-
return handle_record_mode(
141-
original_function_call=original_call,
142-
record_mode_handler=lambda is_pre_app_start: instrumentation_self._handle_record_send(
143-
session_self, request, is_pre_app_start, original_send, **kwargs
144-
),
145-
span_kind=OTelSpanKind.CLIENT,
146-
)
152+
finally:
153+
calling_library_context.reset(context_token)
147154

148155
# Apply patch
149156
module.Session.send = patched_send

0 commit comments

Comments
 (0)