feat: add /debug/pprof endpoint when using --pprof flag

cmd/waku/flags.go

@@ -427,4 +427,9 @@ var (
 		Usage:       "Enable access to REST HTTP Private API",
 		Destination: &options.RESTServer.Private,
 	}
+	PProf = &cli.BoolFlag{
+		Name:        "pprof",
+		Usage:       "provides runtime profiling data at /debug/pprof in both REST and RPC servers if they're enabled",
+		Destination: &options.PProf,
+	}
 )

cmd/waku/main.go

@@ -87,6 +87,7 @@ func main() {
 		RESTRelayCacheCapacity,
 		RESTAdmin,
 		RESTPrivate,
+		PProf,
 	}
 
 	rlnFlags := rlnFlags()

waku/node.go

@@ -360,7 +360,7 @@ func Execute(options Options) {
 
 	var rpcServer *rpc.WakuRpc
 	if options.RPCServer.Enable {
-		rpcServer = rpc.NewWakuRpc(wakuNode, options.RPCServer.Address, options.RPCServer.Port, options.RPCServer.Admin, options.RPCServer.Private, options.RPCServer.RelayCacheCapacity, logger)
+		rpcServer = rpc.NewWakuRpc(wakuNode, options.RPCServer.Address, options.RPCServer.Port, options.RPCServer.Admin, options.RPCServer.Private, options.PProf, options.RPCServer.RelayCacheCapacity, logger)
 		rpcServer.Start()
 	}
 
@@ -369,7 +369,7 @@ func Execute(options Options) {
 	var restServer *rest.WakuRest
 	if options.RESTServer.Enable {
 		wg.Add(1)
-		restServer = rest.NewWakuRest(wakuNode, options.RESTServer.Address, options.RESTServer.Port, options.RESTServer.Admin, options.RESTServer.Private, options.RESTServer.RelayCacheCapacity, logger)
+		restServer = rest.NewWakuRest(wakuNode, options.RESTServer.Address, options.RESTServer.Port, options.RESTServer.Admin, options.RESTServer.Private, options.PProf, options.RESTServer.RelayCacheCapacity, logger)
 		restServer.Start(ctx, &wg)
 	}
 

waku/options.go

@@ -159,6 +159,7 @@ type Options struct {
 	NAT              string
 	PersistPeers     bool
 	UserAgent        string
+	PProf            bool
 
 	PeerExchange PeerExchangeOptions
 	Websocket    WSOptions

waku/v2/protocol/relay/waku_relay.go

@@ -247,7 +247,7 @@ func (w *WakuRelay) SubscribeToTopic(ctx context.Context, topic string) (*Subscr
 		w.bcaster.Register(&topic, subscription.C)
 	}
 
-	go w.subscribeToTopic(topic, subscription, sub)
+	go w.subscribeToTopic(ctx, topic, subscription, sub)
 
 	return subscription, nil
 }
@@ -307,8 +307,8 @@ func (w *WakuRelay) nextMessage(ctx context.Context, sub *pubsub.Subscription) <
 	return msgChannel
 }
 
-func (w *WakuRelay) subscribeToTopic(t string, subscription *Subscription, sub *pubsub.Subscription) {
+func (w *WakuRelay) subscribeToTopic(ctx context.Context, t string, subscription *Subscription, sub *pubsub.Subscription) {
-	ctx, err := tag.New(context.Background(), tag.Insert(metrics.KeyType, "relay"))
+	ctx, err := tag.New(ctx, tag.Insert(metrics.KeyType, "relay"))
 	if err != nil {
 		w.log.Error("creating tag map", zap.Error(err))
 		return

waku/v2/rest/waku_rest.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/http/pprof"
 	"sync"
 
 	"github.com/gorilla/mux"
@@ -20,12 +21,17 @@ type WakuRest struct {
 	relayService *RelayService
 }
 
-func NewWakuRest(node *node.WakuNode, address string, port int, enableAdmin bool, enablePrivate bool, relayCacheCapacity int, log *zap.Logger) *WakuRest {
+func NewWakuRest(node *node.WakuNode, address string, port int, enableAdmin bool, enablePrivate bool, enablePProf bool, relayCacheCapacity int, log *zap.Logger) *WakuRest {
 	wrpc := new(WakuRest)
 	wrpc.log = log.Named("rest")
 
 	mux := mux.NewRouter()
 
+	if enablePProf {
+		mux.PathPrefix("/debug/").Handler(http.DefaultServeMux)
+		mux.HandleFunc("/debug/pprof/", pprof.Index)
+	}
+
 	_ = NewDebugService(node, mux)
 	relayService := NewRelayService(node, mux, relayCacheCapacity, log)
 

waku/v2/rest/waku_rest_test.go

@@ -14,7 +14,7 @@ func TestWakuRest(t *testing.T) {
 	n, err := node.New(context.Background(), options)
 	require.NoError(t, err)
 
-	rpc := NewWakuRest(n, "127.0.0.1", 8080, true, true, 10, utils.Logger())
+	rpc := NewWakuRest(n, "127.0.0.1", 8080, true, true, false, 10, utils.Logger())
 	require.NotNil(t, rpc.server)
 	require.Equal(t, rpc.server.Addr, "127.0.0.1:8080")
 }

waku/v2/rpc/waku_rpc.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"net/http/pprof"
 	"time"
 
 	"github.com/gorilla/mux"
@@ -24,7 +25,7 @@ type WakuRpc struct {
 	adminService   *AdminService
 }
 
-func NewWakuRpc(node *node.WakuNode, address string, port int, enableAdmin bool, enablePrivate bool, cacheCapacity int, log *zap.Logger) *WakuRpc {
+func NewWakuRpc(node *node.WakuNode, address string, port int, enableAdmin bool, enablePrivate bool, enablePProf bool, cacheCapacity int, log *zap.Logger) *WakuRpc {
 	wrpc := new(WakuRpc)
 	wrpc.log = log.Named("rpc")
 
@@ -39,6 +40,11 @@ func NewWakuRpc(node *node.WakuNode, address string, port int, enableAdmin bool,
 		wrpc.log.Info("served request", zap.String("path", r.URL.Path), zap.Duration("duration", time.Since(t)))
 	})
 
+	if enablePProf {
+		mux.PathPrefix("/debug/").Handler(http.DefaultServeMux)
+		mux.HandleFunc("/debug/pprof/", pprof.Index)
+	}
+
 	debugService := NewDebugService(node)
 	err := s.RegisterService(debugService, "Debug")
 	if err != nil {
@@ -109,6 +115,7 @@ func (r *WakuRpc) Start() {
 	if r.privateService != nil {
 		go r.privateService.Start()
 	}
+
 	go func() {
 		_ = r.server.ListenAndServe()
 	}()

waku/v2/rpc/waku_rpc_test.go

@@ -14,7 +14,7 @@ func TestWakuRpc(t *testing.T) {
 	n, err := node.New(context.Background(), options)
 	require.NoError(t, err)
 
-	rpc := NewWakuRpc(n, "127.0.0.1", 8080, true, true, 30, utils.Logger())
+	rpc := NewWakuRpc(n, "127.0.0.1", 8080, true, true, false, 30, utils.Logger())
 	require.NotNil(t, rpc.server)
 	require.Equal(t, rpc.server.Addr, "127.0.0.1:8080")
 }