fix: preserve cancellation evidence across Python 3.10 task boundaries

Read attached diagnostics through CancelledError context wrappers in iteration handling, checkpoints and partial reports. Retain cancelled task semantics and test real Task/wait_for boundaries on supported Python versions.

Author: O2eg

src/pg_perf_bench/benchmark.py

@@ -46,6 +46,7 @@
     CommandExecutionError,
     CommandFailure,
     ConfigurationError,
+    exception_evidence,
 )
 from pg_perf_bench.executors.process import ProcessResult
 from pg_perf_bench.init_policy import (
@@ -547,8 +548,9 @@ async def finish_sampling():
         except BaseException as exc:
             result['status'] = 'cancelled' if isinstance(exc, asyncio.CancelledError) else 'failed'
             result['error'] = str(exc) or type(exc).__name__
-            if hasattr(exc, 'failure') and 'workload' not in result:
-                result['workload'] = exc.failure.as_dict()
+            failure = exception_evidence(exc, 'failure')
+            if failure is not None and 'workload' not in result:
+                result['workload'] = failure.as_dict()
             # Failed stdout can contain apparently valid TPS. Keep it only as raw
             # evidence, never as a point in the successful benchmark series.
             exc.benchmark_run = redact_mapping(result, secrets=secrets)
@@ -704,10 +706,11 @@ def write_checkpoint(workload_conf, db_conf, results, status, error=None):
         }
         if error is not None:
             payload['error'] = str(error) or type(error).__name__
