1 files changed, 10 insertions, 9 deletions
diff --git a/cmd/dmap/main.go b/cmd/dmap/main.go
index 83dad50..f3f706a 100644
--- a/cmd/dmap/main.go
+++ b/cmd/dmap/main.go
@@ -1,13 +1,15 @@
 package main
 
 import (
+	"context"
 	"flag"
+	"os"
 
-	"github.com/mimecast/dtail/internal/omode"
 	"github.com/mimecast/dtail/internal/clients"
 	"github.com/mimecast/dtail/internal/color"
 	"github.com/mimecast/dtail/internal/config"
-	"github.com/mimecast/dtail/internal/logger"
+	"github.com/mimecast/dtail/internal/io/logger"
+	"github.com/mimecast/dtail/internal/omode"
 	"github.com/mimecast/dtail/internal/pprof"
 	"github.com/mimecast/dtail/internal/user"
 	"github.com/mimecast/dtail/internal/version"
@@ -29,7 +31,6 @@ func main() {
 	var sshPort int
 	var trustAllHosts bool
 
-	pingTimeoutS := 900
 	userName := user.Name()
 
 	flag.BoolVar(&debugEnable, "debug", false, "Activate debug messages")
@@ -39,7 +40,6 @@ func main() {
 	flag.BoolVar(&silentEnable, "silent", false, "Reduce output")
 	flag.BoolVar(&trustAllHosts, "trustAllHosts", false, "Auto trust all unknown host keys")
 	flag.IntVar(&connectionsPerCPU, "cpc", 10, "How many connections established per CPU core concurrently")
-	flag.IntVar(&pingTimeoutS, "pingTimeout", 10, "The server ping timeout (0 means disable pings)")
 	flag.IntVar(&sshPort, "port", 2222, "SSH server port")
 	flag.StringVar(&cfgFile, "cfg", "", "Config file path")
 	flag.StringVar(&discovery, "discovery", "", "Server discovery method")
@@ -57,10 +57,10 @@ func main() {
 		version.PrintAndExit()
 	}
 
+	ctx := context.Background()
 	serverEnable := false
-	logger.Start(serverEnable, debugEnable, silentEnable, silentEnable)
-	defer logger.Stop()
 
+	logger.Start(ctx, serverEnable, debugEnable, silentEnable, silentEnable)
 	if pprofEnable || config.Common.PProfEnable {
 		pprof.Start()
 	}
@@ -70,9 +70,8 @@ func main() {
 		ServersStr:        serversStr,
 		Discovery:         discovery,
 		UserName:          userName,
-		Files:             files,
+		What:              files,
 		TrustAllHosts:     trustAllHosts,
-		PingTimeout:       pingTimeoutS,
 		Mode:              omode.MapClient,
 	}
 
@@ -81,5 +80,7 @@ func main() {
 		panic(err)
 	}
 
-	client.Start()
+	status := client.Start(ctx)
+	logger.Flush()
+	os.Exit(status)
 }