Claim security outbox rows with expiring D1 leases before delivery. Add the schema migration and concurrent-drain regression coverage.
- commit
- 2767ca77a9d7eab8a2ec52e33e89b97054075a49
- parent
- 701fedd83b0fef008054b3dc694daad8305e6a4d
- author
- Dami <dami.etoile@gmail.com>
| A | worker/ | 2 | + |
| M | worker/ | 4 | +- |
| M | worker/ | 36 | ++++++++++--- |
| M | worker/ | 82 | +++++++++++++++++++++++++++--- |
@@ -102,7 +102,9 @@ CREATE TABLE IF NOT EXISTS security_notifications (
102102 created_at INTEGER NOT NULL,103103 sent_at INTEGER,104104 attempts INTEGER NOT NULL DEFAULT 0,105- last_error TEXT105+ last_error TEXT,106+ lease_owner TEXT,107+ lease_expires_at INTEGER106108 );107109 108110 CREATE INDEX IF NOT EXISTS idx_security_notifications_pending@@ -9,6 +9,8 @@ interface NotificationRow {
99 attempts: number;1010 }1111 12+const DELIVERY_LEASE_SECONDS = 60;13+1214 function subject(action: Exclude<AlertAction, null>, severity: number): string {1315 if (action === "resolve") return "[nzip] Enumeration incident resolved";1416 if (action === "summary") return "[nzip] Enumeration incident summary";@@ -87,12 +89,26 @@ export async function drainSecurityNotifications(
8789 env: Env,8890 now = Math.floor(Date.now() / 1000),8991 ): Promise<void> {90- const pending = await env.DB.prepare(92+ const candidates = await env.DB.prepare(9193 `SELECT id, subject, text, html, attempts FROM security_notifications92- WHERE sent_at IS NULL ORDER BY created_at, id LIMIT 10`,93- ).all<NotificationRow>();94+ WHERE sent_at IS NULL AND (lease_expires_at IS NULL OR lease_expires_at < ?)95+ ORDER BY created_at, id LIMIT 10`,96+ ).bind(now).all<NotificationRow>();97+98+ const owner = crypto.randomUUID();99+ for (const notification of candidates.results) {100+ const claim = await env.DB.prepare(101+ `UPDATE security_notifications SET lease_owner = ?, lease_expires_at = ?102+ WHERE id = ? AND sent_at IS NULL103+ AND (lease_expires_at IS NULL OR lease_expires_at < ?)`,104+ ).bind(105+ owner,106+ now + DELIVERY_LEASE_SECONDS,107+ notification.id,108+ now,109+ ).run();110+ if (claim.meta.changes !== 1) continue;94111 95- for (const notification of pending.results) {96112 const attempt = notification.attempts + 1;97113 try {98114 await env.EMAIL.send({@@ -104,8 +120,9 @@ export async function drainSecurityNotifications(
104120 });105121 await env.DB.prepare(106122 `UPDATE security_notifications SET sent_at = ?, attempts = attempts + 1,107- last_error = NULL WHERE id = ? AND sent_at IS NULL`,108- ).bind(now, notification.id).run();123+ last_error = NULL, lease_owner = NULL, lease_expires_at = NULL124+ WHERE id = ? AND sent_at IS NULL AND lease_owner = ?`,125+ ).bind(now, notification.id, owner).run();109126 console.log({110127 event: "security.notification_sent",111128 notificationId: notification.id,@@ -114,9 +131,10 @@ export async function drainSecurityNotifications(
114131 } catch (error) {115132 const message = errorMessage(error);116133 await env.DB.prepare(117- `UPDATE security_notifications SET attempts = attempts + 1, last_error = ?118- WHERE id = ? AND sent_at IS NULL`,119- ).bind(message, notification.id).run();134+ `UPDATE security_notifications SET attempts = attempts + 1, last_error = ?,135+ lease_owner = NULL, lease_expires_at = NULL136+ WHERE id = ? AND sent_at IS NULL AND lease_owner = ?`,137+ ).bind(message, notification.id, owner).run();120138 console.error({121139 event: "security.notification_failed",122140 notificationId: notification.id,@@ -147,9 +147,13 @@ class FakeStatement {
147147 return { results: this.db.probes };148148 }149149 if (this.sql.includes("FROM security_notifications")) {150+ const [now] = this.args;150151 return {151152 results: [...this.db.notifications.values()]152- .filter((row) => row.sent_at === null)153+ .filter((row) =>154+ row.sent_at === null &&155+ (row.lease_expires_at === null || row.lease_expires_at < now)156+ )153157 .sort((a, b) => a.created_at - b.created_at),154158 };155159 }@@ -227,6 +231,8 @@ class FakeStatement {
227231 sent_at: null,228232 attempts: 0,229233 last_error: null,234+ lease_owner: null,235+ lease_expires_at: null,230236 });231237 }232238 return { success: true };@@ -255,28 +261,47 @@ class FakeStatement {
255261 sent_at: null,256262 attempts: 0,257263 last_error: null,264+ lease_owner: null,265+ lease_expires_at: null,258266 });259267 }260268 return { success: true };261269 }270+ if (this.sql.includes("SET lease_owner = ?")) {271+ const [owner, leaseExpiresAt, id, now] = this.args;272+ const row = this.db.notifications.get(id);273+ if (274+ row?.sent_at === null &&275+ (row.lease_expires_at === null || row.lease_expires_at < now)276+ ) {277+ row.lease_owner = owner;278+ row.lease_expires_at = leaseExpiresAt;279+ return { success: true, meta: { changes: 1 } };280+ }281+ return { success: true, meta: { changes: 0 } };282+ }262283 if (this.sql.includes("SET sent_at = ?")) {263- const [sentAt, id] = this.args;284+ const [sentAt, id, owner] = this.args;264285 const row = this.db.notifications.get(id);265- if (row?.sent_at === null) {286+ if (row?.sent_at === null && row.lease_owner === owner) {266287 row.sent_at = sentAt;267288 row.attempts += 1;268289 row.last_error = null;290+ row.lease_owner = null;291+ row.lease_expires_at = null;269292 }270- return { success: true };293+ return { success: true, meta: { changes: 1 } };271294 }272295 if (this.sql.includes("SET attempts = attempts + 1")) {273- const [lastError, id] = this.args;296+ const [lastError, id, owner] = this.args;274297 const row = this.db.notifications.get(id);275- if (row?.sent_at === null) {298+ if (row?.sent_at === null && row.lease_owner === owner) {276299 row.attempts += 1;277300 row.last_error = lastError;301+ row.lease_owner = null;302+ row.lease_expires_at = null;278303 }279- return { success: true };304+ return { success: true, meta: { changes: 1 } };280305 }281306 throw new Error(`unexpected run(): ${this.sql}`);282307 }@@ -389,3 +414,46 @@ test("the daily digest is durably deduplicated by UTC day", async () => {
389414 assert.equal(notification.sent_at, now + 60);390415 assert.equal(notification.attempts, 2);391416 });417+418+test("concurrent cron drains claim a security email only once", async () => {419+ const db = new FakeDb();420+ db.notifications.set("security:daily-digest:0", {421+ id: "security:daily-digest:0",422+ incident_name: "security",423+ action: "daily-digest",424+ window_bucket: 0,425+ subject: "[nzip] Daily security activity summary",426+ text: "activity",427+ html: null,428+ created_at: 100,429+ sent_at: null,430+ attempts: 0,431+ last_error: null,432+ lease_owner: null,433+ lease_expires_at: null,434+ });435+ let sends = 0;436+ const env = {437+ DB: db,438+ ALERT_EMAIL_FROM: "security@example.com",439+ ALERT_EMAIL_TO: "operator@example.com",440+ EMAIL: {441+ async send() {442+ sends += 1;443+ await new Promise((resolve) => setImmediate(resolve));444+ },445+ },446+ };447+448+ await Promise.all([449+ drainSecurityNotifications(env, 200),450+ drainSecurityNotifications(env, 200),451+ ]);452+453+ assert.equal(sends, 1);454+ const notification = db.notifications.get("security:daily-digest:0");455+ assert.equal(notification.sent_at, 200);456+ assert.equal(notification.attempts, 1);457+ assert.equal(notification.lease_owner, null);458+ assert.equal(notification.lease_expires_at, null);459+});