Implement graceful termination

config/ssh.go

@@ -43,6 +43,10 @@ type SSHConfig struct {
 	// allowed to be sent without a response being received. If this number
 	// is exceeded the connection is considered dead
 	ClientAliveCountMax int `json:"clientAliveCountMax" yaml:"clientAliveCountMax" default:"3" comment:"Maximum number of failed keepalives"`
+	// GracefulTerminationDeadline is the amount of time ContainerSSH will
+	// wait after receiving a SIGTERM for all the clients to disconnect.
+	// After this duraction all connections are forcefully terminated.
+	GracefulTerminationDeadline time.Duration `json:"gracefulTerminationDeadline" yaml:"gracefulTerminationDeadline" default:"0"`
 }
 
 // GenerateHostKey generates a random host key and adds it to SSHConfig

internal/sshserver/serverImpl.go

@@ -86,12 +86,13 @@ func (s *serverImpl) RunWithLifecycle(lifecycle service.Lifecycle) error {
 		s.wg.Add(1)
 		go s.handleConnection(tcpConn)
 	}
+
 	lifecycle.Stopping()
 	s.shuttingDown = true
 	allClientsExited := make(chan struct{})
 	shutdownHandlerExited := make(chan struct{}, 1)
+	s.disconnectClients(lifecycle, allClientsExited)
 	go s.shutdownHandlers.Shutdown(lifecycle.ShutdownContext())
-	go s.disconnectClients(lifecycle, allClientsExited)
 	go s.shutdownHandler(lifecycle, shutdownHandlerExited)
 
 	s.wg.Wait()
@@ -322,6 +323,7 @@ func (s *serverImpl) createConfiguration(
 		ServerVersion:               s.cfg.ServerVersion.String(),
 		BannerCallback:              func(conn ssh.ConnMetadata) string { return s.cfg.Banner },
 	}
+
 	for _, key := range s.hostKeys {
 		serverConfig.AddHostKey(key)
 	}

main.go

@@ -182,10 +182,10 @@ func startServices(cfg config.AppConfig, loggerFactory log.LoggerFactory) error
 		return err
 	}
 
-	return startPool(pool, lifecycle)
+	return startPool(pool, lifecycle, cfg)
 }
 
-func startPool(pool Service, lifecycle service.Lifecycle) error {
+func startPool(pool Service, lifecycle service.Lifecycle, cfg config.AppConfig) error {
 	starting := make(chan struct{})
 	lifecycle.OnStarting(
 		func(s service.Service, l service.Lifecycle) {
@@ -204,12 +204,15 @@ func startPool(pool Service, lifecycle service.Lifecycle) error {
 	rotateSignals := make(chan os.Signal, 1)
 	signal.Notify(exitSignals, exitSignalList...)
 	signal.Notify(rotateSignals, rotateSignalList...)
+
+	deadline := cfg.SSH.GracefulTerminationDeadline + 5 * time.Second
+
 	go func() {
 		if _, ok := <-exitSignals; ok {
 			// ok means the channel wasn't closed
 			shutdownContext, cancelFunc := context.WithTimeout(
 				context.Background(),
-				20*time.Second,
+				deadline,
 			)
 			defer cancelFunc()
 			lifecycle.Stop(