Add CPU pprof flag

Related: #23
Modified from: taoky/ayano@b26dcc8

Author: taoky

main.go

@@ -5,7 +5,10 @@ import (
 	"flag"
 	"log"
 	"os"
+	"os/signal"
+	"runtime/pprof"
 	"strings"
+	"syscall"
 
 	"github.com/pelletier/go-toml/v2"
 )
@@ -35,8 +38,29 @@ func sshmuxServer(configFile string) (*Server, error) {
 
 func main() {
 	var configFile string
+	var pprofFile string
 	flag.StringVar(&configFile, "c", "/etc/sshmux/config.toml", "config file")
+	flag.StringVar(&pprofFile, "pprof", "", "write pprof data to file")
 	flag.Parse()
+	if pprofFile != "" {
+		f, err := os.Create(pprofFile)
+		if err != nil {
+			log.Fatal("failed to create pprof file: %w", err)
+		}
+		defer f.Close()
+		if err := pprof.StartCPUProfile(f); err != nil {
+			log.Fatal("failed to start pprof: %w", err)
+		}
+		defer pprof.StopCPUProfile()
+		// Handle SIGINT/SIGTERM
+		signalChan := make(chan os.Signal, 1)
+		signal.Notify(signalChan, syscall.SIGINT, syscall.SIGTERM)
+		go func() {
+			<-signalChan
+			pprof.StopCPUProfile()
+			os.Exit(0)
+		}()
+	}
 	sshmux, err := sshmuxServer(configFile)
 	if err != nil {
 		log.Fatal(err)