Support stopping async run closures before cancellation compensation starts #238

# Run: uv run --no-project --python 3.13 --with restate-sdk==1.0.5 --with hypercorn --with httpx python repro.py
# Requires Docker. Runs both server versions and removes its own containers.
import asyncio
import json
import socket
import subprocess
import time
import uuid
from datetime import timedelta
from importlib.metadata import version

import httpx
import restate
from hypercorn.asyncio import serve
from hypercorn.config import Config

EVENTS = []


def emit(case, event, **fields):
    EVENTS.append(dict(case=case, event=event, t=time.monotonic(), **fields))


def port():
    with socket.socket() as sock:
        sock.bind(('127.0.0.1', 0))
        return sock.getsockname()[1]


active = restate.Service('ActiveClosure', inactivity_timeout=timedelta(seconds=60))
short = restate.Service('ShortClosure', inactivity_timeout=timedelta(milliseconds=200), abort_timeout=timedelta(seconds=60))


async def body(ctx, case):
    emit(case, 'handler_enter')

    async def work():
        emit(case, 'closure_start')
        try:
            await asyncio.sleep(20)
            emit(case, 'closure_complete')
            return 'complete'
        except asyncio.CancelledError:
            emit(case, 'closure_cancelled')
            await asyncio.sleep(0.1)
            emit(case, 'closure_drained')
            raise

    try:
        await ctx.run('work', work)
    except restate.TerminalError as exc:
        emit(case, 'handler_terminal', code=exc.status_code, message=exc.message)

        async def compensate():
            emit(case, 'compensation_start')
            await asyncio.sleep(0.5)
            emit(case, 'compensation_end')
            return 'compensated'

        await ctx.run('compensate', compensate)
        emit(case, 'compensation_returned')
        raise
    finally:
        emit(case, 'handler_finally')
    return 'done'


@active.handler()
async def run(ctx: restate.Context, case: str) -> str:
    return await body(ctx, case)


@short.handler(name='run')
async def short_run(ctx: restate.Context, case: str) -> str:
    return await body(ctx, case)


@active.handler()
async def cancel(ctx: restate.Context, invocation: str) -> str:
    ctx.cancel_invocation(invocation)
    return 'sent'


async def wait_for(predicate, seconds=40):
    async with asyncio.timeout(seconds):
        while not predicate():
            await asyncio.sleep(0.02)


async def test(image):
    EVENTS.clear()
    ingress, admin, fabric, worker = [port() for _ in range(4)]
    name = 'run-closure-repro-' + uuid.uuid4().hex
    command = ['docker', 'run', '--rm', '-d', '--name', name, '--network', 'host']
    for key, value in dict(RESTATE_BIND_IP='127.0.0.1', RESTATE_BIND_PORT=fabric,
                           RESTATE_INGRESS__BIND_PORT=ingress, RESTATE_ADMIN__BIND_PORT=admin,
                           RESTATE_ADVERTISED_ADDRESS=f'http://127.0.0.1:{fabric}/',
                           RESTATE_INGRESS__ADVERTISED_ADDRESS=f'http://127.0.0.1:{ingress}/').items():
        command.extend(['-e', f'{key}={value}'])
    subprocess.run(command + [f'docker.io/restatedev/restate:{image}'], check=True, capture_output=True)
    shutdown = asyncio.Event()
    config = Config()
    config.bind = [f'127.0.0.1:{worker}']
    server = asyncio.create_task(serve(restate.app(services=[active, short]), config, shutdown_trigger=shutdown.wait))
    summaries = []
    try:
        async with httpx.AsyncClient(timeout=45, trust_env=False) as client:
            async with asyncio.timeout(45):
                while True:
                    try:
                        if (await client.get(f'http://127.0.0.1:{admin}/health')).status_code == 200:
                            break
                    except httpx.TransportError:
                        pass
                    await asyncio.sleep(0.1)
            response = await client.post(f'http://127.0.0.1:{admin}/deployments', json={'uri': f'http://127.0.0.1:{worker}'})
            response.raise_for_status()
            for case, service in [('active', 'ActiveClosure'), ('short-inactivity', 'ShortClosure')]:
                response = await client.post(f'http://127.0.0.1:{ingress}/{service}/run/send', json=case)
                response.raise_for_status()
                invocation = response.json()['invocationId']
                await wait_for(lambda: any(e['case'] == case and e['event'] == 'closure_start' for e in EVENTS))
                await asyncio.sleep(1)
                emit(case, 'cancel_sent')
                response = await client.post(f'http://127.0.0.1:{ingress}/ActiveClosure/cancel', json=invocation)
                response.raise_for_status()
                await wait_for(lambda: any(e['case'] == case and e['event'] == 'compensation_returned' for e in EVENTS))
                await asyncio.sleep(0.5)
                rows = [e for e in EVENTS if e['case'] == case]
                first = {}
                for event in rows:
                    first.setdefault(event['event'], event['t'])
                timings = {key: round(value - first['cancel_sent'], 6) for key, value in first.items()}
                terminal = next(e for e in rows if e['event'] == 'handler_terminal')
                if case == 'active':
                    reproduced = first.get('closure_cancelled', 0) > first['compensation_returned']
                else:
                    reproduced = first['handler_terminal'] - first['cancel_sent'] > 10 and 'closure_complete' in first and 'closure_cancelled' not in first
                summary = dict(server=image, sdk=version('restate-sdk'), case=case,
                               code=terminal['code'], message=terminal['message'], seconds_from_cancel=timings,
                               verdict='STILL-REPRODUCES' if reproduced else 'CANNOT-REPRODUCE')
                summaries.append(summary)
                print(json.dumps(summary), flush=True)
    finally:
        shutdown.set()
        try:
            await asyncio.wait_for(server, 10)
        finally:
            subprocess.run(['docker', 'rm', '-f', name], check=False, capture_output=True)
    return summaries


async def main():
    summaries = []
    for image in ['1.7.13', '1.8.0']:
        summaries.extend(await test(image))
    print('VERDICT: ' + ('STILL-REPRODUCES (4/4 cases)' if all(s['verdict'] == 'STILL-REPRODUCES' for s in summaries) else 'CONFLICTING'))


if __name__ == '__main__':
    asyncio.run(main())