refactor

1 files changed, 19 insertions, 1 deletions

diff --git a/cmd/dtailhealthcheck/main.go b/cmd/dtailhealthcheck/main.go
index e0cb795..71c162e 100644
--- a/cmd/dtailhealthcheck/main.go
+++ b/cmd/dtailhealthcheck/main.go
@@ -3,9 +3,14 @@ package main
 import (
 	"context"
 	"flag"
+	"fmt"
 	"os"
 	"sync"
 
+	"net/http"
+	_ "net/http"
+	_ "net/http/pprof"
+
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/dlog"
@@ -16,8 +21,14 @@ import (
 // The evil begins here.
 func main() {
 	var args config.Args
+	var pprof int
+
+	flag.IntVar(&pprof, "pprof", -1, "Start PProf server this port")
+	flag.StringVar(&args.Logger, "logger", config.DefaultHealthCheckLogger, "Logger name")
+	flag.StringVar(&args.LogLevel, "logLevel", "", "Log level")
 	flag.StringVar(&args.ServersStr, "server", "", "Remote server to connect")
 	flag.Parse()
+
 	config.Setup(source.HealthCheck, &args, flag.Args())
 
 	ctx, cancel := context.WithCancel(context.Background())
@@ -25,7 +36,14 @@ func main() {
 	var wg sync.WaitGroup
 	wg.Add(1)
 
-	dlog.Start(ctx, &wg, source.HealthCheck, config.Common.LogLevel)
+	if pprof > -1 {
+		// For debugging purposes only
+		pprofArgs := fmt.Sprintf("0.0.0.0:%d", pprof)
+		go http.ListenAndServe(pprofArgs, nil)
+		dlog.Client.Info("Started PProf", pprofArgs)
+	}
+
+	dlog.Start(ctx, &wg, source.HealthCheck)
 	healthClient, _ := clients.NewHealthClient(args)
 	os.Exit(healthClient.Start(ctx, signal.NoCh(ctx)))
 }