Skip to content

Commit eaaee12

Browse files
NathanFlurryNathanFlurry
committed
feat: redesign workflow actor integration
1 parent b86692e commit eaaee12

15 files changed

Lines changed: 608 additions & 172 deletions

File tree

‎README.md‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,14 @@ pnpm add @rivet-dev/workflows rivetkit
99
```
1010

1111
```ts
12-
import { actor } from "rivetkit";
1312
import { workflow } from "@rivet-dev/workflows";
1413

15-
export const report = actor({
16-
run: workflow(async (ctx) => {
14+
export const report = workflow({
15+
run: async (ctx) => {
1716
await ctx.step("generate", async (step) => {
1817
step.log.info("generating report");
1918
});
20-
}),
19+
},
2120
});
2221
```
2322

‎packages/workflows/README.md‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,12 @@ Durable, replayable workflows for Rivet Actors.
55
[Documentation](https://rivet.dev/workflows/docs)
66

77
```ts
8-
import { actor } from "rivetkit";
98
import { workflow } from "@rivet-dev/workflows";
109

11-
export const example = actor({
12-
run: workflow(async (ctx) => {
10+
export const example = workflow({
11+
run: async (ctx) => {
1312
await ctx.step("hello", async () => "world");
14-
}),
13+
},
1514
});
1615
```
1716

‎packages/workflows/src/rivetkit/driver.ts‎

Lines changed: 178 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,5 @@
11
import type { ActorQueue, ActorRun, RunContext } from "rivetkit";
2-
import {
3-
WORKFLOW_STORAGE_V1,
4-
type WorkflowStorageHandle,
5-
} from "rivetkit/storage";
2+
import type { RawAccess } from "rivetkit/db";
63
import type {
74
EngineDriver,
85
KVEntry,
@@ -12,6 +9,14 @@ import type {
129
WorkflowMessageIdentity,
1310
} from "../index.js";
1411

12+
const WORKFLOW_STORAGE_PREFIX = new Uint8Array([6, 1]);
13+
const WORKFLOW_UPSERT_SQL =
14+
"INSERT INTO _rivet_wf_kv (key, value) VALUES (?, ?) ON CONFLICT(key) DO UPDATE SET value = excluded.value";
15+
16+
const WORKFLOW_SQLITE_MAX_VALUE_BYTES = 256 * 1024;
17+
const WORKFLOW_SQLITE_MAX_BATCH_ROWS = 128;
18+
const WORKFLOW_SQLITE_MAX_BATCH_BYTES = 512 * 1024;
19+
1520
function track<T>(
1621
runCtx: RunContext<any, any, any, any, any, any, any, any>,
1722
promise: Promise<T>,
@@ -25,6 +30,169 @@ function track<T>(
2530
return promise;
2631
}
2732

33+
function prefixWorkflowKey(key: Uint8Array): Uint8Array {
34+
const prefixed = new Uint8Array(
35+
WORKFLOW_STORAGE_PREFIX.byteLength + key.byteLength,
36+
);
37+
prefixed.set(WORKFLOW_STORAGE_PREFIX);
38+
prefixed.set(key, WORKFLOW_STORAGE_PREFIX.byteLength);
39+
return prefixed;
40+
}
41+
42+
function stripWorkflowKey(key: Uint8Array): Uint8Array {
43+
if (
44+
key.byteLength < WORKFLOW_STORAGE_PREFIX.byteLength ||
45+
!WORKFLOW_STORAGE_PREFIX.every((byte, index) => key[index] === byte)
46+
) {
47+
throw new Error("workflow SQLite key escaped the [6, 1] namespace");
48+
}
49+
return key.slice(WORKFLOW_STORAGE_PREFIX.byteLength);
50+
}
51+
52+
function computeUpperBound(prefix: Uint8Array): Uint8Array {
53+
const upperBound = prefix.slice();
54+
for (let index = upperBound.length - 1; index >= 0; index--) {
55+
if (upperBound[index] !== 0xff) {
56+
upperBound[index]++;
57+
return upperBound.slice(0, index + 1);
58+
}
59+
}
60+
61+
// Every workflow key begins with 6, so a finite upper bound always exists.
62+
throw new Error("workflow storage prefix has no upper bound");
63+
}
64+
65+
function normalizeSqlBlob(value: unknown): Uint8Array {
66+
if (value instanceof Uint8Array) {
67+
return value;
68+
}
69+
if (value instanceof ArrayBuffer) {
70+
return new Uint8Array(value);
71+
}
72+
if (ArrayBuffer.isView(value)) {
73+
return new Uint8Array(value.buffer, value.byteOffset, value.byteLength);
74+
}
75+
if (Array.isArray(value)) {
76+
const bytes = new Uint8Array(value.length);
77+
for (const [index, byte] of value.entries()) {
78+
if (!Number.isInteger(byte) || byte < 0 || byte > 255) {
79+
throw new Error("workflow SQLite value was not a byte array");
80+
}
81+
bytes[index] = byte;
82+
}
83+
return bytes;
84+
}
85+
throw new Error("workflow SQLite value was not a blob");
86+
}
87+
88+
function validateWrites(writes: KVWrite[]): void {
89+
if (writes.length > WORKFLOW_SQLITE_MAX_BATCH_ROWS) {
90+
throw new Error(
91+
`Workflow batch contains ${writes.length} rows, exceeding the ${WORKFLOW_SQLITE_MAX_BATCH_ROWS} row limit`,
92+
);
93+
}
94+
95+
let batchBytes = 0;
96+
for (const write of writes) {
97+
if (write.value.byteLength > WORKFLOW_SQLITE_MAX_VALUE_BYTES) {
98+
throw new Error(
99+
`Workflow value is ${write.value.byteLength} bytes, exceeding the ${WORKFLOW_SQLITE_MAX_VALUE_BYTES} byte limit`,
100+
);
101+
}
102+
batchBytes +=
103+
WORKFLOW_STORAGE_PREFIX.byteLength +
104+
write.key.byteLength +
105+
write.value.byteLength;
106+
}
107+
108+
if (batchBytes > WORKFLOW_SQLITE_MAX_BATCH_BYTES) {
109+
throw new Error(
110+
`Workflow batch is ${batchBytes} bytes, exceeding the ${WORKFLOW_SQLITE_MAX_BATCH_BYTES} byte limit`,
111+
);
112+
}
113+
}
114+
115+
class WorkflowStorage {
116+
#db: RawAccess;
117+
118+
constructor(db: RawAccess) {
119+
this.#db = db;
120+
}
121+
122+
async get(key: Uint8Array): Promise<Uint8Array | null> {
123+
const rows = await this.#db.execute<{ value: unknown }>(
124+
"SELECT value FROM _rivet_wf_kv WHERE key = ?",
125+
prefixWorkflowKey(key),
126+
);
127+
const value = rows[0]?.value;
128+
return value == null ? null : normalizeSqlBlob(value);
129+
}
130+
131+
async set(key: Uint8Array, value: Uint8Array): Promise<void> {
132+
await this.batch([{ key, value }], false);
133+
}
134+
135+
async delete(key: Uint8Array): Promise<void> {
136+
await this.#db.execute(
137+
"DELETE FROM _rivet_wf_kv WHERE key = ?",
138+
prefixWorkflowKey(key),
139+
);
140+
}
141+
142+
async deletePrefix(prefix: Uint8Array): Promise<void> {
143+
const start = prefixWorkflowKey(prefix);
144+
await this.#db.execute(
145+
"DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?",
146+
start,
147+
computeUpperBound(start),
148+
);
149+
}
150+
151+
async deleteRange(start: Uint8Array, end: Uint8Array): Promise<void> {
152+
await this.#db.execute(
153+
"DELETE FROM _rivet_wf_kv WHERE key >= ? AND key < ?",
154+
prefixWorkflowKey(start),
155+
prefixWorkflowKey(end),
156+
);
157+
}
158+
159+
async list(prefix: Uint8Array): Promise<KVEntry[]> {
160+
const start = prefixWorkflowKey(prefix);
161+
const rows = await this.#db.execute<{ key: unknown; value: unknown }>(
162+
"SELECT key, value FROM _rivet_wf_kv WHERE key >= ? AND key < ? ORDER BY key ASC",
163+
start,
164+
computeUpperBound(start),
165+
);
166+
return rows.map((row) => ({
167+
key: stripWorkflowKey(normalizeSqlBlob(row.key)),
168+
value: normalizeSqlBlob(row.value),
169+
}));
170+
}
171+
172+
async batch(writes: KVWrite[], includeState: boolean): Promise<void> {
173+
if (writes.length === 0) return;
174+
validateWrites(writes);
175+
176+
const commit = async (tx: RawAccess) => {
177+
for (const write of writes) {
178+
await tx.execute(
179+
WORKFLOW_UPSERT_SQL,
180+
prefixWorkflowKey(write.key),
181+
write.value,
182+
);
183+
}
184+
};
185+
186+
if (includeState) {
187+
await this.#db.transaction(commit, {
188+
experimental: { includeState: true },
189+
});
190+
} else {
191+
await this.#db.transaction(commit);
192+
}
193+
}
194+
}
195+
28196
class ActorWorkflowMessageDriver implements WorkflowMessageDriver {
29197
#runCtx: RunContext<any, any, any, any, any, any, any, any>;
30198
#queue: ActorQueue;
@@ -95,15 +263,15 @@ export class ActorWorkflowDriver implements EngineDriver {
95263
readonly workerPollInterval = 100;
96264
readonly messageDriver: WorkflowMessageDriver;
97265
#runCtx: RunContext<any, any, any, any, any, any, any, any>;
98-
#storage: WorkflowStorageHandle;
266+
#storage: WorkflowStorage;
99267
#queue: ActorQueue;
100268
#run: ActorRun;
101269

102270
constructor(runCtx: RunContext<any, any, any, any, any, any, any, any>) {
103271
this.#runCtx = runCtx;
104272
this.messageDriver = new ActorWorkflowMessageDriver(runCtx);
105273
this.#queue = runCtx.queue;
106-
this.#storage = runCtx.storage.open(WORKFLOW_STORAGE_V1);
274+
this.#storage = new WorkflowStorage(runCtx.db);
107275
this.#run = runCtx.run;
108276
}
109277

@@ -132,9 +300,7 @@ export class ActorWorkflowDriver implements EngineDriver {
132300
}
133301

134302
async batch(writes: KVWrite[]): Promise<void> {
135-
if (writes.length === 0) return;
136-
137-
await track(this.#runCtx, this.#storage.flushWithState(writes));
303+
await track(this.#runCtx, this.#storage.batch(writes, true));
138304
}
139305

140306
async setAlarm(_workflowId: string, wakeAt: number): Promise<void> {
@@ -184,11 +350,11 @@ export class ActorWorkflowControlDriver implements EngineDriver {
184350
readonly workerPollInterval = 100;
185351
readonly messageDriver: WorkflowMessageDriver =
186352
new NoopWorkflowMessageDriver();
187-
#storage: WorkflowStorageHandle;
353+
#storage: WorkflowStorage;
188354
#run: ActorRun;
189355

190356
constructor(runCtx: RunContext<any, any, any, any, any, any, any, any>) {
191-
this.#storage = runCtx.storage.open(WORKFLOW_STORAGE_V1);
357+
this.#storage = new WorkflowStorage(runCtx.db);
192358
this.#run = runCtx.run;
193359
}
194360

@@ -217,11 +383,7 @@ export class ActorWorkflowControlDriver implements EngineDriver {
217383
}
218384

219385
async batch(writes: KVWrite[]): Promise<void> {
220-
if (writes.length === 0) {
221-
return;
222-
}
223-
224-
await this.#storage.batch(writes);
386+
await this.#storage.batch(writes, false);
225387
}
226388

227389
async setAlarm(_workflowId: string, wakeAt: number): Promise<void> {

‎packages/workflows/src/rivetkit/inspector.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
import * as transport from "rivetkit/inspector/workflow";
1+
import * as transport from "rivetkit/experimental/inspector/workflow";
22
import {
33
encodeWorkflowHistoryTransport,
44
encodeWorkflowInspectorValue,
55
type WorkflowInspectorAdapter,
6-
} from "rivetkit/inspector/workflow";
6+
} from "rivetkit/experimental/inspector/workflow";
77
import type {
88
BranchStatus,
99
BranchStatusType,

0 commit comments

Comments
 (0)