-            payload['failed_iteration'] = getattr(error, 'benchmark_run', None)
-            if hasattr(error, 'failure'):
-                payload['failed_command'] = error.failure.as_dict()
-                payload['statement_timeout_messages'] = error.failure.stderr.count(
+            payload['failed_iteration'] = exception_evidence(error, 'benchmark_run')
+            failure = exception_evidence(error, 'failure')
+            if failure is not None:
+                payload['failed_command'] = failure.as_dict()
+                payload['statement_timeout_messages'] = failure.stderr.count(
                     'canceling statement due to statement timeout'
                 )
         payload = redact_mapping(payload, secrets=(db_conf.get('password'),))
@@ -901,9 +904,10 @@ async def _run_benchmark_iterations(
                     result['existing_dataset_validation'] = validation
             except BaseException as exc:
                 primary_error = exc
-                evidence = getattr(exc, 'benchmark_run', {})
-                if hasattr(exc, 'failure') and 'workload' not in evidence:
-                    evidence['failed_command'] = exc.failure.as_dict()
+                evidence = exception_evidence(exc, 'benchmark_run', {})
+                failure = exception_evidence(exc, 'failure')
+                if failure is not None and 'workload' not in evidence:
+                    evidence['failed_command'] = failure.as_dict()
                 evidence.update(
                     iteration=iteration,
                     init_policy=init_policy,
@@ -1023,8 +1027,8 @@ def partial_report(report, report_data, error, *, started_at, started_clock, sec
         """Render existing evidence without reconnecting to a failed database."""
         safe = redact_mapping(
             {
-                'runs': getattr(error, 'completed_runs', []),
-                'failed': getattr(error, 'benchmark_run', None),
+                'runs': exception_evidence(error, 'completed_runs', []),
+                'failed': exception_evidence(error, 'benchmark_run'),
                 'message': str(error) or type(error).__name__,
             },
             secrets=secrets,

src/pg_perf_bench/errors.py

@@ -2,10 +2,34 @@
 
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass
 from typing import Any
 
 
+def exception_evidence(error: BaseException, name: str, default: Any = None) -> Any:
+    """Read attached evidence through Task cancellation wrappers on Python 3.10.
+
+    Older asyncio Tasks create a new CancelledError and keep the original in
+    __context__. Do not follow unrelated exception contexts: they may describe
+    a different operation. The closest explicit value, including None, wins.
+    """
+    pending = [error]
+    seen = set()
+    while pending:
+        current = pending.pop()
+        if id(current) in seen:
+            continue
+        seen.add(id(current))
+        if hasattr(current, name):
+            return getattr(current, name)
+        if isinstance(current, asyncio.CancelledError):
+            for nested in (current.__context__, current.__cause__):
+                if isinstance(nested, asyncio.CancelledError):
+                    pending.append(nested)
+    return default
+
+
 class PgPerfBenchError(Exception):
     """Base class for user-facing failures."""
 

tests/test_cancellation_evidence.py

@@ -0,0 +1,92 @@
+"""Cancelled task diagnostics must survive the Python 3.10 Task boundary."""
+
+import asyncio
+import json
+import time
+from copy import deepcopy
+
+import pytest
+
+from pg_perf_bench.benchmark import BenchmarkRunner
+from pg_perf_bench.errors import CommandFailure, exception_evidence
+from tests.test_interrupted_benchmark import partial_report
+
+
+@pytest.mark.parametrize('wait_for', [False, True])
+def test_cancelled_task_retains_checkpoint_and_partial_report_evidence(tmp_path, wait_for):
+    async def scenario():
+        reference = partial_report(cancelled=True)
+        started = asyncio.Event()
+        command = CommandFailure('pgbench', -15, 'partial secret-value', 'cancelled', 1)
+
+        async def child():
+            try:
+                started.set()
+                await asyncio.Event().wait()
+            except asyncio.CancelledError as exc:
+                exc.failure = command
+                exc.benchmark_run = reference['failed_iteration']
+                exc.completed_runs = reference['benchmark_runs']
+                raise
+
+        task = asyncio.create_task(child())
+        await started.wait()
+        task.cancel()
+        with pytest.raises(asyncio.CancelledError) as caught:
+            if wait_for:
+                await asyncio.wait_for(task, timeout=2)
+            else:
+                await task
+        error = caught.value
+        assert task.cancelled()  # Evidence handling must not turn cancellation into success.
+        assert exception_evidence(error, 'failure') is command
+        results = exception_evidence(error, 'completed_runs')
+        assert results == reference['benchmark_runs']
+        path = tmp_path / 'checkpoint.json'
+        BenchmarkRunner.write_checkpoint(
+            {'checkpoint_path': str(path)},
+            {'password': 'secret-value'},
+            results,
+            'cancelled',
+            error,
+        )
+        checkpoint = json.loads(path.read_text())
+        assert checkpoint['completed_iterations'] == 1
+        assert checkpoint['failed_iteration'] == reference['failed_iteration']
+        assert checkpoint['failed_command']['returncode'] == -15
+        assert 'partial' in checkpoint['failed_command']['stdout']
+        assert 'secret-value' not in path.read_text()
+        report = BenchmarkRunner.partial_report(
+            deepcopy(reference),
+            {
+                'workload_conf': {
+                    'pgbench_iter_name': 'pgbench_clients',
+                    'pgbench_iter_list': [8, 16, 64],
+                },
+                'report_conf': {'report_name': 'cancellation'},
+            },
+            error,
+            started_at='start',
+            started_clock=time.monotonic(),
+        )
+        assert report['benchmark_status'] == 'cancelled'
+        assert report['benchmark_runs'] == reference['benchmark_runs']
+        assert report['failed_iteration'] == reference['failed_iteration']
+        assert report['maximum_tps']['tps'] == 10
+
+    asyncio.run(scenario())
+
+
+def test_cancellation_lookup_respects_current_value_and_ignores_unrelated_errors():
+    inner = asyncio.CancelledError()
+    inner.benchmark_run = {'workload': 'original'}
+    outer = asyncio.CancelledError()
+    outer.__context__ = inner
+    assert exception_evidence(outer, 'benchmark_run') == inner.benchmark_run
+    outer.benchmark_run = None
+    assert exception_evidence(outer, 'benchmark_run', {}) is None
+    unrelated = RuntimeError('new operation')
+    unrelated.__context__ = inner
+    assert exception_evidence(unrelated, 'benchmark_run') is None
+    inner.__context__ = outer  # A malformed context cycle must not hang reporting.
+    assert exception_evidence(outer, 'missing', 'fallback') == 'fallback'

tests/test_core_review_regressions.py

@@ -7,7 +7,7 @@
 from pg_diag.sampler_runtime import SamplerCollection
 
 from pg_perf_bench.benchmark import BenchmarkRunner
-from pg_perf_bench.errors import CommandExecutionError, CommandFailure
+from pg_perf_bench.errors import CommandExecutionError, CommandFailure, exception_evidence
 from pg_perf_bench.executors import ProcessResult
 from pg_perf_bench.report.benchmark_status import build_execution_section
 from pg_perf_bench.system_metrics import collect_system_metrics
@@ -113,7 +113,7 @@ async def command(logger, cmd, **kwargs):
                 release_window.set()
                 with pytest.raises(asyncio.CancelledError) as caught:
                     await asyncio.wait_for(task, timeout=2)
-                evidence = caught.value.benchmark_run
+                evidence = exception_evidence(caught.value, 'benchmark_run')
                 assert evidence['status'] == 'cancelled'
                 assert evidence['system_metrics']['samples']['os.memory'] == [sample(1), sample(2)]
                 assert evidence['system_metrics']['samples']['os.disk'] == [sample(1)]

tests/test_interrupted_benchmark.py

@@ -11,7 +11,7 @@
 from pg_perf_bench.benchmark import BenchmarkRunner
 from pg_perf_bench.cli import main
 from pg_perf_bench.const import BENCHMARK_TEMPLATE_JSON_PATH
-from pg_perf_bench.errors import CommandExecutionError, CommandTimeoutError
+from pg_perf_bench.errors import CommandExecutionError, CommandTimeoutError, exception_evidence
 from pg_perf_bench.executors import ProcessResult, run_local_process
 from pg_perf_bench.report.processing import get_report_structure
 from pg_perf_bench.system_metrics import collect_system_metrics
@@ -344,8 +344,8 @@ async def scenario():
                     command_timeout=1,
                     initialization_plan=MagicMock(),
                 )
-        assert caught.value.benchmark_run['initialization'] == init
-        assert caught.value.benchmark_run['init']['returncode'] == 0
+        assert exception_evidence(caught.value, 'benchmark_run')['initialization'] == init
+        assert exception_evidence(caught.value, 'benchmark_run')['init']['returncode'] == 0
 
     asyncio.run(scenario())
 
@@ -411,7 +411,7 @@ async def command(logger, cmd, **kwargs):
                 task.cancel()
                 with pytest.raises(asyncio.CancelledError) as caught:
                     await task
-                result = caught.value.benchmark_run
+                result = exception_evidence(caught.value, 'benchmark_run')
                 assert 'tps = 1.0' in result['workload']['stdout']
                 assert 'metrics' not in result
             else: