diff --git a/packages/plugin-eval-addons/src/index.ts b/packages/plugin-eval-addons/src/index.ts index a5d726b312..868fb4a202 100644 --- a/packages/plugin-eval-addons/src/index.ts +++ b/packages/plugin-eval-addons/src/index.ts @@ -15,7 +15,7 @@ export interface Config { exclude?: RegExp } -declare module 'koishi-plugin-eval' { +declare module 'koishi-plugin-eval/dist/main' { interface MainConfig extends AddonConfig {} } @@ -42,12 +42,12 @@ interface Manifest { } export function apply(ctx: Context, config: Config) { - const { evalConfig } = ctx.app - Object.assign(evalConfig, config) - const root = resolve(process.cwd(), assertProperty(evalConfig, 'moduleRoot')) - evalConfig.moduleRoot = root - evalConfig.dataKeys.push('addonNames', 'moduleRoot') - evalConfig.setupFiles['koishi/addons.ts'] = resolve(__dirname, 'worker.js') + const { worker } = ctx.app + Object.assign(worker.config, config) + const root = resolve(process.cwd(), assertProperty(worker.config, 'moduleRoot')) + worker.config.moduleRoot = root + worker.config.dataKeys.push('addonNames', 'moduleRoot') + worker.config.setupFiles['koishi/addons.ts'] = resolve(__dirname, 'worker.js') const git = Git(root) @@ -55,9 +55,9 @@ export function apply(ctx: Context, config: Config) { .option('update', '-u 更新扩展模块', { authority: 3 }) .action(async ({ options, session }) => { if (options.update) { - const { files, summary } = await git.pull(evalConfig.gitRemote) + const { files, summary } = await git.pull(worker.config.gitRemote) if (!files.length) return '所有模块均已是最新。' - await session.$app.evalWorker.terminate() + await session.$app.worker.restart() return `更新成功！(${summary.insertions}A ${summary.deletions}D ${summary.changes}M)` } return session.$execute('help addon') @@ -69,11 +69,11 @@ export function apply(ctx: Context, config: Config) { }) let manifests: Record> - const { exclude = /^(\..+|node_modules)$/ } = evalConfig + const { exclude = /^(\..+|node_modules)$/ } = worker.config ctx.on('worker/start', async () => { const dirents = await fs.readdir(root, { withFileTypes: true }) - const paths = evalConfig.addonNames = dirents + const paths = worker.config.addonNames = dirents .filter(dir => dir.isDirectory() && !exclude.test(dir.name)) .map(dir => dir.name) // cmd.dispose() may affect addon.children, so here we make a slice @@ -87,7 +87,7 @@ export function apply(ctx: Context, config: Config) { } ctx.on('worker/ready', (response) => { - evalConfig.addonNames.map(async (path) => { + worker.config.addonNames.map(async (path) => { const manifest = await manifests[path] if (!manifest) return const { commands = [] } = manifest @@ -105,7 +105,7 @@ export function apply(ctx: Context, config: Config) { UserTrap.attach(cmd, userFields, async ({ session, command, options, user, writable }, ...args) => { const { $app, $uuid } = session const { name } = command - const result = await $app.evalRemote.callAddon($uuid, user, writable, { name, args, options }) + const result = await $app.worker.remote.callAddon($uuid, user, writable, { name, args, options }) return result }) diff --git a/packages/plugin-eval/src/index.ts b/packages/plugin-eval/src/index.ts index c985d4d381..3cb870f495 100644 --- a/packages/plugin-eval/src/index.ts +++ b/packages/plugin-eval/src/index.ts @@ -1,19 +1,13 @@ -import { Context, Session, User } from 'koishi-core' -import { Worker, ResourceLimits } from 'worker_threads' -import { CQCode, Logger, defineProperty, Random, pick } from 'koishi-utils' -import { WorkerAPI, WorkerConfig, WorkerData, Response } from './worker' -import { wrap, expose, Remote } from './transfer' -import { MainAPI, Access, UserTrap } from './main' -import { resolve } from 'path' +import { Context, Session } from 'koishi-core' +import { CQCode, Logger, defineProperty, Random } from 'koishi-utils' +import { EvalWorker, UserTrap, EvalConfig, Config } from './main' -export * from './main' +export { UserTrap, MainAPI } from './main' declare module 'koishi-core/dist/app' { interface App { _sessions: Record - evalConfig: EvalConfig - evalWorker: Worker - evalRemote: Remote + worker: EvalWorker } } @@ -23,14 +17,6 @@ declare module 'koishi-core/dist/command' { } } -declare module 'koishi-core/dist/context' { - interface EventMap { - 'worker/start'(): void | Promise - 'worker/ready'(response: Response): void - 'worker/exit'(): void - } -} - declare module 'koishi-core/dist/session' { interface Session { $uuid: string @@ -39,19 +25,6 @@ declare module 'koishi-core/dist/session' { } } -export interface MainConfig { - prefix?: string - timeout?: number - maxLogs?: number - userFields?: Access - resourceLimits?: ResourceLimits - dataKeys?: (keyof WorkerData)[] -} - -export interface EvalConfig extends MainConfig, WorkerData {} - -export interface Config extends MainConfig, WorkerConfig {} - const defaultConfig: EvalConfig = { prefix: '>', timeout: 1000, @@ -63,48 +36,14 @@ const defaultConfig: EvalConfig = { const logger = new Logger('eval') -export const workerScript = `require(${JSON.stringify(resolve(__dirname, 'worker.js'))});` +export const name = 'eval' export function apply(ctx: Context, config: Config = {}) { const { prefix } = config = { ...defaultConfig, ...config } const { app } = ctx + const worker = new EvalWorker(app, config) defineProperty(app, '_sessions', {}) - defineProperty(app, 'evalConfig', config) - defineProperty(app, 'evalRemote', null) - defineProperty(app, 'evalWorker', null) - - let restart = true - const api = new MainAPI(app) - async function createWorker() { - await app.parallel('worker/start') - - const worker = app.evalWorker = new Worker(workerScript, { - eval: true, - workerData: { - logLevels: Logger.levels, - ...pick(config, config.dataKeys), - }, - resourceLimits: config.resourceLimits, - }) - - expose(worker, api) - - app.evalRemote = wrap(worker) - await app.evalRemote.start().then((response) => { - app.emit('worker/ready', response) - logger.info('worker started') - - worker.on('exit', (code) => { - ctx.emit('worker/exit') - logger.info('exited with code', code) - if (restart) createWorker() - }) - }) - } - - process.on('beforeExit', () => { - restart = false - }) + defineProperty(app, 'worker', worker) app.prependMiddleware((session, next) => { app._sessions[session.$uuid = Random.uuid()] = session @@ -116,7 +55,7 @@ export function apply(ctx: Context, config: Config = {}) { }) app.on('before-connect', () => { - return createWorker() + return worker.start() }) ctx.on('before-command', ({ command, session }) => { @@ -136,7 +75,7 @@ export function apply(ctx: Context, config: Config = {}) { UserTrap.attach(cmd, config.userFields, async ({ session, options, user, writable }, expr) => { if (options.restart) { - await session.$app.evalWorker.terminate() + await app.worker.restart() return '子线程已重启。' } @@ -149,27 +88,26 @@ export function apply(ctx: Context, config: Config = {}) { const _resolve = (result?: string) => { clearTimeout(timer) - app.evalWorker.off('error', listener) session._isEval = false + dispose() resolve(result) } const timer = setTimeout(async () => { - await app.evalWorker.terminate() _resolve(!session._sendCount && '执行超时。') + app.worker.restart() }, config.timeout) - const listener = (error: Error) => { + const dispose = app.worker.onError((error: Error) => { let message = ERROR_CODES[error['code']] if (!message) { logger.warn(error) message = '执行过程中遇到错误。' } _resolve(message) - } + }) - app.evalWorker.on('error', listener) - app.evalRemote.eval({ + app.worker.remote.eval({ user, writable, sid: session.$uuid, diff --git a/packages/plugin-eval/src/main.ts b/packages/plugin-eval/src/main.ts index 6346d9f887..ad4ffc87d8 100644 --- a/packages/plugin-eval/src/main.ts +++ b/packages/plugin-eval/src/main.ts @@ -1,4 +1,32 @@ import { App, Command, CommandAction, ParsedArgv, User } from 'koishi-core' +import { Logger, pick } from 'koishi-utils' +import { resolve } from 'path' +import { WorkerAPI, WorkerConfig, WorkerData, Response } from './worker' +import { Worker, ResourceLimits } from 'worker_threads' +import { expose, Remote, wrap } from './transfer' + +declare module 'koishi-core/dist/context' { + interface EventMap { + 'worker/start'(): void | Promise + 'worker/ready'(response: Response): void + 'worker/exit'(): void + } +} + +const logger = new Logger('eval') + +export interface MainConfig { + prefix?: string + timeout?: number + maxLogs?: number + userFields?: Access + resourceLimits?: ResourceLimits + dataKeys?: (keyof WorkerData)[] +} + +export interface EvalConfig extends MainConfig, WorkerData {} + +export interface Config extends MainConfig, WorkerConfig {} interface TrappedArgv extends ParsedArgv { user: Partial @@ -77,7 +105,7 @@ export class MainAPI { async send(uuid: string, message: string) { const session = this.getSession(uuid) if (!session._sendCount) session._sendCount = 0 - if (this.app.evalConfig.maxLogs > session._sendCount++) { + if (this.app.worker.config.maxLogs > session._sendCount++) { return await session.$sendQueued(message) } } @@ -87,3 +115,60 @@ export class MainAPI { return UserTrap.set(session.$user, data) } } + +export const workerScript = `require(${JSON.stringify(resolve(__dirname, 'worker.js'))});` + +export class EvalWorker { + static restart = true + + private worker: Worker + private promise: Promise + + public local: MainAPI + public remote: Remote + + constructor(public app: App, public config: EvalConfig) { + this.local = new MainAPI(app) + } + + async start() { + await this.app.parallel('worker/start') + + this.worker = new Worker(workerScript, { + eval: true, + workerData: { + logLevels: Logger.levels, + ...pick(this.config, this.config.dataKeys), + }, + resourceLimits: this.config.resourceLimits, + }) + + expose(this.worker, this.local) + this.remote = wrap(this.worker) + + await this.remote.start().then((response) => { + this.app.emit('worker/ready', response) + logger.info('worker started') + + this.worker.on('exit', (code) => { + this.app.emit('worker/exit') + logger.info('exited with code', code) + if (EvalWorker.restart) this.promise = this.start() + }) + }) + } + + async restart() { + await this.worker.terminate() + await this.promise + } + + onError(listener: (error: Error) => void) { + this.worker.on('error', listener) + return () => this.worker.off('error', listener) + } +} + +process.on('beforeExit', () => { + EvalWorker.restart = false +}) diff --git a/packages/plugin-eval/tests/index.spec.ts b/packages/plugin-eval/tests/index.spec.ts index 71ba193881..a05ae8b5c0 100644 --- a/packages/plugin-eval/tests/index.spec.ts +++ b/packages/plugin-eval/tests/index.spec.ts @@ -1,18 +1,17 @@ -import { defineProperty } from 'koishi-utils' import { App } from 'koishi-test-utils' import { inspect } from 'util' import { resolve } from 'path' -import * as _eval from 'koishi-plugin-eval' +import * as pluginEval from 'koishi-plugin-eval' -defineProperty(_eval, 'workerScript', [ +require('koishi-plugin-eval/src/main').workerScript = [ 'require("ts-node/register/transpile-only");', 'require("tsconfig-paths/register");', `require(${JSON.stringify(resolve(__dirname, '../src/worker.ts'))})`, -].join('\n')) +].join('\n') const app = new App() -app.plugin(_eval, { +app.plugin(pluginEval, { setupFiles: { 'test-worker': resolve(__dirname, 'worker.ts'), },