Repository navigation
feat(worker): retain raw agent output in R2 with a per-attempt cap #142
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
5 commits
Select commit
Hold shift + click to select a range
ba1bd10
chore: claim quest/m0/kix/output
kixelated 0245679
feat(worker): retain raw agent output in R2 with a per-attempt cap
kixelated b0975be
fix(worker): never page raw output past a chunk that may still land
kixelated ef024f4
fix(worker): time output retention from the upload and requeue failed…
kixelated 7a179e8
Merge remote-tracking branch 'origin/main' into ready-output
kixelated File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| -- Raw agent output moves to R2, one object per chunk at | ||
| -- output/<job>/<attempt>/<seq>. Nothing is deployed yet, so the D1 text is | ||
| -- dropped rather than copied. | ||
| DROP TABLE job_output; | ||
| -- Each attempt's output totals. It lives as long as the job, so expiring | ||
| -- chunks never give capacity back and an expired `seq` is never reused. | ||
| CREATE TABLE job_output_attempts ( | ||
| job TEXT NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, | ||
| attempt INTEGER NOT NULL, | ||
| runner TEXT NOT NULL, | ||
| -- The only `seq` the attempt accepts next; also its count of chunks. | ||
| next INTEGER NOT NULL DEFAULT 0, | ||
| bytes INTEGER NOT NULL DEFAULT 0, | ||
| -- Set once a chunk is refused for the attempt's cap; later chunks are too. | ||
| truncated INTEGER NOT NULL DEFAULT 0 CHECK(truncated IN (0,1)), | ||
| PRIMARY KEY(job,attempt) | ||
| ); | ||
| -- The index of stored chunks, readable only by current project maintainers. | ||
| -- A row reserves its `seq` for one text, by hash, before the object is | ||
| -- written; `complete` marks the object written. The daily cron deletes rows | ||
| -- older than the bucket's 30-day lifecycle. | ||
| CREATE TABLE job_output ( | ||
| job TEXT NOT NULL REFERENCES jobs(id) ON DELETE CASCADE, | ||
| attempt INTEGER NOT NULL, | ||
| seq INTEGER NOT NULL, | ||
| hash TEXT NOT NULL, | ||
| complete INTEGER NOT NULL DEFAULT 0 CHECK(complete IN (0,1)), | ||
| at INTEGER NOT NULL, | ||
| PRIMARY KEY(job,attempt,seq) | ||
| ); | ||
| CREATE INDEX job_output_at ON job_output(at); | ||
| -- Jobs whose output objects a disconnect still has to delete. Written in the | ||
| -- same batch that deletes the project, and retried by the daily cron. | ||
| CREATE TABLE output_cleanup ( | ||
| job TEXT PRIMARY KEY NOT NULL, | ||
| at INTEGER NOT NULL | ||
| ); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,183 @@ | ||
| // Raw agent output: one R2 object per chunk, indexed in D1. | ||
| // | ||
| // An attempt's chunks arrive in order, `seq` 0, 1, 2, ... Each is reserved | ||
| // first: an index row records its text's hash while the job exists, so a | ||
| // conflicting chunk is refused before anything reaches R2. The object is | ||
| // written next and the row marked complete, only while the attempt is | ||
| // current. A chunk is reserved only once the one before it is complete, so an | ||
| // attempt's one pending chunk is its last, and readers that see complete rows | ||
| // never page past a chunk that may still land. R2 cannot join a D1 batch, so | ||
| // a disconnect or reassignment between those steps leaves no row to complete, | ||
| // and the upload deletes its own object. | ||
| import { JobError, type Attempt, type Runner } from "./store"; | ||
|
|
||
| type Storage = Pick<Env, "DB" | "OUTPUT">; | ||
|
|
||
| /** Raw output kept per attempt, in UTF-8 bytes. */ | ||
| export const MAX_BYTES = 64 * 1024 * 1024; | ||
| /** Chunks kept per attempt, so tiny chunks cannot multiply rows and objects. */ | ||
| export const MAX_CHUNKS = 4096; | ||
| /** How long output is kept after each chunk's upload; the bucket lifecycle rule matches it. */ | ||
| export const RETENTION = 30 * 24 * 60 * 60_000; | ||
|
|
||
| const key = (job: string, attempt: number, seq: number) => `output/${job}/${attempt}/${seq}`; | ||
|
|
||
| async function digest(text: string) { | ||
| const bytes = await crypto.subtle.digest("SHA-256", new TextEncoder().encode(text)); | ||
| return Array.from(new Uint8Array(bytes), (b) => b.toString(16).padStart(2, "0")).join(""); | ||
| } | ||
|
|
||
| /** | ||
| * Store a chunk of the current attempt's output, in any job state. Only the | ||
| * attempt's next `seq` is accepted, or an identical retry of one still | ||
| * indexed. Past the attempt's cap, chunks are refused as `truncated`; the | ||
| * job carries on. | ||
| */ | ||
| export async function append( | ||
| { DB: db, OUTPUT: bucket }: Storage, | ||
| runner: Runner, | ||
| { job, attempt, seq, text }: Attempt & { seq: number; text: string }, | ||
| now: number, | ||
| ): Promise<{ truncated?: true }> { | ||
| const bytes = new TextEncoder().encode(text).length; | ||
| const hash = await digest(text); | ||
| const current = "EXISTS(SELECT 1 FROM jobs WHERE id=?1 AND attempt=?2 AND runner=?3)"; | ||
| const next = "job=?1 AND attempt=?2 AND runner=?3 AND next=?4"; | ||
| const params = [job, attempt, runner.id, seq, bytes, hash, now]; | ||
| const statement = (sql: string, n: number) => db.prepare(sql).bind(...params.slice(0, n)); | ||
| const [, , , , read] = await db.batch([ | ||
| statement( | ||
| `INSERT INTO job_output_attempts(job,attempt,runner) SELECT ?1,?2,?3 WHERE ${current} ON CONFLICT DO NOTHING`, | ||
| 3, | ||
| ), | ||
| // The truncation is recorded once, by the first chunk over the cap. | ||
| statement( | ||
| `UPDATE job_output_attempts SET truncated=1 WHERE ${next} AND truncated=0 AND (bytes+?5>${MAX_BYTES} OR next>=${MAX_CHUNKS}) AND ${current}`, | ||
| 5, | ||
| ), | ||
| statement( | ||
| `INSERT INTO job_output(job,attempt,seq,hash,at) SELECT ?1,?2,?4,?6,?7 FROM job_output_attempts WHERE ${next} AND truncated=0 AND ${current} | ||
| AND NOT EXISTS(SELECT 1 FROM job_output WHERE job=?1 AND attempt=?2 AND seq=?4-1 AND complete=0)`, | ||
| 7, | ||
| ), | ||
| statement( | ||
| `UPDATE job_output_attempts SET next=next+1,bytes=bytes+?5 WHERE ${next} AND EXISTS(SELECT 1 FROM job_output WHERE job=?1 AND attempt=?2 AND seq=?4)`, | ||
| 5, | ||
| ), | ||
| statement( | ||
| `SELECT a.runner,a.next,a.truncated,o.hash,o.complete,${current} AS current FROM (SELECT 1) | ||
| LEFT JOIN job_output_attempts a ON a.job=?1 AND a.attempt=?2 LEFT JOIN job_output o ON o.job=?1 AND o.attempt=?2 AND o.seq=?4`, | ||
| 4, | ||
| ), | ||
| ]); | ||
| const row = read!.results[0] as { | ||
| runner: string | null; | ||
| next: number | null; | ||
| truncated: number | null; | ||
| hash: string | null; | ||
| complete: number | null; | ||
| current: number; | ||
| }; | ||
| if (row.hash !== null && row.runner === runner.id) { | ||
| if (row.hash !== hash) throw new JobError("A different chunk has that sequence number"); | ||
| // A retry of a completed chunk writes nothing, so its lifecycle does not restart. | ||
| if (row.complete) return {}; | ||
| const name = key(job, attempt, seq); | ||
| await bucket.put(name, text); | ||
| const done = await db | ||
| .prepare( | ||
| `UPDATE job_output SET complete=1,at=?6 WHERE job=?1 AND attempt=?2 AND seq=?4 AND hash=?5 AND ${current}`, | ||
| ) | ||
| // Retention starts with the upload, as the bucket lifecycle's does. | ||
| .bind(job, attempt, runner.id, seq, hash, now) | ||
| .run(); | ||
| if (done.meta.changes === 1) return {}; | ||
| // The job was disconnected or reassigned while the object was written. A | ||
| // disconnect's cleanup may already have listed the prefix, so queue | ||
| // another in case this delete fails. | ||
| await db | ||
| .prepare( | ||
| "INSERT OR IGNORE INTO output_cleanup(job,at) SELECT ?1,?2 WHERE NOT EXISTS(SELECT 1 FROM jobs WHERE id=?1)", | ||
| ) | ||
| .bind(job, now) | ||
| .run(); | ||
| await bucket.delete(name); | ||
|
kixelated marked this conversation as resolved.
Resolved
|
||
| throw new JobError("Attempt is no longer current"); | ||
| } | ||
| if (!row.current || row.runner !== runner.id || row.next === null) | ||
| throw new JobError("Attempt is no longer current"); | ||
| if (seq < row.next) throw new JobError("That chunk has expired"); | ||
| if (row.truncated) return { truncated: true }; | ||
| if (seq === row.next) throw new JobError(`Chunk ${seq - 1} is not stored yet`); | ||
| throw new JobError(`Expected sequence ${row.next}`); | ||
| } | ||
|
|
||
| /** | ||
| * A page of stored raw output, at most a few MiB, in order; `next` continues | ||
| * it. Only stored chunks are listed: an attempt's one pending chunk is its | ||
| * last and lands only while the attempt is current, so a cursor never passes | ||
| * a chunk that may still land. | ||
| * `truncated` lists the attempts whose output reached the cap. | ||
| */ | ||
| export async function read( | ||
| { DB: db, OUTPUT: bucket }: Storage, | ||
| job: string, | ||
| after: { attempt: number; seq: number } | null, | ||
| ) { | ||
| const page = 4; | ||
| const [rows, capped] = await db.batch([ | ||
| db | ||
| .prepare( | ||
| "SELECT attempt,seq,at FROM job_output WHERE job=? AND complete=1 AND (attempt,seq)>(?,?) ORDER BY attempt,seq LIMIT ?", | ||
| ) | ||
| .bind(job, after?.attempt ?? 0, after?.seq ?? -1, page), | ||
| db.prepare("SELECT attempt FROM job_output_attempts WHERE job=? AND truncated=1 ORDER BY attempt").bind(job), | ||
| ]); | ||
| const index = rows!.results as { attempt: number; seq: number; at: number }[]; | ||
| const chunks = await Promise.all( | ||
| index.map(async (chunk) => { | ||
| const object = await bucket.get(key(job, chunk.attempt, chunk.seq)); | ||
| // The lifecycle rule may delete an object before the cron deletes its row. | ||
| return object && { ...chunk, text: await object.text() }; | ||
| }), | ||
| ); | ||
| const last = index.length === page ? index.at(-1)! : null; | ||
| return { | ||
| output: chunks.filter((chunk) => chunk !== null), | ||
| next: last && { attempt: last.attempt, seq: last.seq }, | ||
| truncated: (capped!.results as { attempt: number }[]).map((row) => row.attempt), | ||
| }; | ||
| } | ||
|
|
||
| /** Delete the output objects of jobs a disconnect removed. A failure stays queued for the next run. */ | ||
| export async function purge({ DB: db, OUTPUT: bucket }: Storage) { | ||
| const { results } = await db.prepare("SELECT job FROM output_cleanup ORDER BY at,job").all<{ job: string }>(); | ||
| for (const { job } of results) { | ||
| try { | ||
| const prefix = `output/${job}/`; | ||
| for ( | ||
| let listed = await bucket.list({ prefix }); | ||
| listed.objects.length; | ||
| listed = await bucket.list({ prefix }) | ||
| ) | ||
| await bucket.delete(listed.objects.map((object) => object.key)); | ||
| await db.prepare("DELETE FROM output_cleanup WHERE job=?").bind(job).run(); | ||
| } catch (error) { | ||
| console.error({ event: "output_cleanup_failed", name: error instanceof Error ? error.name : "unknown" }); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /** The daily cron: forget expired chunks, whose objects the lifecycle rule deletes, and retry cleanups. */ | ||
| export async function retain(storage: Storage, now: number) { | ||
| const batch = 1000; | ||
| for (;;) { | ||
| const deleted = await storage.DB.prepare( | ||
| "DELETE FROM job_output WHERE rowid IN (SELECT rowid FROM job_output WHERE at<=? LIMIT ?)", | ||
| ) | ||
| .bind(now - RETENTION, batch) | ||
| .run(); | ||
| if (deleted.meta.changes < batch) break; | ||
| } | ||
| await purge(storage); | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.