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