Make unread Web Push delivery multi-worker safe #552

Merged
timmy merged 1 commits from timmy/551-push-delivery-lease into main 2026-08-11 07:27:55 +00:00
5 changed files with 229 additions and 33 deletions

View File

@ -192,8 +192,12 @@ export STACKCHAIN_TRUSTED_PROXY_CIDRS='127.0.0.0/8'
export STACKCHAIN_VAPID_PUBLIC_KEY='<url-safe-public-key>' export STACKCHAIN_VAPID_PUBLIC_KEY='<url-safe-public-key>'
export STACKCHAIN_VAPID_PRIVATE_KEY='<private-key-from-secret-manager>' export STACKCHAIN_VAPID_PRIVATE_KEY='<private-key-from-secret-manager>'
export STACKCHAIN_VAPID_SUBJECT='mailto:ops@example.com' export STACKCHAIN_VAPID_SUBJECT='mailto:ops@example.com'
# Optional; defaults to 30 seconds and STACKCHAIN_STATE_DIR/push-subscriptions.sqlite3. # Optional; defaults to a 30-second poll, 10-second endpoint deadline,
# 60-second renewable cross-worker lease, and
# STACKCHAIN_STATE_DIR/push-subscriptions.sqlite3.
export STACKCHAIN_PUSH_POLL_SECONDS=30 export STACKCHAIN_PUSH_POLL_SECONDS=30
export STACKCHAIN_PUSH_SEND_TIMEOUT_SECONDS=10
export STACKCHAIN_PUSH_LEASE_SECONDS=60
export STACKCHAIN_PUSH_DB='/var/lib/stackchain-dashboard/push-subscriptions.sqlite3' export STACKCHAIN_PUSH_DB='/var/lib/stackchain-dashboard/push-subscriptions.sqlite3'
uvicorn src.main:app --host 127.0.0.1 --port 8000 uvicorn src.main:app --host 127.0.0.1 --port 8000
``` ```

View File

@ -80,6 +80,13 @@ async def _drain_authored_action_operations() -> None:
async def _push_poll_loop() -> None: async def _push_poll_loop() -> None:
interval = max(5.0, float(os.getenv("STACKCHAIN_PUSH_POLL_SECONDS", "30"))) interval = max(5.0, float(os.getenv("STACKCHAIN_PUSH_POLL_SECONDS", "30")))
send_timeout = max(
1.0, float(os.getenv("STACKCHAIN_PUSH_SEND_TIMEOUT_SECONDS", "10"))
)
lease_seconds = max(
send_timeout + 5.0,
float(os.getenv("STACKCHAIN_PUSH_LEASE_SECONDS", "60")),
)
while True: while True:
await asyncio.sleep(interval) await asyncio.sleep(interval)
try: try:
@ -87,6 +94,8 @@ async def _push_poll_loop() -> None:
_push_subscription_store, _push_subscription_store,
_push_configuration(), _push_configuration(),
notifications, notifications,
lease_seconds=lease_seconds,
send_timeout_seconds=send_timeout,
) )
except asyncio.CancelledError: except asyncio.CancelledError:
raise raise

View File

@ -1,5 +1,7 @@
import asyncio import asyncio
import json import json
import secrets
import time
from dataclasses import dataclass from dataclasses import dataclass
from typing import Awaitable, Callable from typing import Awaitable, Callable
@ -37,38 +39,68 @@ async def dispatch_unread_updates(
configuration: PushConfiguration, configuration: PushConfiguration,
unread: Callable[[], Awaitable[dict]], unread: Callable[[], Awaitable[dict]],
send: Callable[[dict, str], Awaitable[None]] | None = None, send: Callable[[dict, str], Awaitable[None]] | None = None,
*,
lease_seconds: float = 60.0,
send_timeout_seconds: float = 10.0,
) -> int: ) -> int:
if not configuration.enabled: if not configuration.enabled:
return 0 return 0
page = await unread() owner = secrets.token_urlsafe(18)
thread_ids = { acquired = await asyncio.to_thread(
int(item["notification_id"]) store.acquire_dispatch_lease,
for item in page.get("items", []) owner,
if isinstance(item, dict) and str(item.get("notification_id", "")).isdigit() now=time.time(),
} lease_seconds=lease_seconds,
count = 0 )
for delivery in store.claim_unseen(thread_ids): if not acquired:
for thread_id in delivery.thread_ids: return 0
payload = json.dumps( try:
{ page = await unread()
"title": "New work update", thread_ids = {
"body": "Tap to review it in Stackchain.", int(item["notification_id"])
"route": f"#/my-work/update/{thread_id}", for item in page.get("items", [])
"tag": f"stackchain-update-{thread_id}", if isinstance(item, dict) and str(item.get("notification_id", "")).isdigit()
}, }
separators=(",", ":"), count = 0
) for delivery in await asyncio.to_thread(store.claim_unseen, thread_ids):
try: for thread_id in delivery.thread_ids:
if send is None: still_owner = await asyncio.to_thread(
await send_web_push(delivery.subscription, payload, configuration) store.acquire_dispatch_lease,
else: owner,
await send(delivery.subscription, payload) now=time.time(),
except Exception as error: lease_seconds=lease_seconds,
status = getattr(getattr(error, "response", None), "status_code", None) )
if status in {404, 410}: if not still_owner:
store.delete_session(delivery.session_id) return count
break payload = json.dumps(
raise {
store.mark_delivered(delivery.session_id, (thread_id,)) "title": "New work update",
count += 1 "body": "Tap to review it in Stackchain.",
return count "route": f"#/my-work/update/{thread_id}",
"tag": f"stackchain-update-{thread_id}",
},
separators=(",", ":"),
)
try:
if send is None:
operation = send_web_push(
delivery.subscription, payload, configuration
)
else:
operation = send(delivery.subscription, payload)
await asyncio.wait_for(operation, timeout=send_timeout_seconds)
except Exception as error:
status = getattr(getattr(error, "response", None), "status_code", None)
if status in {404, 410}:
await asyncio.to_thread(store.delete_session, delivery.session_id)
break
# A transient provider failure belongs to this endpoint;
# leave it unseen for a later poll and continue fan-out.
continue
await asyncio.to_thread(
store.mark_delivered, delivery.session_id, (thread_id,)
)
count += 1
return count
finally:
await asyncio.to_thread(store.release_dispatch_lease, owner)

