From 928184c5814ab300484a44c7d714a9c90668d006 Mon Sep 17 00:00:00 2001 From: timmy Date: Thu, 20 Aug 2026 06:38:13 +0000 Subject: [PATCH] feat: make start-day reminders lease-safe (Closes #1168) --- src/main.py | 1 + src/push_notifications.py | 69 ++++++++----- tests/test_start_day_reminders.py | 165 ++++++++++++++++++++++++++++++ 3 files changed, 209 insertions(+), 26 deletions(-) diff --git a/src/main.py b/src/main.py index 3844b75..0c1a42a 100644 --- a/src/main.py +++ b/src/main.py @@ -180,6 +180,7 @@ async def _push_poll_loop() -> None: session_statuses=dashboard_auth.managed_session_statuses, send_timeout_seconds=send_timeout, lease_seconds=lease_seconds, + max_concurrency=max_concurrency, ) channel_tasks = ( diff --git a/src/push_notifications.py b/src/push_notifications.py index 87bbfa0..0975670 100644 --- a/src/push_notifications.py +++ b/src/push_notifications.py @@ -366,6 +366,7 @@ async def dispatch_start_day_reminders( session_statuses: Callable[[list[str]], Awaitable[dict[str, str]]] | None = None, send_timeout_seconds: float = 10.0, lease_seconds: float = 60.0, + max_concurrency: int = 8, ) -> int: if not configuration.enabled: return 0 @@ -426,34 +427,50 @@ async def dispatch_start_day_reminders( "tag": f"stackchain-start-day-{plan_date}", "plan_date": plan_date, }, separators=(",", ":")) - delivered = 0 - for device, _local_day in due_devices: - try: - operation = ( - send(device.subscription, payload) if send is not None - else send_web_push(device.subscription, payload, configuration) + semaphore = asyncio.Semaphore(max(1, max_concurrency)) + + async def dispatch_device(device) -> int: + async with semaphore: + still_owner = await asyncio.to_thread( + store.acquire_dispatch_lease, + owner, + channel="start-day", + now=time.time(), + lease_seconds=max(15.0, lease_seconds, send_timeout_seconds + 5.0), ) - await asyncio.wait_for(operation, timeout=send_timeout_seconds) - except Exception as error: - status = getattr(getattr(error, "response", None), "status_code", None) - if isinstance(error, UnsafePushEndpoint) or status in {404, 410}: - await asyncio.to_thread(store.delete_session, device.session_id) - else: - await asyncio.to_thread( - store.mark_delivery_failed, - device.session_id, - "start-day", - _delivery_failure_reason(error), + if not still_owner: + return 0 + try: + operation = ( + send(device.subscription, payload) if send is not None + else send_web_push(device.subscription, payload, configuration) ) - continue - await asyncio.to_thread( - store.mark_delivery_succeeded, device.session_id, "start-day" - ) - await asyncio.to_thread( - store.mark_start_day_reminder_delivered, device.session_id, plan_date - ) - delivered += 1 - return delivered + await asyncio.wait_for(operation, timeout=send_timeout_seconds) + except Exception as error: + status = getattr(getattr(error, "response", None), "status_code", None) + if isinstance(error, UnsafePushEndpoint) or status in {404, 410}: + await asyncio.to_thread(store.delete_session, device.session_id) + else: + await asyncio.to_thread( + store.mark_delivery_failed, + device.session_id, + "start-day", + _delivery_failure_reason(error), + ) + return 0 + await asyncio.to_thread( + store.mark_delivery_succeeded, device.session_id, "start-day" + ) + await asyncio.to_thread( + store.mark_start_day_reminder_delivered, device.session_id, plan_date + ) + return 1 + + results = await asyncio.gather( + *(dispatch_device(device) for device, _local_day in due_devices), + return_exceptions=True, + ) + return sum(result for result in results if isinstance(result, int)) finally: await asyncio.to_thread(store.release_dispatch_lease, owner, channel="start-day") diff --git a/tests/test_start_day_reminders.py b/tests/test_start_day_reminders.py index b127467..6744a31 100644 --- a/tests/test_start_day_reminders.py +++ b/tests/test_start_day_reminders.py @@ -1,3 +1,4 @@ +import asyncio import json from datetime import datetime, timezone from pathlib import Path @@ -52,6 +53,170 @@ async def test_start_day_reminder_sends_one_private_prompt_for_due_plan(tmp_path assert "private/repo" not in json.dumps(sent) +@pytest.mark.anyio +async def test_start_day_reminders_bound_parallel_device_delivery(tmp_path): + store = PushSubscriptionStore(tmp_path / "push.sqlite3") + for index in range(6): + session_id = f"device-{index}" + store.upsert(session_id, { + "endpoint": f"https://push.example/{session_id}", + "keys": {"p256dh": "public-key", "auth": "auth-secret"}, + }) + store.set_start_day_preferences( + session_id, enabled=True, timezone="UTC", reminder_hour=9 + ) + + active = 0 + peak = 0 + first_wave = asyncio.Event() + release = asyncio.Event() + + async def tomorrow(): + return {"ids": ["one"], "plan_date": "2026-08-20"} + + async def send(_subscription, _payload): + nonlocal active, peak + active += 1 + peak = max(peak, active) + if active == 2: + first_wave.set() + await release.wait() + active -= 1 + + dispatch = asyncio.create_task(dispatch_start_day_reminders( + store, + PushConfiguration("public", "private", "mailto:ops@example.com"), + tomorrow, + send, + now=datetime(2026, 8, 20, 9, 0, tzinfo=timezone.utc), + max_concurrency=2, + )) + await asyncio.wait_for(first_wave.wait(), timeout=0.5) + assert peak == 2 + release.set() + assert await dispatch == 6 + assert peak == 2 + + +@pytest.mark.anyio +async def test_start_day_reminder_stops_queued_send_after_lease_loss( + tmp_path, monkeypatch +): + store = PushSubscriptionStore(tmp_path / "push.sqlite3") + for session_id in ("device-a", "device-b"): + store.upsert(session_id, { + "endpoint": f"https://push.example/{session_id}", + "keys": {"p256dh": "public-key", "auth": "auth-secret"}, + }) + store.set_start_day_preferences( + session_id, enabled=True, timezone="UTC", reminder_hour=9 + ) + + lease_checks = 0 + + def acquire(_owner, *, channel, now, lease_seconds): + nonlocal lease_checks + assert channel == "start-day" + assert now > 0 + assert lease_seconds >= 15 + lease_checks += 1 + return lease_checks <= 2 + + monkeypatch.setattr(store, "acquire_dispatch_lease", acquire) + sent = [] + + async def tomorrow(): + return {"ids": ["one"], "plan_date": "2026-08-20"} + + async def send(subscription, _payload): + sent.append(subscription["endpoint"]) + + delivered = await dispatch_start_day_reminders( + store, + PushConfiguration("public", "private", "mailto:ops@example.com"), + tomorrow, + send, + now=datetime(2026, 8, 20, 9, 0, tzinfo=timezone.utc), + max_concurrency=1, + ) + + assert delivered == 1 + assert lease_checks == 3 + assert len(sent) == 1 + + +@pytest.mark.anyio +async def test_start_day_reminder_contains_one_device_persistence_failure( + tmp_path, monkeypatch +): + store = PushSubscriptionStore(tmp_path / "push.sqlite3") + for session_id in ("device-a", "device-b"): + store.upsert(session_id, { + "endpoint": f"https://push.example/{session_id}", + "keys": {"p256dh": "public-key", "auth": "auth-secret"}, + }) + store.set_start_day_preferences( + session_id, enabled=True, timezone="UTC", reminder_hour=9 + ) + + mark_succeeded = store.mark_delivery_succeeded + + def fail_one_device(session_id, channel): + if session_id == "device-a": + raise RuntimeError("device checkpoint unavailable") + return mark_succeeded(session_id, channel) + + monkeypatch.setattr(store, "mark_delivery_succeeded", fail_one_device) + sent = [] + + async def tomorrow(): + return {"ids": ["one"], "plan_date": "2026-08-20"} + + async def send(subscription, _payload): + sent.append(subscription["endpoint"]) + + delivered = await dispatch_start_day_reminders( + store, + PushConfiguration("public", "private", "mailto:ops@example.com"), + tomorrow, + send, + now=datetime(2026, 8, 20, 9, 0, tzinfo=timezone.utc), + max_concurrency=2, + ) + + assert delivered == 1 + assert len(sent) == 2 + + +@pytest.mark.anyio +async def test_push_poll_applies_configured_concurrency_to_start_day( + monkeypatch +): + captured = {} + hold = asyncio.Event() + + async def no_wait(_seconds): + return None + + async def hold_dispatch(*_args, **_kwargs): + await hold.wait() + + async def capture_start_day(*_args, **kwargs): + captured.update(kwargs) + raise asyncio.CancelledError + + monkeypatch.setenv("STACKCHAIN_PUSH_MAX_CONCURRENCY", "3") + monkeypatch.setattr(main.asyncio, "sleep", no_wait) + monkeypatch.setattr(main, "dispatch_unread_updates", hold_dispatch) + monkeypatch.setattr(main, "dispatch_deadline_reminders", hold_dispatch) + monkeypatch.setattr(main, "dispatch_start_day_reminders", capture_start_day) + + with pytest.raises(asyncio.CancelledError): + await main._push_poll_loop() + + assert captured["max_concurrency"] == 3 + + @pytest.mark.anyio async def test_start_day_reminder_waits_for_hour_and_nonempty_due_plan(tmp_path): store = PushSubscriptionStore(tmp_path / "push.sqlite3")