| 1 | export type PollErrorMapperData<E = string> |
| 2 | = | { type: 'failed', taskName: string, error: E } |
| 3 | | { type: 'completed_without_data', taskName: string } |
| 4 | | { type: 'timeout', taskName: string, maxPollingMs: number } |
| 5 | |
| 6 | export interface PollOptions<E = string> { |
| 7 | /** 最大轮询时长（毫秒），默认 15 分钟 */ |
| 8 | maxPollingMs?: number |
| 9 | /** 轮询间隔（毫秒），默认 10 秒 */ |
| 10 | intervalMs?: number |
| 11 | /** 任务名称，用于超时错误信息 */ |
| 12 | taskName?: string |
| 13 | /** 映射 poll 自身产生的失败数据 */ |
| 14 | errorMapper?: (data: PollErrorMapperData<E>) => Error |
| 15 | } |
| 16 | |
| 17 | export interface PollResult<T, E = string> { |
| 18 | /** 是否完成（成功或失败） */ |
| 19 | done: boolean |
| 20 | /** 完成时的返回数据 */ |
| 21 | data?: T |
| 22 | /** 失败时的原始错误数据 */ |
| 23 | error?: E |
| 24 | } |
| 25 | |
| 26 | const DEFAULT_POLL_OPTIONS = { |
| 27 | maxPollingMs: 15 * 60 * 1000, |
| 28 | intervalMs: 10_000, |
| 29 | taskName: 'Task', |
| 30 | } |
| 31 | |
| 32 | function defaultPollErrorMapper<E>(data: PollErrorMapperData<E>): Error { |
| 33 | switch (data.type) { |
| 34 | case 'failed': |
| 35 | return new Error(`${data.taskName} failed: ${data.error}`) |
| 36 | case 'completed_without_data': |
| 37 | return new Error(`${data.taskName} completed without data`) |
| 38 | case 'timeout': |
| 39 | return new Error(`${data.taskName} timed out after ${Math.round(data.maxPollingMs / 60_000)} minutes`) |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | /** |
| 44 | * 通用轮询工具：按固定间隔调用 pollFn 直到完成或超时 |
| 45 | * |
| 46 | * @param pollFn - 每次轮询调用的函数，返回 PollResult |
| 47 | * @param options - 轮询配置 |
| 48 | * @returns 轮询成功后的数据 |
| 49 | * |
| 50 | * @example |
| 51 | * const videoUrl = await poll( |
| 52 | * async () => { |
| 53 | * const result = await videoService.getTask(taskId) |
| 54 | * if (result.videoUrl) return { done: true, data: result.videoUrl } |
| 55 | * if (result.error) return { done: true, error: result.error } |
| 56 | * return { done: false } |
| 57 | * }, |
| 58 | * { maxPollingMs: 10 * 60 * 1000, taskName: 'Video generation' }, |
| 59 | * ) |
| 60 | */ |
| 61 | export async function poll<T, E = string>( |
| 62 | pollFn: () => Promise<PollResult<T, E>>, |
| 63 | options?: PollOptions<E>, |
| 64 | ): Promise<T> { |
| 65 | const { maxPollingMs, intervalMs, taskName, errorMapper } = { ...DEFAULT_POLL_OPTIONS, ...options } |
| 66 | const mapError = errorMapper || defaultPollErrorMapper<E> |
| 67 | const startTime = Date.now() |
| 68 | |
| 69 | while (Date.now() - startTime < maxPollingMs) { |
| 70 | await new Promise(resolve => setTimeout(resolve, intervalMs)) |
| 71 | const { done, data, error } = await pollFn() |
| 72 | |
| 73 | if (error !== undefined) { |
| 74 | throw mapError({ type: 'failed', taskName, error }) |
| 75 | } |
| 76 | if (done && data !== undefined) { |
| 77 | return data |
| 78 | } |
| 79 | if (done) { |
| 80 | throw mapError({ type: 'completed_without_data', taskName }) |
| 81 | } |
| 82 | } |
| 83 | |
| 84 | throw mapError({ type: 'timeout', taskName, maxPollingMs }) |
| 85 | } |
| 86 |