SnowLuma/packages/onebot/src/message-store-migration-task.ts

import { createRequire } from 'node:module';
import path from 'node:path';
import { Worker, type WorkerOptions } from 'node:worker_threads';
import { createLogger } from '@snowluma/common/logger';
import type { DatabaseMigrationCallbacks, DatabaseMigrationTask } from './manager';
import type { MessageStoreMigrationStatus } from './message-store-migration';
import {
    MESSAGE_STORE_WORKER_KIND,
    type MessageStoreMigrationWorkerData,
    type MessageStoreMigrationWorkerMessage,
} from './message-store-migration-worker';

export interface MessageStoreMigrationWorkerHandle {
    on(event: 'message', listener: (message: MessageStoreMigrationWorkerMessage) => void): this;
    once(event: 'error', listener: (error: Error) => void): this;
    once(event: 'exit', listener: (code: number) => void): this;
    postMessage(message: 'start' | 'cancel'): void;
    unref(): void;
}

type CreateMigrationWorker = (
    filename: string,
    options: WorkerOptions,
) => MessageStoreMigrationWorkerHandle;

export function createMessageStoreMigrationTask(
    uin: string,
    createWorker: CreateMigrationWorker = (filename, options) => new Worker(filename, options),
): DatabaseMigrationTask {
    const log = createLogger('OneBot.MessageMigration').child({ uin: Number(uin) });
    let worker: MessageStoreMigrationWorkerHandle | null = null;
    let cancelled = false;
    let finished = false;
    let failureReported = false;

    return {
        start(callbacks: DatabaseMigrationCallbacks): void {
            const reportFailure = (error: Error): void => {
                if (cancelled || failureReported) {
                    return;
                }

                failureReported = true;
                callbacks.onFailed(error);
            };

            const data: MessageStoreMigrationWorkerData = {
                kind: MESSAGE_STORE_WORKER_KIND,
                dbPath: path.join('data', uin, 'messages.db'),
            };
            const entryPath = process.argv[1];
            const sourceEntry = /\.(?:cts|mts|ts|tsx)$/.test(entryPath);
            worker = createWorker(sourceEntry ? sourceWorkerBootstrap(entryPath) : entryPath, {
                workerData: data,
                eval: sourceEntry,
                execArgv: sourceEntry ? [] : process.execArgv,
            });
            worker.on('message', (message: MessageStoreMigrationWorkerMessage) => {
                if (cancelled) {
                    return;
                }

                if (message.kind === 'ready') {
                    callbacks.onReady();
                } else if (message.kind === 'preparation') {
                    const progress = message.progress;

                    if (progress.stage === 'started') {
                        log.info('preparing message history migration');
                    } else if (progress.stage === 'indexed') {
                        log.info(
                            'message history migration preparation finished in %dms',
                            Math.round(progress.elapsedMs),
                        );
                    } else {
                        log.info(
                            'message history migration: pending=%d checked in %dms',
                            progress.total,
                            Math.round(progress.elapsedMs),
                        );
                    }
                } else if (message.kind === 'progress') {
                    callbacks.onProgress(message.status, message.elapsedMs);

                    if (message.status.phase === 'complete') {
                        finished = true;
                    }
                } else {
                    reportFailure(new Error(message.message));
                }
            });
            worker.once('error', reportFailure);
            worker.once('exit', (code) => {
                if (!cancelled && !finished) {
                    reportFailure(
                        new Error(
                            `database migration worker exited before completion with code ${code}`,
                        ),
                    );
                }
            });
        },
        beginMigration(): void {
            if (!cancelled) {
                worker?.postMessage('start');
            }
        },
        cancel(): void {
            cancelled = true;
            worker?.postMessage('cancel');
            worker?.unref();
        },
    };
}

function sourceWorkerBootstrap(entryPath: string): string {
    const require = createRequire(path.resolve(entryPath));
    return [
        `const { register } = require(${JSON.stringify(require.resolve('tsx/esm/api'))});`,
        'register();',
        `import(${JSON.stringify(path.resolve(entryPath))});`,
    ].join('\n');
}

export function estimateRemainingSeconds(
    status: MessageStoreMigrationStatus,
    recentRowsPerSecond: number | null,
): number | null {
    if (status.phase === 'complete') {
        return 0;
    }

    if (!recentRowsPerSecond || recentRowsPerSecond <= 0) {
        return null;
    }

    return Math.max(0, Math.ceil((status.total - status.processed) / recentRowsPerSecond));
}