cherry pick #824 to release-4.0 (#839)

Signed-off-by: ti-srebot <[email protected]> Co-authored-by: 山岚 <[email protected]> Co-authored-by: Neil Shen <[email protected]>
pingcap · Mar 11, 2021 · 34c2240 · 34c2240
1 parent c8e15bf
commit 34c2240
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 39 deletions.
diff --git a/cmd/br/cmd.go b/cmd/br/cmd.go
@@ -11,6 +11,8 @@ import (
 	"sync/atomic"
 	"time"
 
+	tidbutils "github.com/pingcap/tidb-tools/pkg/utils"
+
 	"github.com/pingcap/errors"
 	"github.com/pingcap/log"
 	"github.com/pingcap/tidb/util/logutil"
@@ -150,22 +152,34 @@ func Init(cmd *cobra.Command) (err error) {
 			return
 		}
 
-		// Initialize the pprof server.
-		// TODO: Support TLS.
-		statusAddr, e := cmd.Flags().GetString(FlagStatusAddr)
-		if e != nil {
-			err = e
-			return
-		}
-		if statusAddr != "" {
-			utils.StartPProfListener(statusAddr)
-		} else {
-			utils.StartDynamicPProfListener()
-		}
+		err = startPProf(cmd)
 	})
 	return errors.Trace(err)
 }
 
