From 32624697fd2d74d4a6e23549b647d323b210fb2a Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 11 May 2020 18:21:39 +0100
Subject: [PATCH] Add PPROFLISTEN (#1019)

* Add PPROFLISTEN env var

* Direct logging to more useful places

* Space
---
 cmd/dendrite-room-server/main.go |  2 --
 common/basecomponent/base.go     |  3 +++
 common/log.go                    | 12 ++++++++++++
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/cmd/dendrite-room-server/main.go b/cmd/dendrite-room-server/main.go
index 41149ad9..17246844 100644
--- a/cmd/dendrite-room-server/main.go
+++ b/cmd/dendrite-room-server/main.go
@@ -15,8 +15,6 @@
 package main
 
 import (
-	_ "net/http/pprof"
-
 	"github.com/matrix-org/dendrite/common/basecomponent"
 	"github.com/matrix-org/dendrite/common/keydb"
 	"github.com/matrix-org/dendrite/roomserver"
diff --git a/common/basecomponent/base.go b/common/basecomponent/base.go
index a7e6736a..cb04a308 100644
--- a/common/basecomponent/base.go
+++ b/common/basecomponent/base.go
@@ -42,6 +42,8 @@ import (
 	federationSenderAPI "github.com/matrix-org/dendrite/federationsender/api"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
 	"github.com/sirupsen/logrus"
+
+	_ "net/http/pprof"
 )
 
 // BaseDendrite is a base for creating new instances of dendrite. It parses
@@ -71,6 +73,7 @@ const HTTPClientTimeout = time.Second * 30
 func NewBaseDendrite(cfg *config.Dendrite, componentName string) *BaseDendrite {
 	common.SetupStdLogging()
 	common.SetupHookLogging(cfg.Logging, componentName)
+	common.SetupPprof()
 
 	closer, err := cfg.SetupTracing("Dendrite" + componentName)
 	if err != nil {
diff --git a/common/log.go b/common/log.go
index 11339ada..60e96965 100644
--- a/common/log.go
+++ b/common/log.go
@@ -18,6 +18,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"path"
 	"path/filepath"
@@ -79,6 +80,17 @@ func callerPrettyfier(f *runtime.Frame) (string, string) {
 	return funcname, filename
 }
 
+// SetupPprof starts a pprof listener. We use the DefaultServeMux here because it is
+// simplest, and it gives us the freedom to run pprof on a separate port.
+func SetupPprof() {
+	if hostPort := os.Getenv("PPROFLISTEN"); hostPort != "" {
+		logrus.Warn("Starting pprof on ", hostPort)
+		go func() {
+			logrus.WithError(http.ListenAndServe(hostPort, nil)).Error("Failed to setup pprof listener")
+		}()
+	}
+}
+
 // SetupStdLogging configures the logging format to standard output. Typically, it is called when the config is not yet loaded.
 func SetupStdLogging() {
 	logrus.SetReportCaller(true)