fix: address PR review findings — forward skip_context_check in

recursive fallback, reduce retry backoff, clean up batch state

  - Forward skip_context_check parameter in the 401/403 ephemeral
    fallback recursive call to prevent silent early return when
    is_tracing_enabled_in_context() is False
  - Reduce retry backoff from 1s to 200ms to minimize lock hold time
    on the non-first-time path (worst case 400ms vs 2s)
  - Add batch state cleanup after _finalize_backend_batch in the
    first-time handler, mirroring finalize_batch: reset current_batch,
    event_buffer, trace_batch_id, is_current_batch_ephemeral,
    batch_owner_type, batch_owner_id, and call _cleanup_batch_data()
This commit is contained in:
Tiago Freire
2026-03-18 21:29:35 -03:00
parent 7f2eda2ac6
commit 4859dc66a5
3 changed files with 16 additions and 4 deletions

View File

@@ -122,6 +122,15 @@ class FirstTimeTraceHandler:
self.batch_manager._finalize_backend_batch(events_count)
self.ephemeral_url = self.batch_manager.ephemeral_trace_url
# Clean up batch state (mirrors finalize_batch cleanup)
self.batch_manager.batch_owner_type = None
self.batch_manager.batch_owner_id = None
self.batch_manager.current_batch = None
self.batch_manager.event_buffer.clear()
self.batch_manager.trace_batch_id = None
self.batch_manager.is_current_batch_ephemeral = False
self.batch_manager._cleanup_batch_data()
if not self.ephemeral_url:
self._show_local_trace_message()

View File

@@ -162,14 +162,14 @@ class TraceBatchManager:
f"Trace batch init attempt {attempt + 1} failed "
f"(status={response.status_code if response else 'None'}), retrying..."
)
time.sleep(1)
time.sleep(0.2)
except Exception as e:
last_exception = e
if attempt < max_retries:
logger.debug(
f"Trace batch init attempt {attempt + 1} raised {type(e).__name__}, retrying..."
)
time.sleep(1)
time.sleep(0.2)
if last_exception and response is None:
logger.warning(
@@ -192,7 +192,10 @@ class TraceBatchManager:
)
self.is_current_batch_ephemeral = True
return self._initialize_backend_batch(
user_context, execution_metadata, use_ephemeral=True
user_context,
execution_metadata,
use_ephemeral=True,
skip_context_check=skip_context_check,
)
if response.status_code in [201, 200]:

View File

@@ -1099,7 +1099,7 @@ class TestInitializeBackendBatchRetry:
assert bm.trace_batch_id == server_id
assert mock_init.call_count == 2
mock_sleep.assert_called_once_with(1)
mock_sleep.assert_called_once_with(0.2)
def test_retries_on_5xx_then_succeeds(self):
"""Retries on 500 server error, succeeds on second attempt."""