@@ -338,49 +338,64 @@ async def __call__(self, scope, receive, send):
338
338
if callable (self .server_request_hook ):
339
339
self .server_request_hook (span , scope )
340
340
341
+ otel_receive = self ._get_otel_receive (
342
+ span_name , scope , receive
343
+ )
344
+
341
345
server_span_context = trace .context_api .get_current ()
346
+ otel_send = self ._get_otel_send (
347
+ span , span_name , server_span_context , scope , send
348
+ )
342
349
343
- @wraps (receive )
344
- async def wrapped_receive ():
345
- with self .tracer .start_as_current_span (
346
- " " .join ((span_name , scope ["type" ], "receive" ))
347
- ) as receive_span :
348
- if callable (self .client_request_hook ):
349
- self .client_request_hook (receive_span , scope )
350
- message = await receive ()
351
- if receive_span .is_recording ():
352
- if message ["type" ] == "websocket.receive" :
353
- set_status_code (receive_span , 200 )
354
- receive_span .set_attribute ("type" , message ["type" ])
355
- return message
356
-
357
- @wraps (send )
358
- async def wrapped_send (message ):
359
- with self .tracer .start_as_current_span (
360
- " " .join ((span_name , scope ["type" ], "send" ))
361
- ) as send_span :
362
- if callable (self .client_response_hook ):
363
- self .client_response_hook (send_span , message )
364
- if send_span .is_recording ():
365
- if message ["type" ] == "http.response.start" :
366
- status_code = message ["status" ]
367
- set_status_code (span , status_code )
368
- set_status_code (send_span , status_code )
369
- elif message ["type" ] == "websocket.send" :
370
- set_status_code (span , 200 )
371
- set_status_code (send_span , 200 )
372
- send_span .set_attribute ("type" , message ["type" ])
373
-
374
- propagator = get_global_response_propagator ()
375
- if propagator :
376
- propagator .inject (
377
- message ,
378
- context = server_span_context ,
379
- setter = asgi_setter ,
380
- )
381
-
382
- await send (message )
383
-
384
- await self .app (scope , wrapped_receive , wrapped_send )
350
+ await self .app (scope , otel_receive , otel_send )
385
351
finally :
386
352
context .detach (token )
353
+
354
+ def _get_otel_receive (self , span_name , scope , receive ):
355
+ @wraps (receive )
356
+ async def wrapped_receive ():
357
+ with self .tracer .start_as_current_span (
358
+ " " .join ((span_name , scope ["type" ], "receive" ))
359
+ ) as receive_span :
360
+ if callable (self .client_request_hook ):
361
+ self .client_request_hook (receive_span , scope )
362
+ message = await receive ()
363
+ if receive_span .is_recording ():
364
+ if message ["type" ] == "websocket.receive" :
365
+ set_status_code (receive_span , 200 )
366
+ receive_span .set_attribute ("type" , message ["type" ])
367
+ return message
368
+
369
+ return wrapped_receive
370
+
371
+ async def _get_otel_send (
372
+ self , span , span_name , server_span_context , scope , send
373
+ ):
374
+ @wraps (send )
375
+ async def wrapped_send (message ):
376
+ with self .tracer .start_as_current_span (
377
+ " " .join ((span_name , scope ["type" ], "send" ))
378
+ ) as send_span :
379
+ if callable (self .client_response_hook ):
380
+ self .client_response_hook (send_span , message )
381
+ if send_span .is_recording ():
382
+ if message ["type" ] == "http.response.start" :
383
+ status_code = message ["status" ]
384
+ set_status_code (span , status_code )
385
+ set_status_code (send_span , status_code )
386
+ elif message ["type" ] == "websocket.send" :
387
+ set_status_code (span , 200 )
388
+ set_status_code (send_span , 200 )
389
+ send_span .set_attribute ("type" , message ["type" ])
390
+
391
+ propagator = get_global_response_propagator ()
392
+ if propagator :
393
+ propagator .inject (
394
+ message ,
395
+ context = server_span_context ,
396
+ setter = asgi_setter ,
397
+ )
398
+
399
+ await send (message )
400
+
401
+ return wrapped_send
0 commit comments