refactor

2 files changed, 25 insertions, 45 deletions

diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index 843eabc..db917bd 100644
--- a/internal/server/handlers/serverhandler.go
+++ b/internal/server/handlers/serverhandler.go
@@ -29,36 +29,34 @@ const (
 // the Bi-directional communication between SSH client and server.
 // This handler implements the handler of the SSH server.
 type ServerHandler struct {
-	done                *internal.Done
-	lines               chan line.Line
-	regex               string
-	aggregate           *server.Aggregate
-	aggregatedMessages  chan string
-	serverMessages      chan string
-	payload             []byte
-	hostname            string
-	user                *user.User
-	catLimiter          chan struct{}
-	tailLimiter         chan struct{}
-	globalServerWaitFor chan struct{}
-	ackCloseReceived    chan struct{}
-	activeCommands      int32
-	activeReaders       int32
+	done               *internal.Done
+	lines              chan line.Line
+	regex              string
+	aggregate          *server.Aggregate
+	aggregatedMessages chan string
+	serverMessages     chan string
+	payload            []byte
+	hostname           string
+	user               *user.User
+	catLimiter         chan struct{}
+	tailLimiter        chan struct{}
+	ackCloseReceived   chan struct{}
+	activeCommands     int32
+	activeReaders      int32
 }
 
 // NewServerHandler returns the server handler.
-func NewServerHandler(user *user.User, catLimiter, tailLimiter, globalServerWaitFor chan struct{}) *ServerHandler {
+func NewServerHandler(user *user.User, catLimiter, tailLimiter chan struct{}) *ServerHandler {
 	h := ServerHandler{
-		done:                internal.NewDone(),
-		lines:               make(chan line.Line, 100),
-		serverMessages:      make(chan string, 10),
-		aggregatedMessages:  make(chan string, 10),
-		ackCloseReceived:    make(chan struct{}),
-		catLimiter:          catLimiter,
-		tailLimiter:         tailLimiter,
-		globalServerWaitFor: globalServerWaitFor,
-		regex:               ".",
-		user:                user,
+		done:               internal.NewDone(),
+		lines:              make(chan line.Line, 100),
+		serverMessages:     make(chan string, 10),
+		aggregatedMessages: make(chan string, 10),
+		ackCloseReceived:   make(chan struct{}),
+		catLimiter:         catLimiter,
+		tailLimiter:        tailLimiter,
+		regex:              ".",
+		user:               user,
 	}
 
 	fqdn, err := os.Hostname()
diff --git a/internal/server/server.go b/internal/server/server.go
index 31fa85d..d8d43c9 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -7,7 +7,6 @@ import (
 	"io"
 	"net"
 	"strings"
-	"time"
 
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
@@ -33,9 +32,6 @@ type Server struct {
 	sched *scheduler
 	// Mointor log files for pattern (if configured)
 	cont *continuous
-	// Wait counter, e.g. there might be still subprocesses (forked by drun) to be killed.
-	// TODO: Remove this counter.
-	shutdownWaitFor chan struct{}
 }
 
 // New returns a new server.
@@ -46,7 +42,6 @@ func New() *Server {
 		sshServerConfig: &gossh.ServerConfig{},
 		catLimiter:      make(chan struct{}, config.Server.MaxConcurrentCats),
 		tailLimiter:     make(chan struct{}, config.Server.MaxConcurrentTails),
-		shutdownWaitFor: make(chan struct{}, 1000),
 		sched:           newScheduler(),
 		cont:            newContinuous(),
 	}
@@ -82,25 +77,12 @@ func (s *Server) Start(ctx context.Context) int {
 
 	select {
 	case <-ctx.Done():
-		// Wait until all commands/jobs/children are no more!
-		s.wait()
 	}
 
 	// For future use.
 	return 0
 }
 
-func (s *Server) wait() {
-	for {
-		num := len(s.shutdownWaitFor)
-		logger.Debug("Waiting for stuff to finish", num)
-		if num <= 0 {
-			return
-		}
-		time.Sleep(time.Second)
-	}
-}
-
 func (s *Server) listenerLoop(ctx context.Context, listener net.Listener) {
 	logger.Debug("Starting listener loop")
 
@@ -180,7 +162,7 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch
 			case config.ControlUser:
 				handler = handlers.NewControlHandler(user)
 			default:
-				handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor)
+				handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter)
 			}
 
 			terminate := func() {