Просмотр исходного кода

feat(cli): add native CPU profiling (#42862)

Dax 17 часов назад
Родитель
Сommit
b080f216a5

+ 5 - 2
packages/cli/src/commands/commands.ts

@@ -1,5 +1,6 @@
-import { Argument, Flag } from "effect/unstable/cli"
+import { Argument, Command, Flag } from "effect/unstable/cli"
 import { Spec } from "../framework/spec"
+import { GlobalFlags } from "./global-flags"
 
 declare const OPENCODE_CLI_NAME: string | undefined
 
@@ -26,7 +27,7 @@ const PermissionParams = {
   ),
 }
 
-export const Commands = Spec.make(typeof OPENCODE_CLI_NAME === "string" ? OPENCODE_CLI_NAME : "opencode", {
+const Root = Spec.make(typeof OPENCODE_CLI_NAME === "string" ? OPENCODE_CLI_NAME : "opencode", {
   description: "OpenCode 2.0 preview command line interface",
   params: {
     ...ServerParams,
@@ -277,3 +278,5 @@ export const Commands = Spec.make(typeof OPENCODE_CLI_NAME === "string" ? OPENCO
     }),
   ],
 })
+
+export const Commands = { ...Root, spec: Root.spec.pipe(Command.withGlobalFlags(GlobalFlags.all)) }

+ 12 - 0
packages/cli/src/commands/global-flags.ts

@@ -0,0 +1,12 @@
+export * as GlobalFlags from "./global-flags"
+
+import { Flag, GlobalFlag } from "effect/unstable/cli"
+
+export const CpuProfile = GlobalFlag.setting("cpu-profile")({
+  flag: Flag.string("cpu-profile").pipe(
+    Flag.withDescription("Write a CPU profile to this path when the process stops"),
+    Flag.optional,
+  ),
+})
+
+export const all = [CpuProfile] as const

+ 45 - 0
packages/cli/src/cpu-profile.ts

@@ -0,0 +1,45 @@
+export * as CpuProfile from "./cpu-profile"
+
+import { Effect, FileSystem } from "effect"
+import { Session } from "node:inspector"
+import path from "node:path"
+
+export function run<A, E, R>(file: string, effect: Effect.Effect<A, E, R>) {
+  const target = path.resolve(file)
+  return Effect.acquireUseRelease(
+    Effect.gen(function* () {
+      const fs = yield* FileSystem.FileSystem
+      yield* fs.makeDirectory(path.dirname(target), { recursive: true })
+      const session = new Session()
+      session.connect()
+      yield* command(session, "Profiler.enable")
+      yield* command(session, "Profiler.start")
+      yield* Effect.logInfo("CPU profile started", { path: target })
+      return session
+    }),
+    () => effect,
+    (session) =>
+      Effect.tryPromise(
+        () =>
+          new Promise<void>((resolve, reject) => {
+            session.post("Profiler.stop", (error, result) => {
+              session.disconnect()
+              if (error) return reject(error)
+              Bun.write(target, JSON.stringify(result.profile)).then(() => resolve(), reject)
+            })
+          }),
+      ).pipe(
+        Effect.andThen(Effect.logInfo("CPU profile written", { path: target })),
+        Effect.catchCause((cause) => Effect.logError("Failed to write CPU profile", { path: target, cause })),
+      ),
+  )
+}
+
+function command(session: Session, method: "Profiler.enable" | "Profiler.start") {
+  return Effect.tryPromise(
+    () =>
+      new Promise<void>((resolve, reject) => {
+        session.post(method, (error) => (error ? reject(error) : resolve()))
+      }),
+  )
+}

+ 18 - 2
packages/cli/src/framework/runtime.ts

@@ -1,10 +1,13 @@
-import { Effect, FileSystem, Scope } from "effect"
+import { Effect, FileSystem, Option, Scope } from "effect"
 import { Command } from "effect/unstable/cli"
 import { Spec } from "./spec"
 import { Global } from "@opencode-ai/util/global"
 import { Updater } from "../services/updater"
 import { Config } from "../config"
 import { Npm } from "@opencode-ai/util/npm"
+import { GlobalFlags } from "../commands/global-flags"
+import { CpuProfile } from "../cpu-profile"
+import path from "node:path"
 
 export type Input<Value> =
   Value extends Spec.Node<infer _Name, infer Command, infer _Commands>
@@ -86,7 +89,20 @@ function provide(node: Spec.Any, handlers: ReadonlyArray<LazyHandler>): Provided
     ? node.spec.pipe(
         Command.withHandler((input) =>
           Effect.gen(function* () {
-            yield* Effect.flatMap(Effect.promise(handler.load), (module) => module.default(input))
+            const module = yield* Effect.promise(handler.load)
+            const cpuProfile = Option.getOrUndefined(yield* GlobalFlags.CpuProfile)
+            if (!cpuProfile) return yield* module.default(input)
+            const target = path.resolve(cpuProfile)
+            const previous = process.env.OPENCODE_CPU_PROFILE
+            process.env.OPENCODE_CPU_PROFILE = target
+            return yield* (node.name === "serve" ? CpuProfile.run(target, module.default(input)) : module.default(input)).pipe(
+              Effect.ensuring(
+                Effect.sync(() => {
+                  if (previous === undefined) delete process.env.OPENCODE_CPU_PROFILE
+                  else process.env.OPENCODE_CPU_PROFILE = previous
+                }),
+              ),
+            )
           }),
         ),
       )

+ 6 - 1
packages/cli/src/services/service-config.ts

@@ -104,7 +104,12 @@ export const options = Effect.fnUntraced(function* (input: { readonly checkVersi
   return {
     file,
     version: input.checkVersion ? OPENCODE_VERSION : undefined,
-    command: [...selfCommand(), "serve", "--service"],
+    command: [
+      ...selfCommand(),
+      "serve",
+      "--service",
+      ...(process.env.OPENCODE_CPU_PROFILE ? ["--cpu-profile", process.env.OPENCODE_CPU_PROFILE] : []),
+    ],
   }
 })
 

+ 23 - 0
packages/cli/test/service.test.ts

@@ -19,6 +19,29 @@ test("managed service ports are stable per installation channel", () => {
   expect(ServiceConfig.defaultPort("preview-a")).not.toBe(ServiceConfig.defaultPort("preview-b"))
 })
 
+test("managed service forwards the CPU profile path to the server", async () => {
+  const root = await fs.mkdtemp(path.join(os.tmpdir(), "opencode-service-profile-"))
+  const profile = path.join(root, "server.cpuprofile")
+  try {
+    const previous = process.env.OPENCODE_CPU_PROFILE
+    process.env.OPENCODE_CPU_PROFILE = profile
+    try {
+      const options = await Effect.runPromise(
+        ServiceConfig.options().pipe(
+          Effect.provide(Global.layerWith({ config: path.join(root, "config"), state: path.join(root, "state") })),
+          Effect.provide(NodeFileSystem.layer),
+        ),
+      )
+      expect(options.command.slice(-2)).toEqual(["--cpu-profile", profile])
+    } finally {
+      if (previous === undefined) delete process.env.OPENCODE_CPU_PROFILE
+      else process.env.OPENCODE_CPU_PROFILE = previous
+    }
+  } finally {
+    await fs.rm(root, { recursive: true, force: true })
+  }
+})
+
 test("local channel stores service config with the local service filename", async () => {
   const root = await fs.mkdtemp(path.join(os.tmpdir(), "opencode-service-"))
   try {