Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions alias.ts
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ export const alias = {
'devframe/rpc/transports/ws-server': r('devframe/src/rpc/transports/ws-server.ts'),
'devframe/rpc/transports/ws-client': r('devframe/src/rpc/transports/ws-client.ts'),
'devframe/rpc/client': r('devframe/src/rpc/client.ts'),
'devframe/rpc/dump': r('devframe/src/rpc/dump/index.ts'),
'devframe/rpc/server': r('devframe/src/rpc/server.ts'),
'devframe/rpc': r('devframe/src/rpc'),
'devframe/types': r('devframe/src/types/index.ts'),
Expand Down
1 change: 1 addition & 0 deletions packages/devframe/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
"./recipes/open-helpers": "./dist/recipes/open-helpers.mjs",
"./rpc": "./dist/rpc/index.mjs",
"./rpc/client": "./dist/rpc/client.mjs",
"./rpc/dump": "./dist/rpc/dump.mjs",
"./rpc/server": "./dist/rpc/server.mjs",
"./rpc/transports/ws-client": "./dist/rpc/transports/ws-client.mjs",
"./rpc/transports/ws-server": "./dist/rpc/transports/ws-server.mjs",
Expand Down
278 changes: 278 additions & 0 deletions packages/devframe/src/rpc/dump/collect.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,278 @@
import type {
BirpcReturn,
RpcDefinitionsToFunctions,
RpcDumpClientOptions,
RpcDumpCollectionOptions,
RpcDumpDefinition,
RpcDumpStore,
RpcFunctionDefinitionAny,
} from '../types'
import { hash } from 'devframe/utils/hash'
import pLimit from 'p-limit'
import { logger } from '../diagnostics'
import { validateDefinitions } from '../validation'
import { reviveDumpError, serializeDumpError } from './error'

function getDumpRecordKey(functionName: string, args: any[]): string {
const argsHash = hash(args)
return `${functionName}---${argsHash}`
}

function getDumpFallbackKey(functionName: string): string {
return `${functionName}---fallback`
}

async function resolveGetter<T>(valueOrGetter: T | (() => Promise<T>)): Promise<T> {
return typeof valueOrGetter === 'function'
? await (valueOrGetter as () => Promise<T>)()
: valueOrGetter
}

/**
* Collects pre-computed dumps by executing functions with their defined input combinations.
* Static functions without dump config automatically get `{ inputs: [[]] }`.
*
* @example
* ```ts
* const store = await dumpFunctions([greet], context, { concurrency: 10 })
* ```
*/
export async function dumpFunctions<
T extends readonly RpcFunctionDefinitionAny[],
>(
definitions: T,
context?: any,
options: RpcDumpCollectionOptions = {},
): Promise<RpcDumpStore<RpcDefinitionsToFunctions<T>>> {
validateDefinitions(definitions)
const concurrency = options.concurrency === true
? 5
: options.concurrency === false || options.concurrency == null
? 1
: options.concurrency

const store: RpcDumpStore = {
definitions: {},
records: {},
}

// #region Definition resolution
interface TaskResolution {
handler: (...args: any[]) => any
dump: RpcDumpDefinition
definition: RpcFunctionDefinitionAny
}

const tasksResolutions: (() => Promise<undefined | TaskResolution>)[] = definitions.map(definition => async () => {
if (definition.type === 'event' || definition.type === 'action') {
return undefined
}

// Fresh setup results for each context to avoid caching issues
const setupResult = definition.setup
? await Promise.resolve(definition.setup(context))
: {}

const handler = setupResult.handler || definition.handler
if (!handler) {
throw logger.DF0024({ name: definition.name }).throw()
}

let dump = setupResult.dump ?? definition.dump
if (!dump && definition.type === 'static') {
dump = { inputs: [[]] }
}
if (!dump && definition.snapshot) {
// Sugar: run the handler once with no args, store the result as
// both the no-args record and the fallback. Any client call then
// resolves to the same snapshot — matching NMI's "getPayload()
// always returns the baked dump" shape.
dump = async (_ctx, h) => {
const output = await Promise.resolve(h(...([] as unknown as any[])))
return {
records: [{ inputs: [] as any, output }],
fallback: output,
}
}
}

if (!dump) {
return undefined
}

if (typeof dump === 'function') {
dump = await Promise.resolve(dump(context, handler))
}

// Only add to definitions if it has a dump
store.definitions[definition.name] = {
name: definition.name,
type: definition.type,
}

return {
handler,
dump,
definition,
}
})

let functionsToDump: TaskResolution[] = []
if (concurrency <= 1) {
for (const task of tasksResolutions) {
const resolution = await task()
if (resolution) {
functionsToDump.push(resolution)
}
}
}
else {
const limit = pLimit(concurrency)
functionsToDump = (await Promise.all(tasksResolutions.map(task => limit(task)))).filter(x => !!x)
}
// #endregion

// #region Dump execution
const dumpTasks: Array<() => Promise<void>> = []
for (const { definition, handler, dump } of functionsToDump) {
const { inputs, records, fallback } = dump

// Add pre-defined records
if (records) {
for (const record of records) {
const recordKey = getDumpRecordKey(definition.name, record.inputs)
store.records[recordKey] = record
}
}

// Add fallback record
if ('fallback' in dump) {
const fallbackKey = getDumpFallbackKey(definition.name)
store.records[fallbackKey] = {
inputs: [],
output: fallback,
}
}

// Add input records execution tasks
if (inputs) {
for (const input of inputs) {
dumpTasks.push(async () => {
const recordKey = getDumpRecordKey(definition.name, input)

try {
const output = await Promise.resolve(handler(...input))
store.records[recordKey] = {
inputs: input,
output,
}
}
catch (error: unknown) {
store.records[recordKey] = {
inputs: input,
error: serializeDumpError(error),
}
}
})
}
}
}

if (concurrency <= 1) {
for (const task of dumpTasks) {
await task()
}
}
else {
const limit = pLimit(concurrency)
await Promise.all(dumpTasks.map(task => limit(task)))
}
// #endregion

return store
}

/**
* Creates a client that serves pre-computed results from a dump store.
* Uses argument hashing to match calls to stored records.
*
* @example
* ```ts
* const client = createClientFromDump(store)
* await client.greet('Alice')
* ```
*/
export function createClientFromDump<T extends Record<string, any>>(
store: RpcDumpStore<T>,
options: RpcDumpClientOptions = {},
): BirpcReturn<T> {
const { onMiss } = options

const client = new Proxy({} as T, {
get(_, functionName: string) {
if (!(functionName in store.definitions)) {
throw logger.DF0025({ name: functionName }).throw()
}

return async (...args: any[]) => {
const recordKey = getDumpRecordKey(functionName, args)

const recordOrGetter = store.records[recordKey]

if (recordOrGetter) {
const record = await resolveGetter(recordOrGetter)

if (record.error) {
throw reviveDumpError(record.error)
}

if (typeof record.output === 'function') {
return await record.output()
}

return record.output
}

onMiss?.(functionName, args)

const fallbackKey = getDumpFallbackKey(functionName)
if (fallbackKey in store.records) {
const fallbackOrGetter = store.records[fallbackKey]

const fallbackRecord = await resolveGetter(fallbackOrGetter)

if (fallbackRecord && typeof fallbackRecord.output === 'function') {
return await fallbackRecord.output()
}
if (fallbackRecord)
return fallbackRecord.output
}

throw logger.DF0026({ name: functionName, args: JSON.stringify(args) }).throw()
}
},
has(_, functionName: string) {
return functionName in store.definitions
},
ownKeys() {
return Object.keys(store.definitions)
},
getOwnPropertyDescriptor(_, functionName: string) {
return functionName in store.definitions
? { configurable: true, enumerable: true, value: undefined }
: undefined
},
})

return client as any as BirpcReturn<T>
}

/**
* Filters function definitions to only those with dump definitions.
* Note: Only checks the definition itself, not setup results.
*/
export function getDefinitionsWithDumps<T extends readonly RpcFunctionDefinitionAny[]>(
definitions: T,
): RpcFunctionDefinitionAny[] {
return definitions.filter(def => def.dump !== undefined)
}
Loading
Loading