Add prof_tags for annotating CPU and goroutine profiles with accoun…

…t/stream/consumer info Signed-off-by: Neil Twigg <neil@nats.io>
nats-io · Jun 5, 2023 · b5ed7d9 · b5ed7d9
1 parent af318be
commit b5ed7d9
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 9 deletions.
diff --git a/server/jetstream_cluster.go b/server/jetstream_cluster.go
@@ -833,7 +833,11 @@ func (js *jetStream) setupMetaGroup() error {
 	atomic.StoreInt32(&js.clustered, 1)
 	c.registerWithAccount(sacc)
 
-	js.srv.startGoRoutine(js.monitorCluster)
+	js.srv.startGoRoutine(
+		js.monitorCluster,
+		"type", "metaleader",
+		"account", sacc.Name,
+	)
 	return nil
 }
 
@@ -3309,7 +3313,12 @@ func (js *jetStream) processClusterUpdateStream(acc *Account, osa, sa *streamAss
 			}
 			mset.monitorWg.Add(1)
 			// Start monitoring..
-			s.startGoRoutine(func() { js.monitorStream(mset, sa, needsNode) })
+			s.startGoRoutine(
+				func() { js.monitorStream(mset, sa, needsNode) },
+				"type", "stream",
+				"account", mset.accName(),
+				"stream", mset.name(),
+			)
 		} else if numReplicas == 1 && alreadyRunning {
 			// We downgraded to R1. Make sure we cleanup the raft node and the stream monitor.
 			mset.removeNode()
@@ -3532,7 +3541,12 @@ func (js *jetStream) processClusterCreateStream(acc *Account, sa *streamAssignme
 			if mset != nil {
 				mset.monitorWg.Add(1)
 			}
-			s.startGoRoutine(func() { js.monitorStream(mset, sa, false) })
+			s.startGoRoutine(
+				func() { js.monitorStream(mset, sa, false) },
+				"type", "stream",
+				"account", mset.accName(),
+				"stream", mset.name(),
+			)
 		}
 	} else {
 		// Single replica stream, process manually here.
@@ -4138,7 +4152,13 @@ func (js *jetStream) processClusterCreateConsumer(ca *consumerAssignment, state
 			// Start our monitoring routine if needed.
 			if !alreadyRunning && !o.isMonitorRunning() {
 				o.monitorWg.Add(1)
-				s.startGoRoutine(func() { js.monitorConsumer(o, ca) })
+				s.startGoRoutine(
+					func() { js.monitorConsumer(o, ca) },
+					"type", "consumer",
+					"account", mset.accName(),
+					"stream", mset.name(),
+					"consumer", ca.Name,
+				)
 			}
 			// For existing consumer, only send response if not recovering.
 			if wasExisting && !js.isMetaRecovering() {

diff --git a/server/opts.go b/server/opts.go
@@ -303,6 +303,7 @@ type Options struct {
 	Websocket             WebsocketOpts     `json:"-"`
 	MQTT                  MQTTOpts          `json:"-"`
 	ProfPort              int               `json:"-"`
+	ProfTags              bool              `json:"-"`
 	PidFile               string            `json:"-"`
 	PortsFileDir          string            `json:"-"`
 	LogFile               string            `json:"-"`
@@ -986,6 +987,8 @@ func (o *Options) processConfigFileLine(k string, v interface{}, errors *[]error
 		o.PortsFileDir = v.(string)
 	case "prof_port":
 		o.ProfPort = int(v.(int64))
+	case "prof_tags":
+		o.ProfTags = v.(bool)
 	case "max_control_line":
 		if v.(int64) > 1<<31-1 {
 			err := &configErr{tk, fmt.Sprintf("%s value is too big", k)}

diff --git a/server/server.go b/server/server.go
@@ -28,6 +28,7 @@ import (
 	"net"
 	"net/http"
 	"regexp"
+	"runtime/pprof"
 
 	// Allow dynamic profiling.
 	_ "net/http/pprof"
@@ -190,6 +191,7 @@ type Server struct {
 	grTmpClients map[uint64]*client
 	grRunning    bool
 	grWG         sync.WaitGroup // to wait on various go routines
+	grTags       bool
 
 	cproto     int64     // number of clients supporting async INFO
 	configTime time.Time // last time config was loaded
@@ -2050,6 +2052,7 @@ func (s *Server) Start() {
 
 	s.grMu.Lock()
 	s.grRunning = true
+	s.grTags = opts.ProfTags
 	s.grMu.Unlock()
 
 	s.startRateLimitLogExpiration()
@@ -3501,15 +3504,22 @@ func (s *Server) String() string {
 	return s.info.Name
 }
 
-func (s *Server) startGoRoutine(f func()) bool {
+func (s *Server) startGoRoutine(f func(), labels ...string) bool {
 	var started bool
 	s.grMu.Lock()
+	defer s.grMu.Unlock()
 	if s.grRunning {
 		s.grWG.Add(1)
-		go f()
+		go func(applyLabels bool) {
+			if applyLabels {
+				pprof.SetGoroutineLabels(
+					pprof.WithLabels(context.Background(), pprof.Labels(labels...)),
+				)
+			}
+			f()
+		}(s.grTags && len(labels) > 0)
 		started = true
 	}
-	s.grMu.Unlock()
 	return started
 }
 

diff --git a/server/stream.go b/server/stream.go
@@ -2354,7 +2354,12 @@ func (mset *stream) setupMirrorConsumer() error {
 				mirror.qch = make(chan struct{})
 				mirror.wg.Add(1)
 				ready.Add(1)
-				if !mset.srv.startGoRoutine(func() { mset.processMirrorMsgs(mirror, &ready) }) {
+				if !mset.srv.startGoRoutine(
+					func() { mset.processMirrorMsgs(mirror, &ready) },
+					"type", "mirror",
+					"account", mset.acc.Name,
+					"stream", mset.cfg.Name,
+				) {
 					ready.Done()
 				}
 			}
@@ -2652,7 +2657,12 @@ func (mset *stream) setSourceConsumer(iname string, seq uint64, startTime time.T
 					si.qch = make(chan struct{})
 					si.wg.Add(1)
 					ready.Add(1)
-					if !mset.srv.startGoRoutine(func() { mset.processSourceMsgs(si, &ready) }) {
+					if !mset.srv.startGoRoutine(
+						func() { mset.processSourceMsgs(si, &ready) },
+						"type", "source",
+						"account", mset.acc.Name,
+						"stream", mset.cfg.Name,
+					) {
 						ready.Done()
 					}
 				}