Add support for nats server request profile (nats-io/nats-server#3774)

cli/server_request_command.go

@@ -14,7 +14,11 @@
 package cli
 
 import (
+	"encoding/json"
 	"fmt"
+	"os"
+	"path/filepath"
+	"time"
 
 	"github.com/choria-io/fisk"
 	"github.com/nats-io/nats-server/v2/server"
@@ -50,6 +54,10 @@ type SrvRequestCmd struct {
 
 	jsServerOnly bool
 	jsEnabled    bool
+
+	profileName  string
+	profileDebug int
+	profileDir   string
 }
 
 func configureServerRequestCommand(srv *fisk.CmdClause) {
@@ -113,6 +121,11 @@ func configureServerRequestCommand(srv *fisk.CmdClause) {
 	healthz := req.Command("jetstream-health", "Request JetStream health status").Alias("healthz").Action(c.healthz)
 	healthz.Flag("js-enabled", "Checks that JetStream should be enabled on all servers").Short('J').BoolVar(&c.jsEnabled)
 	healthz.Flag("server-only", "Restricts the health check to the JetStream server only, do not check streams and consumers").Short('S').BoolVar(&c.jsServerOnly)
+
+	profilez := req.Command("profile", "Run a profile").Action(c.profilez)
+	profilez.Arg("profile", "Specify the name of the profile to run (allocs, heap, goroutine, mutex, threadcreate, block)").StringVar(&c.profileName)
+	profilez.Arg("dir", "Set the output directory for profile files").Default(".").ExistingDirVar(&c.profileDir)
+	profilez.Flag("debug", "Set the debug level of the profile").IntVar(&c.profileDebug)
 }
 
 func (c *SrvRequestCmd) healthz(_ *fisk.ParseContext) error {
@@ -138,6 +151,72 @@ func (c *SrvRequestCmd) healthz(_ *fisk.ParseContext) error {
 	return nil
 }
 
+type profilezResponse struct {
+	Server server.ServerInfo     `json:"server"`
+	Resp   server.ProfilezStatus `json:"data"`
+}
+
+func (c *SrvRequestCmd) profilez(_ *fisk.ParseContext) error {
+	nc, _, err := prepareHelper("", natsOpts()...)
+	if err != nil {
+		return err
+	}
+
+	opts := server.ProfilezEventOptions{
+		ProfilezOptions: server.ProfilezOptions{
+			Name:  c.profileName,
+			Debug: c.profileDebug,
+		},
+		EventFilterOptions: c.reqFilter(),
+	}
+
+	res, err := c.doReq("PROFILEZ", &opts, nc)
+	if err != nil {
+		return err
+	}
+
+	prefix := fmt.Sprintf("%s-%s-", c.profileName, time.Now().Format("20060102-150405"))
+	prefix = filepath.Join(c.profileDir, prefix)
+
+	for _, r := range res {
+		var resp profilezResponse
+		if err := json.Unmarshal(r, &resp); err != nil {
+			continue
+		}
+		if resp.Resp.Error != "" {
+			fmt.Fprintf(os.Stderr, "Server %q error: %s\n", resp.Server.Name, resp.Resp.Error)
+			continue
+		}
+
+		filename := prefix + resp.Server.Name
+		if err := c.profilezWrite(filename, &resp); err != nil {
+			fmt.Fprintf(os.Stderr, "Server %q error: %s\n", resp.Server.Name, err)
+		} else {
+			fmt.Fprintf(os.Stdout, "Server %q profile written: %s\n", resp.Server.Name, filename)
+		}
+	}
+
+	return nil
+}
+
+func (c *SrvRequestCmd) profilezWrite(filename string, resp *profilezResponse) error {
+	f, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, 0o644)
+	if err != nil {
+		return fmt.Errorf("failed to open file: %w", err)
+	}
+	defer f.Close()
+
+	n, err := f.Write(resp.Resp.Profile)
+	if err != nil {
+		return fmt.Errorf("failed to write file: %w", err)
+	}
+	if n != len(resp.Resp.Profile) {
+		return fmt.Errorf("short write")
+	}
+
+	return nil
+}
+
 func (c *SrvRequestCmd) jsz(_ *fisk.ParseContext) error {
 	nc, _, err := prepareHelper("", natsOpts()...)
 	if err != nil {
@@ -289,8 +368,8 @@ func (c *SrvRequestCmd) routez(_ *fisk.ParseContext) error {
 	}
 
 	return nil
-
 }
+
 func (c *SrvRequestCmd) conns(_ *fisk.ParseContext) error {
 	opts := &server.ConnzEventOptions{
 		ConnzOptions: server.ConnzOptions{
@@ -332,7 +411,6 @@ func (c *SrvRequestCmd) conns(_ *fisk.ParseContext) error {
 	}
 
 	return nil
-
 }
 
 func (c *SrvRequestCmd) varz(_ *fisk.ParseContext) error {
@@ -356,7 +434,6 @@ func (c *SrvRequestCmd) varz(_ *fisk.ParseContext) error {
 	}
 
 	return nil
-
 }
 
 func (c *SrvRequestCmd) subs(_ *fisk.ParseContext) error {

go.mod

@@ -16,7 +16,7 @@ require (
 	github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51
 	github.com/klauspost/compress v1.15.13
 	github.com/nats-io/jsm.go v0.0.36-0.20221218234554-efc05fda17d5
-	github.com/nats-io/nats-server/v2 v2.9.11-0.20230106144522-cf93b0535977
+	github.com/nats-io/nats-server/v2 v2.9.12-0.20230111160909-68953678bbd2
 	github.com/nats-io/nats.go v1.22.1
 	github.com/nats-io/nuid v1.0.1
 	github.com/prometheus/client_golang v1.14.0

go.sum

@@ -80,8 +80,8 @@ github.com/nats-io/jsm.go v0.0.36-0.20221218234554-efc05fda17d5 h1:c8/AsTy82nPvd
 github.com/nats-io/jsm.go v0.0.36-0.20221218234554-efc05fda17d5/go.mod h1:IuS2F1KLw5u6A+r8Fn/FWd+tY7+KG+OEmDoPrChp0pI=
 github.com/nats-io/jwt/v2 v2.3.1-0.20221227170542-bdf40fa3627b h1:exHeHbghpBp1JvdYq7muaKFvJgLD93UDcmoIbFu/9PA=
 github.com/nats-io/jwt/v2 v2.3.1-0.20221227170542-bdf40fa3627b/go.mod h1:DYujvzCMZzUuqB3i1Pnpf1YtkuTwhdI84Aah9wRXkK0=
-github.com/nats-io/nats-server/v2 v2.9.11-0.20230106144522-cf93b0535977 h1:srZ3zCeMoq7UoBBtYk+Jm8kklkt9wts+PZogDtCt8MI=
-github.com/nats-io/nats-server/v2 v2.9.11-0.20230106144522-cf93b0535977/go.mod h1:ibVHvIWZwqnarh51bnfR3zZWtlL3SjG9X49ocsfFUm4=
+github.com/nats-io/nats-server/v2 v2.9.12-0.20230111160909-68953678bbd2 h1:RgR+HViLAMt0cc6E6Wk0ZdVkX9vWvED7ZNitXs/Svrs=
+github.com/nats-io/nats-server/v2 v2.9.12-0.20230111160909-68953678bbd2/go.mod h1:ibVHvIWZwqnarh51bnfR3zZWtlL3SjG9X49ocsfFUm4=
 github.com/nats-io/nats.go v1.22.1 h1:XzfqDspY0RNufzdrB8c4hFR+R3dahkxlpWe5+IWJzbE=
 github.com/nats-io/nats.go v1.22.1/go.mod h1:tLqubohF7t4z3du1QDPYJIQQyhb4wl6DhjxEajSI7UA=
 github.com/nats-io/nkeys v0.3.0/go.mod h1:gvUNGjVcM2IPr5rCsRsC6Wb3Hr2CQAm08dsxtV6A5y4=