RPC: continuous profiling for Python/JS/Go subprocesses

rewrite-go/cmd/rpc/main.go

@@ -36,6 +36,7 @@ import (
 	"time"
 
 	"github.com/google/uuid"
+	"github.com/grafana/pyroscope-go"
 
 	goparser "github.com/openrewrite/rewrite/rewrite-go/pkg/parser"
 	"github.com/openrewrite/rewrite/rewrite-go/pkg/printer"
@@ -257,7 +258,34 @@ func parseFlags() serverConfig {
 	return cfg
 }
 
+// initPyroscope starts continuous profiling when PYROSCOPE_SERVER_ADDRESS is
+// set. Tags inherited via PYROSCOPE_TAGS (k=v,k=v) are forwarded verbatim; a
+// runtime=go tag is added so flame graphs in the shared modcli application
+// can be sliced by which RPC subprocess produced them.
+func initPyroscope() {
+	server := os.Getenv("PYROSCOPE_SERVER_ADDRESS")
+	if server == "" {
+		return
+	}
+	appName := os.Getenv("PYROSCOPE_APPLICATION_NAME")
+	if appName == "" {
+		appName = "modcli"
+	}
+	tags := map[string]string{"runtime": "go"}
+	for _, pair := range strings.Split(os.Getenv("PYROSCOPE_TAGS"), ",") {
+		if i := strings.Index(pair, "="); i > 0 {
+			tags[strings.TrimSpace(pair[:i])] = strings.TrimSpace(pair[i+1:])
+		}
+	}
+	_, _ = pyroscope.Start(pyroscope.Config{
+		ApplicationName: appName,
+		ServerAddress:   server,
+		Tags:            tags,
+	})
+}
+
 func main() {
+	initPyroscope()
 	cfg := parseFlags()
 	s := newServer(cfg)
 	s.logger.Println("Go RPC server starting...")

rewrite-go/go.mod

@@ -4,4 +4,9 @@ go 1.25.0
 
 require github.com/google/uuid v1.6.0
 
-require golang.org/x/mod v0.35.0 // indirect
+require (
+	github.com/grafana/pyroscope-go v1.2.8 // indirect
+	github.com/grafana/pyroscope-go/godeltaprof v0.1.9 // indirect
+	github.com/klauspost/compress v1.17.8 // indirect
+	golang.org/x/mod v0.35.0 // indirect
+)

rewrite-go/go.sum

@@ -1,4 +1,10 @@
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/grafana/pyroscope-go v1.2.8 h1:UvCwIhlx9DeV7F6TW/z8q1Mi4PIm3vuUJ2ZlCEvmA4M=
+github.com/grafana/pyroscope-go v1.2.8/go.mod h1:SSi59eQ1/zmKoY/BKwa5rSFsJaq+242Bcrr4wPix1g8=
+github.com/grafana/pyroscope-go/godeltaprof v0.1.9 h1:c1Us8i6eSmkW+Ez05d3co8kasnuOY813tbMN8i/a3Og=
+github.com/grafana/pyroscope-go/godeltaprof v0.1.9/go.mod h1:2+l7K7twW49Ct4wFluZD3tZ6e0SjanjcUUBPVD/UuGU=
+github.com/klauspost/compress v1.17.8 h1:YcnTYrq7MikUT7k0Yb5eceMmALQPYBW/Xltxn0NAMnU=
+github.com/klauspost/compress v1.17.8/go.mod h1:Di0epgTjJY877eYKx5yC51cX2A2Vl2ibi7bDH9ttBbw=
 golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM=
 golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU=

rewrite-javascript/rewrite/package.json

@@ -86,6 +86,7 @@
     "vitest": "^4.0.18"
   },
   "optionalDependencies": {
+    "@pyroscope/nodejs": "^0.4.0",
     "prettier": "^3.7.4"
   },
   "bin": {

rewrite-javascript/rewrite/src/rpc/server.ts

@@ -31,6 +31,33 @@ import "../javascript";
 // Not possible to set the stack size when executing from npx for security reasons
 require('v8').setFlagsFromString('--stack-size=8000');
 
+function initPyroscope(logger: rpc.Logger): void {
+    const server = process.env.PYROSCOPE_SERVER_ADDRESS;
+    if (!server) {
+        return;
+    }
+    let Pyroscope: any;
+    try {
+        Pyroscope = require('@pyroscope/nodejs');
+    } catch {
+        logger.warn('PYROSCOPE_SERVER_ADDRESS set but @pyroscope/nodejs not installed; profiling disabled');
+        return;
+    }
+    const tags: Record<string, string> = {runtime: 'node'};
+    for (const pair of (process.env.PYROSCOPE_TAGS || '').split(',')) {
+        const eq = pair.indexOf('=');
+        if (eq > 0) {
+            tags[pair.slice(0, eq).trim()] = pair.slice(eq + 1).trim();
+        }
+    }
+    Pyroscope.init({
+        appName: process.env.PYROSCOPE_APPLICATION_NAME || 'modcli',
+        serverAddress: server,
+        tags,
+    });
+    Pyroscope.start();
+}
+
 interface ProgramOptions {
     logFile?: string;
     metricsCsv?: string;
@@ -97,6 +124,8 @@ async function main() {
         log: (msg: string) => log && options.traceRpcMessages && log.write(`[js trace] ${msg}\n`)
     };
 
+    initPyroscope(logger);
+
     // Create the connection with the custom logger
     const connection = rpc.createMessageConnection(
         new rpc.StreamMessageReader(process.stdin),

rewrite-python/rewrite/pyproject.toml

@@ -41,6 +41,11 @@ publish = [
     "build>=1.0.0",
     "twine>=5.0.0",
 ]
+# Continuous profiling for production recipe-worker subprocesses; the RPC
+# server's _init_pyroscope() is a no-op when this isn't installed.
+profiling = [
+    "pyroscope-io>=0.8.0",
+]
 
 [project.urls]
 Homepage = "https://github.com/openrewrite/rewrite"