+func startPProf(cmd *cobra.Command) error {
+	// Initialize the pprof server.
+	statusAddr, err := cmd.Flags().GetString(FlagStatusAddr)
+	if err != nil {
+		return errors.Trace(err)
+	}
+	ca, cert, key, err := task.ParseTLSTripleFromFlags(cmd.Flags())
+	if err != nil {
+		return errors.Trace(err)
+	}
+	// Host isn't used here.
+	tls, err := tidbutils.NewTLS(ca, cert, key, "localhost", nil)
+	if err != nil {
+		return errors.Trace(err)
+	}
+
+	if statusAddr != "" {
+		return utils.StartPProfListener(statusAddr, tls)
+	}
+	utils.StartDynamicPProfListener(tls)
+	return nil
+}
+
 // HasLogFile returns whether we set a log file.
 func HasLogFile() bool {
 	return atomic.LoadUint64(&hasLogFile) != uint64(0)

diff --git a/pkg/task/common.go b/pkg/task/common.go
@@ -195,19 +195,25 @@ func DefineFilterFlags(command *cobra.Command) {
 // ParseFromFlags parses the TLS config from the flag set.
 func (tls *TLSConfig) ParseFromFlags(flags *pflag.FlagSet) error {
 	var err error
-	tls.CA, err = flags.GetString(flagCA)
+	tls.CA, tls.Cert, tls.Key, err = ParseTLSTripleFromFlags(flags)
+	return err
+}
+
+// ParseTLSTripleFromFlags parses the (ca, cert, key) triple from flags.
+func ParseTLSTripleFromFlags(flags *pflag.FlagSet) (ca, cert, key string, err error) {
+	ca, err = flags.GetString(flagCA)
 	if err != nil {
-		return errors.Trace(err)
+		return
 	}
-	tls.Cert, err = flags.GetString(flagCert)
+	cert, err = flags.GetString(flagCert)
 	if err != nil {
-		return errors.Trace(err)
+		return
 	}
-	tls.Key, err = flags.GetString(flagKey)
+	key, err = flags.GetString(flagKey)
 	if err != nil {
-		return errors.Trace(err)
+		return
 	}
-	return nil
+	return
 }
 
 func (cfg *Config) normalizePDURLs() error {

diff --git a/pkg/utils/dyn_pprof_unix.go b/pkg/utils/dyn_pprof_unix.go
@@ -8,25 +8,27 @@ import (
 	"os/signal"
 	"syscall"
 
+	tidbutils "github.com/pingcap/tidb-tools/pkg/utils"
+
 	"github.com/pingcap/log"
 	"go.uber.org/zap"
 )
 
 const startPProfSignal = syscall.SIGUSR1
 
-var signalChan = make(chan os.Signal, 1)
-
 // StartDynamicPProfListener starts the listener that will enable pprof when received `startPProfSignal`.
-func StartDynamicPProfListener() {
+func StartDynamicPProfListener(tls *tidbutils.TLS) {
+	signalChan := make(chan os.Signal, 1)
 	signal.Notify(signalChan, startPProfSignal)
-	go onSignalStartPProf(signalChan)
-}
-
-func onSignalStartPProf(signals <-chan os.Signal) {
-	for sig := range signals {
-		if sig == startPProfSignal {
-			log.Info("signal received, starting pprof...", zap.Stringer("signal", sig))
-			StartPProfListener("0.0.0.0:0")
+	go func() {
+		for sig := range signalChan {
+			if sig == startPProfSignal {
+				log.Info("signal received, starting pprof...", zap.Stringer("signal", sig))
+				if err := StartPProfListener("0.0.0.0:0", tls); err != nil {
+					log.Warn("failed to start pprof", zap.Error(err))
+					return
+				}
+			}
 		}
-	}
+	}()
 }
diff --git a/pkg/utils/pprof.go b/pkg/utils/pprof.go
@@ -7,11 +7,17 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"sync"
+
+	tidbutils "github.com/pingcap/tidb-tools/pkg/utils"
+
+	berrors "github.com/pingcap/br/pkg/errors"
+
+	"github.com/pingcap/errors"
 
 	// #nosec
 	// register HTTP handler for /debug/pprof
 	_ "net/http/pprof"
-	"sync"
 
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/log"
@@ -23,13 +29,12 @@ var (
 	mu           sync.Mutex
 )
 
-// StartPProfListener forks a new goroutine listening on specified port and provide pprof info.
-func StartPProfListener(statusAddr string) {
+func listen(statusAddr string) (net.Listener, error) {
 	mu.Lock()
 	defer mu.Unlock()
 	if startedPProf != "" {
 		log.Warn("Try to start pprof when it has been started, nothing will happen", zap.String("address", startedPProf))
-		return
+		return nil, errors.Annotate(berrors.ErrUnknown, "try to start pprof when it has been started at "+startedPProf)
 	}
 	failpoint.Inject("determined-pprof-port", func(v failpoint.Value) {
 		port := v.(int)
@@ -39,19 +44,29 @@ func StartPProfListener(statusAddr string) {
 	listener, err := net.Listen("tcp", statusAddr)
 	if err != nil {
 		log.Warn("failed to start pprof", zap.String("addr", statusAddr), zap.Error(err))
-		return
+		return nil, errors.Trace(err)
 	}
 	startedPProf = listener.Addr().String()
 	log.Info("bound pprof to addr", zap.String("addr", startedPProf))
 	_, _ = fmt.Fprintf(os.Stderr, "bound pprof to addr %s\n", startedPProf)
+	return listener, nil
+}
+
+// StartPProfListener forks a new goroutine listening on specified port and provide pprof info.
+func StartPProfListener(statusAddr string, wrapper *tidbutils.TLS) error {
+	listener, err := listen(statusAddr)
+	if err != nil {
+		return err
+	}
 
 	go func() {
-		if e := http.Serve(listener, nil); e != nil {
+		if e := http.Serve(wrapper.WrapListener(listener), nil); e != nil {
 			log.Warn("failed to serve pprof", zap.String("addr", startedPProf), zap.Error(e))
 			mu.Lock()
 			startedPProf = ""
 			mu.Unlock()
 			return
 		}
 	}()
+	return nil
 }
diff --git a/tests/br_other/run.sh b/tests/br_other/run.sh
@@ -84,8 +84,7 @@ echo "starting pprof..."
 
 # give the former backup some time to write down lock file (and start pprof server).
 sleep 1
-# TODO Support TLS.
-run_curl "http://localhost:$PPROF_PORT/debug/pprof/trace?seconds=1" &>/dev/null
+run_curl "https://localhost:$PPROF_PORT/debug/pprof/trace?seconds=1" &>/dev/null
 echo "pprof started..."
 
 run_curl https://$PD_ADDR/pd/api/v1/config/schedule | grep '"disable": false'