View File

@ -1,4 +1,5 @@
import json import json
import os
import sqlite3 import sqlite3
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
@ -18,6 +19,7 @@ class PushSubscriptionStore:
def __init__(self, path: str | Path): def __init__(self, path: str | Path):
self.path = Path(path) self.path = Path(path)
self.path.parent.mkdir(parents=True, exist_ok=True) self.path.parent.mkdir(parents=True, exist_ok=True)
os.chmod(self.path.parent, 0o700)
with self._connect() as connection: with self._connect() as connection:
connection.executescript( connection.executescript(
""" """
@ -33,14 +35,47 @@ class PushSubscriptionStore:
FOREIGN KEY (session_id) REFERENCES push_subscriptions(session_id) FOREIGN KEY (session_id) REFERENCES push_subscriptions(session_id)
ON DELETE CASCADE ON DELETE CASCADE
); );
CREATE TABLE IF NOT EXISTS push_dispatch_lease (
singleton INTEGER PRIMARY KEY CHECK (singleton = 1),
owner TEXT NOT NULL,
expires_at REAL NOT NULL
);
""" """
) )
os.chmod(self.path, 0o600)
def _connect(self): def _connect(self):
connection = sqlite3.connect(self.path, timeout=2) connection = sqlite3.connect(self.path, timeout=2)
connection.execute("PRAGMA foreign_keys = ON") connection.execute("PRAGMA foreign_keys = ON")
return connection return connection
def acquire_dispatch_lease(
self, owner: str, *, now: float, lease_seconds: float
) -> bool:
with self._connect() as connection:
connection.execute("BEGIN IMMEDIATE")
current = connection.execute(
"SELECT owner, expires_at FROM push_dispatch_lease WHERE singleton = 1"
).fetchone()
if current is not None and current[0] != owner and current[1] > now:
return False
connection.execute(
"""INSERT INTO push_dispatch_lease(singleton, owner, expires_at)
VALUES (1, ?, ?)
ON CONFLICT(singleton) DO UPDATE SET
owner = excluded.owner, expires_at = excluded.expires_at""",
(owner, now + lease_seconds),
)
return True
def release_dispatch_lease(self, owner: str) -> bool:
with self._connect() as connection:
result = connection.execute(
"DELETE FROM push_dispatch_lease WHERE singleton = 1 AND owner = ?",
(owner,),
)
return result.rowcount == 1
def upsert(self, session_id: str, subscription: dict) -> None: def upsert(self, session_id: str, subscription: dict) -> None:
endpoint = subscription["endpoint"] endpoint = subscription["endpoint"]
encoded = json.dumps(subscription, separators=(",", ":"), sort_keys=True) encoded = json.dumps(subscription, separators=(",", ":"), sort_keys=True)

View File

@ -1,4 +1,6 @@
import json import json
import asyncio
import os
from types import SimpleNamespace from types import SimpleNamespace
import pytest import pytest
@ -8,6 +10,32 @@ from src.push_notifications import PushConfiguration, dispatch_unread_updates
from src.push_subscription_store import PushSubscriptionStore from src.push_subscription_store import PushSubscriptionStore
def test_dispatch_lease_is_exclusive_recoverable_and_owner_fenced(tmp_path):
path = tmp_path / "push.sqlite3"
first = PushSubscriptionStore(path)
second = PushSubscriptionStore(path)
assert first.acquire_dispatch_lease("worker-a", now=100, lease_seconds=30) is True
assert second.acquire_dispatch_lease("worker-b", now=100, lease_seconds=30) is False
assert second.acquire_dispatch_lease("worker-b", now=131, lease_seconds=30) is True
assert first.release_dispatch_lease("worker-a") is False
assert second.release_dispatch_lease("worker-b") is True
def test_subscription_store_uses_private_filesystem_permissions(tmp_path):
state_dir = tmp_path / "push-state"
previous_umask = os.umask(0)
try:
database = state_dir / "push.sqlite3"
PushSubscriptionStore(database)
finally:
os.umask(previous_umask)
assert state_dir.stat().st_mode & 0o777 == 0o700
assert database.stat().st_mode & 0o777 == 0o600
def test_subscription_store_delivers_each_unread_thread_once_per_device(tmp_path): def test_subscription_store_delivers_each_unread_thread_once_per_device(tmp_path):
store = PushSubscriptionStore(tmp_path / "push.sqlite3") store = PushSubscriptionStore(tmp_path / "push.sqlite3")
subscription = { subscription = {
@ -79,6 +107,40 @@ async def test_dispatch_sends_one_privacy_safe_deep_link_per_new_thread(tmp_path
assert "Secret title" not in json.dumps(sent) assert "Secret title" not in json.dumps(sent)
@pytest.mark.anyio
async def test_concurrent_workers_do_not_dispatch_the_same_update(tmp_path):
path = tmp_path / "push.sqlite3"
first = PushSubscriptionStore(path)
second = PushSubscriptionStore(path)
first.upsert("session-a", {
"endpoint": "https://push.example/device-a",
"keys": {"p256dh": "public-key", "auth": "auth-secret"},
})
sending = asyncio.Event()
finish = asyncio.Event()
sent = []
async def unread():
return {"items": [{"notification_id": 42}]}
async def send(_subscription, payload):
sent.append(json.loads(payload)["tag"])
sending.set()
await finish.wait()
config = PushConfiguration("public", "private", "mailto:ops@example.com")
active = asyncio.create_task(dispatch_unread_updates(first, config, unread, send))
await sending.wait()
competing = await asyncio.wait_for(
dispatch_unread_updates(second, config, unread, send), timeout=0.1
)
finish.set()
assert competing == 0
assert await active == 1
assert sent == ["stackchain-update-42"]
@pytest.mark.anyio @pytest.mark.anyio
async def test_dispatch_removes_an_expired_push_endpoint(tmp_path): async def test_dispatch_removes_an_expired_push_endpoint(tmp_path):
store = PushSubscriptionStore(tmp_path / "push.sqlite3") store = PushSubscriptionStore(tmp_path / "push.sqlite3")
@ -101,6 +163,60 @@ async def test_dispatch_removes_an_expired_push_endpoint(tmp_path):
assert store.is_subscribed("session-a") is False assert store.is_subscribed("session-a") is False
@pytest.mark.anyio
async def test_transient_endpoint_failure_does_not_block_healthy_devices(tmp_path):
store = PushSubscriptionStore(tmp_path / "push.sqlite3")
for session_id in ("session-a", "session-b"):
store.upsert(session_id, {
"endpoint": f"https://push.example/{session_id}",
"keys": {"p256dh": "public-key", "auth": "auth-secret"},
})
sent = []
async def unread():
return {"items": [{"notification_id": 8}]}
async def send(subscription, _payload):
if subscription["endpoint"].endswith("session-a"):
raise RuntimeError("provider unavailable")
sent.append(subscription["endpoint"])
config = PushConfiguration("public", "private", "mailto:ops@example.com")
assert await dispatch_unread_updates(store, config, unread, send) == 1
assert sent == ["https://push.example/session-b"]
assert store.claim_unseen({8})[0].session_id == "session-a"
@pytest.mark.anyio
async def test_timed_out_endpoint_does_not_stall_push_fanout(tmp_path):
store = PushSubscriptionStore(tmp_path / "push.sqlite3")
for session_id in ("session-a", "session-b"):
store.upsert(session_id, {
"endpoint": f"https://push.example/{session_id}",
"keys": {"p256dh": "public-key", "auth": "auth-secret"},
})
sent = []
async def unread():
return {"items": [{"notification_id": 13}]}
async def send(subscription, _payload):
if subscription["endpoint"].endswith("session-a"):
await asyncio.Event().wait()
sent.append(subscription["endpoint"])
config = PushConfiguration("public", "private", "mailto:ops@example.com")
delivered = await asyncio.wait_for(
dispatch_unread_updates(
store, config, unread, send, send_timeout_seconds=0.01
),
timeout=0.2,
)
assert delivered == 1
assert sent == ["https://push.example/session-b"]
@pytest.mark.anyio @pytest.mark.anyio
async def test_authenticated_device_can_subscribe_report_status_and_unsubscribe(tmp_path, monkeypatch): async def test_authenticated_device_can_subscribe_report_status_and_unsubscribe(tmp_path, monkeypatch):
store = PushSubscriptionStore(tmp_path / "push.sqlite3") store = PushSubscriptionStore(tmp_path / "push.sqlite3")