From 799b9b69ba08b898e13026b7ecab9f9f58580a82 Mon Sep 17 00:00:00 2001
From: Paul Buetow <pbuetow@mimecast.com>
Date: Tue, 8 Dec 2020 14:49:41 +0000
Subject: merge develop

---
 internal/server/server.go | 44 +++++++++++++++++---------------------------
 1 file changed, 17 insertions(+), 27 deletions(-)

(limited to 'internal/server/server.go')

diff --git a/internal/server/server.go b/internal/server/server.go
index a446738..31fa85d 100644
--- a/internal/server/server.go
+++ b/internal/server/server.go
@@ -11,7 +11,6 @@ import (
 
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
-	"github.com/mimecast/dtail/internal/server/background"
 	"github.com/mimecast/dtail/internal/server/handlers"
 	"github.com/mimecast/dtail/internal/ssh/server"
 	user "github.com/mimecast/dtail/internal/user/server"
@@ -35,9 +34,8 @@ type Server struct {
 	// 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{}
-	// Background jobs
-	background background.Background
 }
 
 // New returns a new server.
@@ -51,7 +49,6 @@ func New() *Server {
 		shutdownWaitFor: make(chan struct{}, 1000),
 		sched:           newScheduler(),
 		cont:            newContinuous(),
-		background:      background.New(),
 	}
 
 	s.sshServerConfig.PasswordCallback = s.Callback
@@ -178,53 +175,46 @@ func (s *Server) handleRequests(ctx context.Context, sshConn gossh.Conn, in <-ch
 
 		switch req.Type {
 		case "shell":
-			handlerCtx, cancel := context.WithCancel(ctx)
-
 			var handler handlers.Handler
-			var done <-chan struct{}
-
 			switch user.Name {
 			case config.ControlUser:
-				handler, done = handlers.NewControlHandler(handlerCtx, user)
+				handler = handlers.NewControlHandler(user)
 			default:
-				handler, done = handlers.NewServerHandler(handlerCtx, ctx, user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor, s.background)
+				handler = handlers.NewServerHandler(user, s.catLimiter, s.tailLimiter, s.shutdownWaitFor)
 			}
 
-			go func() {
-				// Handler finished work, cancel all remaining routines
-				defer cancel()
-
-				<-done
-			}()
+			terminate := func() {
+				handler.Shutdown()
+				sshConn.Close()
+			}
 
 			go func() {
 				// Broken pipe, cancel
-				defer cancel()
-
 				io.Copy(channel, handler)
+				terminate()
 			}()
 
 			go func() {
 				// Broken pipe, cancel
-				defer cancel()
-
 				io.Copy(handler, channel)
+				terminate()
 			}()
 
 			go func() {
-				defer cancel()
+				select {
+				case <-ctx.Done():
+				case <-handler.Done():
+				}
+				terminate()
+			}()
 
+			go func() {
 				if err := sshConn.Wait(); err != nil && err != io.EOF {
 					logger.Error(user, err)
 				}
 				s.stats.decrementConnections()
 				logger.Info(user, "Good bye Mister!")
-			}()
-
-			go func() {
-				<-handlerCtx.Done()
-				sshConn.Close()
-				logger.Info(user, "Closed SSH connection")
+				terminate()
 			}()
 
 			// Only serving shell type
-- 
cgit v1.2.3


From ab676c2b484225ed22765b23d8f0545088ecd610 Mon Sep 17 00:00:00 2001
From: Paul Buetow <git@mx.buetow.org>
Date: Sat, 26 Dec 2020 10:48:51 +0000
Subject: code cleanup and minor refactorings

---
 internal/server/server.go | 24 ++----------------------
 1 file changed, 2 insertions(+), 22 deletions(-)

(limited to 'internal/server/server.go')

diff --git a/internal/server/server.go b/internal/server/server.go
index 31fa85d..a20737e 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(),
 	}
@@ -80,27 +75,12 @@ func (s *Server) Start(ctx context.Context) int {
 	go s.cont.start(ctx)
 	go s.listenerLoop(ctx, listener)
 
-	select {
-	case <-ctx.Done():
-		// Wait until all commands/jobs/children are no more!
-		s.wait()
-	}
+	<-ctx.Done()
 
 	// 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 +160,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() {
-- 
cgit v1.2.3