more on this

1 files changed, 11 insertions, 7 deletions

diff --git a/cmd/dtail/main.go b/cmd/dtail/main.go
index 3746af9..820323c 100644
--- a/cmd/dtail/main.go
+++ b/cmd/dtail/main.go
@@ -66,7 +66,11 @@ func main() {
 	if grep != "" {
 		args.RegexStr = grep
 	}
-	config.Setup(source.Server, &args, flag.Args())
+	sourceProcess := source.Client
+	if checkHealth {
+		sourceProcess = source.HealthCheck
+	}
+	config.Setup(sourceProcess, &args, flag.Args())
 
 	if displayVersion {
 		version.PrintAndExit()
@@ -88,16 +92,16 @@ func main() {
 		defer cancel()
 	}
 
+	var wg sync.WaitGroup
+	wg.Add(1)
+	dlog.Start(ctx, &wg, sourceProcess, config.Common.LogLevel)
+
 	if checkHealth {
-		healthClient, _ := clients.NewHealthClient(omode.HealthClient)
+		healthClient, _ := clients.NewHealthClient(args)
 		cancel()
-		os.Exit(healthClient.Start(ctx))
+		os.Exit(healthClient.Start(ctx, signal.InterruptCh(ctx)))
 	}
 
-	var wg sync.WaitGroup
-	wg.Add(1)
-	dlog.Start(ctx, &wg, source.Client, config.Common.LogLevel)
-
 	if pprof > -1 {
 		// For debugging purposes only
 		pprofArgs := fmt.Sprintf("0.0.0.0:%d", pprof)