feat(eval): worker class with onError & restart methods

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,22 +42,22 @@ 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)
 
   const addon = ctx.command('addon', '扩展功能')
     .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<string, Promise<Manifest>>
-  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
         })
 

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<string, Session>
-    evalConfig: EvalConfig
-    evalWorker: Worker
-    evalRemote: Remote<WorkerAPI>
+    worker: EvalWorker
   }
 }
 
@@ -23,14 +17,6 @@ declare module 'koishi-core/dist/command' {
   }
 }
 
-declare module 'koishi-core/dist/context' {
-  interface EventMap {
-    'worker/start'(): void | Promise<void>
-    '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<User.Field>
-  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,

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<void>
+    'worker/ready'(response: Response): void
+    'worker/exit'(): void
+  }
+}
+
+const logger = new Logger('eval')
+
+export interface MainConfig {
+  prefix?: string
+  timeout?: number
+  maxLogs?: number
+  userFields?: Access<User.Field>
+  resourceLimits?: ResourceLimits
+  dataKeys?: (keyof WorkerData)[]
+}
+
+export interface EvalConfig extends MainConfig, WorkerData {}
+
+export interface Config extends MainConfig, WorkerConfig {}
 
 interface TrappedArgv<O> extends ParsedArgv<never, never, O> {
   user: Partial<User>
@@ -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<void>
+
+  public local: MainAPI
+  public remote: Remote<WorkerAPI>
+
+  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
+})

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'),
   },