merge develop

25 files changed, 241 insertions, 266 deletions

diff --git a/internal/clients/baseclient.go b/internal/clients/baseclient.go
index 008a01e..69055a3 100644
--- a/internal/clients/baseclient.go
+++ b/internal/clients/baseclient.go
@@ -66,7 +66,7 @@ func (c *baseClient) makeConnections(maker maker) {
 	c.stats = newTailStats(len(c.connections))
 }
 
-func (c *baseClient) Start(ctx context.Context, statsCh <-chan struct{}) (status int) {
+func (c *baseClient) Start(ctx context.Context, statsCh <-chan string) (status int) {
 	// Periodically check for unknown hosts, and ask the user whether to trust them or not.
 	go c.hostKeyCallback.PromptAddHosts(ctx)
 	// Print client stats every time something on statsCh is recieved.
@@ -99,7 +99,7 @@ func (c *baseClient) start(ctx context.Context, active chan struct{}, i int, con
 	defer func() { <-active }()
 
 	for {
-		connCtx, cancel := conn.Handler.WithCancel(ctx)
+		connCtx, cancel := context.WithCancel(ctx)
 		defer cancel()
 
 		conn.Start(connCtx, cancel, c.throttleCh, c.stats.connectionsEstCh)
diff --git a/internal/clients/client.go b/internal/clients/client.go
index eb8452d..4a547e8 100644
--- a/internal/clients/client.go
+++ b/internal/clients/client.go
@@ -4,5 +4,5 @@ import "context"
 
 // Client is the interface for the end user command line client.
 type Client interface {
-	Start(ctx context.Context, statsCh <-chan struct{}) int
+	Start(ctx context.Context, statsCh <-chan string) int
 }
diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
index 65bbfd7..b5045e2 100644
--- a/internal/clients/handlers/basehandler.go
+++ b/internal/clients/handlers/basehandler.go
@@ -8,12 +8,13 @@ import (
 	"strings"
 	"time"
 
+	"github.com/mimecast/dtail/internal"
 	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/version"
 )
 
 type baseHandler struct {
-	withCancel
+	done         *internal.Done
 	server       string
 	shellStarted bool
 	commands     chan string
@@ -29,6 +30,14 @@ func (h *baseHandler) Status() int {
 	return h.status
 }
 
+func (h *baseHandler) Done() <-chan struct{} {
+	return h.done.Done()
+}
+
+func (h *baseHandler) Shutdown() {
+	h.done.Shutdown()
+}
+
 // SendMessage to the server.
 func (h *baseHandler) SendMessage(command string) error {
 	encoded := base64.StdEncoding.EncodeToString([]byte(command))
@@ -38,7 +47,8 @@ func (h *baseHandler) SendMessage(command string) error {
 	case h.commands <- fmt.Sprintf("protocol %s base64 %v;", version.ProtocolCompat, encoded):
 	case <-time.After(time.Second * 5):
 		return fmt.Errorf("Timed out sending command '%s' (base64: '%s')", command, encoded)
-	case <-h.ctx.Done():
+	case <-h.Done():
+		return nil
 	}
 
 	return nil
@@ -65,7 +75,7 @@ func (h *baseHandler) Read(p []byte) (n int, err error) {
 	select {
 	case command := <-h.commands:
 		n = copy(p, []byte(command))
-	case <-h.ctx.Done():
+	case <-h.Done():
 		return 0, io.EOF
 	}
 	return
@@ -95,10 +105,11 @@ func (h *baseHandler) handleHiddenMessage(message string) {
 	case strings.HasPrefix(message, ".syn close connection"):
 		h.SendMessage(".ack close connection")
 		select {
-		case <-time.After(time.Second * 1):
+		case <-time.After(time.Second * 5):
 			logger.Debug("Shutting down client after timeout and sending ack to server")
-			h.withCancel.shutdown()
-		case <-h.ctx.Done():
+			h.Shutdown()
+		case <-h.Done():
+			return
 		}
 
 	case strings.HasPrefix(message, ".run exitstatus"):
diff --git a/internal/clients/handlers/clienthandler.go b/internal/clients/handlers/clienthandler.go
index fcd8052..2bcb038 100644
--- a/internal/clients/handlers/clienthandler.go
+++ b/internal/clients/handlers/clienthandler.go
@@ -1,6 +1,7 @@
 package handlers
 
 import (
+	"github.com/mimecast/dtail/internal"
 	"github.com/mimecast/dtail/internal/io/logger"
 )
 
@@ -19,9 +20,7 @@ func NewClientHandler(server string) *ClientHandler {
 			shellStarted: false,
 			commands:     make(chan string),
 			status:       -1,
-			withCancel: withCancel{
-				done: make(chan struct{}),
-			},
+			done:         internal.NewDone(),
 		},
 	}
 }
diff --git a/internal/clients/handlers/handler.go b/internal/clients/handlers/handler.go
index c53ca34..afa87e2 100644
--- a/internal/clients/handlers/handler.go
+++ b/internal/clients/handlers/handler.go
@@ -1,7 +1,6 @@
 package handlers
 
 import (
-	"context"
 	"io"
 )
 
@@ -11,6 +10,6 @@ type Handler interface {
 	SendMessage(command string) error
 	Server() string
 	Status() int
-	WithCancel(ctx context.Context) (context.Context, context.CancelFunc)
+	Shutdown()
 	Done() <-chan struct{}
 }
diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go
index 9051015..08ed137 100644
--- a/internal/clients/handlers/healthhandler.go
+++ b/internal/clients/handlers/healthhandler.go
@@ -4,11 +4,13 @@ import (
 	"errors"
 	"fmt"
 	"time"
+
+	"github.com/mimecast/dtail/internal"
 )
 
 // HealthHandler implements the handler required for health checks.
 type HealthHandler struct {
-	withCancel
+	done *internal.Done
 	// Buffer of incoming data from server.
 	receiveBuf []byte
 	// To send commands to the server.
@@ -27,9 +29,7 @@ func NewHealthHandler(server string, receive chan<- string) *HealthHandler {
 		receive:  receive,
 		commands: make(chan string),
 		status:   -1,
-		withCancel: withCancel{
-			done: make(chan struct{}),
-		},
+		done:     internal.NewDone(),
 	}
 
 	return &h
@@ -45,12 +45,23 @@ func (h *HealthHandler) Status() int {
 	return h.status
 }
 
+// Done returns done channel of the handler.
+func (h *HealthHandler) Done() <-chan struct{} {
+	return h.done.Done()
+}
+
+// Shutdown the handler.
+func (h *HealthHandler) Shutdown() {
+	h.done.Shutdown()
+}
+
 // SendMessage sends a DTail command to the server.
 func (h *HealthHandler) SendMessage(command string) error {
 	select {
 	case h.commands <- fmt.Sprintf("%s;", command):
 	case <-time.NewTimer(time.Second * 10).C:
 		return errors.New("Timed out sending command " + command)
+	case <-h.Done():
 	}
 
 	return nil
diff --git a/internal/clients/handlers/maprhandler.go b/internal/clients/handlers/maprhandler.go
index b908f3b..fb71c8f 100644
--- a/internal/clients/handlers/maprhandler.go
+++ b/internal/clients/handlers/maprhandler.go
@@ -3,6 +3,7 @@ package handlers
 import (
 	"strings"
 
+	"github.com/mimecast/dtail/internal"
 	"github.com/mimecast/dtail/internal/io/logger"
 	"github.com/mimecast/dtail/internal/mapr"
 	"github.com/mimecast/dtail/internal/mapr/client"
@@ -24,9 +25,7 @@ func NewMaprHandler(server string, query *mapr.Query, globalGroup *mapr.GlobalGr
 			shellStarted: false,
 			commands:     make(chan string),
 			status:       -1,
-			withCancel: withCancel{
-				done: make(chan struct{}),
-			},
+			done:         internal.NewDone(),
 		},
 		query:     query,
 		aggregate: client.NewAggregate(server, query, globalGroup),
diff --git a/internal/clients/healthclient.go b/internal/clients/healthclient.go
index 7313583..e93f6be 100644
--- a/internal/clients/healthclient.go
+++ b/internal/clients/healthclient.go
@@ -50,7 +50,7 @@ func (c *HealthClient) Start(ctx context.Context) (status int) {
 	conn.Handler = handlers.NewHealthHandler(c.server, receive)
 	conn.Commands = []string{c.mode.String()}
 
-	connCtx, cancel := conn.Handler.WithCancel(ctx)
+	connCtx, cancel := context.WithCancel(ctx)
 	go conn.Start(connCtx, cancel, throttleCh, statsCh)
 
 	for {
diff --git a/internal/clients/maprclient.go b/internal/clients/maprclient.go
index 581db44..6aadd6b 100644
--- a/internal/clients/maprclient.go
+++ b/internal/clients/maprclient.go
@@ -94,7 +94,7 @@ func NewMaprClient(args Args, queryStr string, maprClientMode MaprClientMode) (*
 }
 
 // Start starts the mapreduce client.
-func (c *MaprClient) Start(ctx context.Context, statsCh <-chan struct{}) (status int) {
+func (c *MaprClient) Start(ctx context.Context, statsCh <-chan string) (status int) {
 	go c.periodicReportResults(ctx)
 
 	status = c.baseClient.Start(ctx, statsCh)
@@ -123,7 +123,7 @@ func (c MaprClient) makeCommands() (commands []string) {
 			commands = append(commands, fmt.Sprintf("timeout %d %s %s %s", c.Timeout, modeStr, file, c.Regex.Serialize()))
 			continue
 		}
-		commands = append(commands, fmt.Sprintf("%s %s regex %s", modeStr, file, c.Regex.Serialize()))
+		commands = append(commands, fmt.Sprintf("%s %s %s", modeStr, file, c.Regex.Serialize()))
 	}
 
 	return
diff --git a/internal/clients/remote/connection.go b/internal/clients/remote/connection.go
index 2d97d14..b29ffed 100644
--- a/internal/clients/remote/connection.go
+++ b/internal/clients/remote/connection.go
@@ -177,21 +177,21 @@ func (c *Connection) handle(ctx context.Context, cancel context.CancelFunc, sess
 	}
 
 	go func() {
-		defer cancel()
 		io.Copy(stdinPipe, c.Handler)
+		cancel()
 	}()
 
 	go func() {
-		defer cancel()
 		io.Copy(c.Handler, stdoutPipe)
+		cancel()
 	}()
 
 	go func() {
-		defer cancel()
 		select {
 		case <-c.Handler.Done():
 		case <-ctx.Done():
 		}
+		cancel()
 	}()
 
 	// Send all commands to client.
@@ -207,5 +207,6 @@ func (c *Connection) handle(ctx context.Context, cancel context.CancelFunc, sess
 	}
 
 	<-ctx.Done()
+	c.Handler.Shutdown()
 	return nil
 }
diff --git a/internal/clients/stats.go b/internal/clients/stats.go
index a6ac0c5..17343b5 100644
--- a/internal/clients/stats.go
+++ b/internal/clients/stats.go
@@ -4,9 +4,11 @@ import (
 	"context"
 	"fmt"
 	"runtime"
+	"strings"
 	"sync"
 	"time"
 
+	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/logger"
 )
 
@@ -32,16 +34,18 @@ func newTailStats(connectionsTotal int) *stats {
 
 // Start starts printing client connection stats every time a signal is recieved or
 // connection count has changed.
-func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{}) {
+func (s *stats) Start(ctx context.Context, throttleCh <-chan struct{}, statsCh <-chan string) {
 	var connectedLast int
 
 	for {
 		var force bool
+		var messages []string
 
 		select {
-		case <-statsCh:
+		case message := <-statsCh:
+			messages = append(messages, message)
 			force = true
-		case <-time.After(time.Second * 2):
+		case <-time.After(time.Second * 10):
 		case <-ctx.Done():
 			return
 		}
@@ -54,7 +58,15 @@ func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{})
 		if connected == connectedLast && !force {
 			continue
 		}
-		s.log(connected, newConnections, throttle)
+
+		stats := s.statsLine(connected, newConnections, throttle)
+		switch force {
+		case true:
+			messages = append(messages, fmt.Sprintf("Connection stats: %s", stats))
+			s.printStatsOnInterrupt(messages)
+		default:
+			logger.Info(stats)
+		}
 
 		connectedLast = connected
 		s.mutex.Lock()
@@ -63,15 +75,25 @@ func (s *stats) Start(ctx context.Context, throttleCh, statsCh <-chan struct{})
 	}
 }
 
-func (s *stats) log(connected, newConnections int, throttle int) {
+func (s *stats) printStatsOnInterrupt(messages []string) {
+	logger.Pause()
+	for _, message := range messages {
+		fmt.Println(fmt.Sprintf(" %s", message))
+	}
+	time.Sleep(time.Second * time.Duration(config.InterruptTimeoutS))
+	logger.Resume()
+}
+
+func (s *stats) statsLine(connected, newConnections int, throttle int) string {
 	percConnected := percentOf(float64(s.connectionsTotal), float64(connected))
 
-	connectedStr := fmt.Sprintf("connected=%d/%d(%d%%)", connected, s.connectionsTotal, int(percConnected))
-	newConnStr := fmt.Sprintf("new=%d", newConnections)
-	throttleStr := fmt.Sprintf("throttle=%d", throttle)
-	cpusGoroutinesStr := fmt.Sprintf("cpus/goroutines=%d/%d", runtime.NumCPU(), runtime.NumGoroutine())
+	var stats []string
+	stats = append(stats, fmt.Sprintf("connected=%d/%d(%d%%)", connected, s.connectionsTotal, int(percConnected)))
+	stats = append(stats, fmt.Sprintf("new=%d", newConnections))
+	stats = append(stats, fmt.Sprintf("throttle=%d", throttle))
+	stats = append(stats, fmt.Sprintf("cpus/goroutines=%d/%d", runtime.NumCPU(), runtime.NumGoroutine()))
 
-	logger.Info("stats", connectedStr, newConnStr, throttleStr, cpusGoroutinesStr)
+	return strings.Join(stats, "|")
 }
 
 func (s *stats) numConnected() int {
diff --git a/internal/config/config.go b/internal/config/config.go
index dc96d6b..276ddcf 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -15,6 +15,9 @@ const ScheduleUser string = "DTAIL-SCHEDULE"
 // ContinuousUser is used for non-interactive continuous mapreduce queries.
 const ContinuousUser string = "DTAIL-CONTINUOUS"
 
+// InterruptTimeoutS is used to terminate DTail when Ctrl+C was pressed twice within a given interval.
+const InterruptTimeoutS int = 3
+
 // Client holds a DTail client configuration.
 var Client *ClientConfig
 
diff --git a/internal/config/server.go b/internal/config/server.go
index db12cec..dc0d587 100644
--- a/internal/config/server.go
+++ b/internal/config/server.go
@@ -61,6 +61,7 @@ type ServerConfig struct {
 	Continuous []Continuous `json:",omitempty"`
 }
 
+// ServerRelaxedAuthEnable should be used for development and testing purposes only.
 var ServerRelaxedAuthEnable bool
 
 // Create a new default server configuration.
diff --git a/internal/io/logger/logger.go b/internal/io/logger/logger.go
index d059cbb..b7db0a7 100644
--- a/internal/io/logger/logger.go
+++ b/internal/io/logger/logger.go
@@ -224,7 +224,7 @@ func log(what string, severity string, args []interface{}) string {
 		return ""
 	}
 
-	messages := []string{severity}
+	messages := []string{}
 
 	for _, arg := range args {
 		switch v := arg.(type) {
diff --git a/internal/io/signal/signal.go b/internal/io/signal/signal.go
index bca7e6e..500c530 100644
--- a/internal/io/signal/signal.go
+++ b/internal/io/signal/signal.go
@@ -5,24 +5,37 @@ import (
 	"os"
 	gosignal "os/signal"
 	"syscall"
+	"time"
+
+	"github.com/mimecast/dtail/internal/config"
 )
 
-// StatsCh returns a channel for "please print stats" signalling.
-func StatsCh(ctx context.Context) <-chan struct{} {
-	sigCh := make(chan os.Signal)
-	gosignal.Notify(sigCh, syscall.SIGINFO, syscall.SIGUSR1)
+// InterruptCh returns a channel for "please print stats" signalling.
+func InterruptCh(ctx context.Context) <-chan string {
+	sigIntCh := make(chan os.Signal)
+	gosignal.Notify(sigIntCh, os.Interrupt)
+
+	sigOtherCh := make(chan os.Signal)
+	gosignal.Notify(sigOtherCh, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGQUIT)
 
-	statsCh := make(chan struct{})
+	statsCh := make(chan string)
 
 	go func() {
 		for {
 			select {
-			case <-sigCh:
+			case <-sigIntCh:
 				select {
-				case statsCh <- struct{}{}:
+				case statsCh <- "Hint: Hit Ctrl+C again to exit":
+					select {
+					case <-sigIntCh:
+						os.Exit(0)
+					case <-time.After(time.Second * time.Duration(config.InterruptTimeoutS)):
+					}
 				default:
-					// Stats currently already printed.
+					// Stats already printed.
 				}
+			case <-sigOtherCh:
+				os.Exit(0)
 			case <-ctx.Done():
 				return
 			}
diff --git a/internal/mapr/server/aggregate.go b/internal/mapr/server/aggregate.go
index 1028943..28bb074 100644
--- a/internal/mapr/server/aggregate.go
+++ b/internal/mapr/server/aggregate.go
@@ -6,6 +6,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/mimecast/dtail/internal"
 	"github.com/mimecast/dtail/internal/config"
 	"github.com/mimecast/dtail/internal/io/line"
 	"github.com/mimecast/dtail/internal/io/logger"
@@ -15,6 +16,7 @@ import (
 
 // Aggregate is for aggregating mapreduce data on the DTail server side.
 type Aggregate struct {
+	done *internal.Done
 	// Log lines to process (parsing MAPREDUCE lines).
 	Lines chan line.Line
 	// Hostname of the current server (used to populate $hostname field).
@@ -23,12 +25,12 @@ type Aggregate struct {
 	serialize chan struct{}
 	// Signals to flush data.
 	flush chan struct{}
+	// Signals that data has been flushed
+	flushed chan struct{}
 	// The mapr query
 	query *mapr.Query
 	// The mapr log format parser
 	parser *logformat.Parser
-	cancel context.CancelFunc
-	ctx    context.Context
 }
 
 // NewAggregate return a new server side aggregator.
@@ -64,56 +66,64 @@ func NewAggregate(queryStr string) (*Aggregate, error) {
 		}
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-
 	a := Aggregate{
+		done:      internal.NewDone(),
 		Lines:     make(chan line.Line, 100),
 		serialize: make(chan struct{}),
 		flush:     make(chan struct{}),
+		flushed:   make(chan struct{}),
 		hostname:  s[0],
 		query:     query,
 		parser:    logParser,
-		ctx:       ctx,
-		cancel:    cancel,
 	}
 
 	return &a, nil
 }
 
+// Shutdown the aggregation engine.
+func (a *Aggregate) Shutdown() {
+	a.Flush()
+	a.done.Shutdown()
+}
+
 // Start an aggregation.
 func (a *Aggregate) Start(ctx context.Context, maprLines chan<- string) {
-	defer a.cancel()
 
-	fieldsCh := a.linesToFields(ctx)
+	myCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	go func() {
+		select {
+		case <-myCtx.Done():
+			a.done.Shutdown()
+		case <-a.done.Done():
+			cancel()
+		}
+	}()
+
+	fieldsCh := a.makeFields(myCtx)
 
 	// Add fields (e.g. via 'set' clause)
 	if len(a.query.Set) > 0 {
-		fieldsCh = a.addMoreFields(ctx, fieldsCh)
+		fieldsCh = a.addFields(myCtx, fieldsCh)
 	}
 
-	go a.fieldsToMaprLines(ctx, fieldsCh, maprLines)
-	a.periodicAggregateTimer(ctx)
+	go a.aggregateTimer(myCtx)
+	a.makeMaprLines(myCtx, fieldsCh, maprLines)
 }
 
-// Cancel the aggregation.
-func (a *Aggregate) Cancel() {
-	a.cancel()
-}
-
-func (a *Aggregate) periodicAggregateTimer(ctx context.Context) {
+func (a *Aggregate) aggregateTimer(ctx context.Context) {
 	for {
 		select {
 		case <-time.After(a.query.Interval):
 			a.Serialize(ctx)
 		case <-ctx.Done():
 			return
-		case <-a.ctx.Done():
-			return
 		}
 	}
 }
 
-func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string {
+func (a *Aggregate) makeFields(ctx context.Context) <-chan map[string]string {
 	ch := make(chan map[string]string)
 
 	go func() {
@@ -144,8 +154,6 @@ func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string
 				}
 			case <-ctx.Done():
 				return
-			case <-a.ctx.Done():
-				return
 			}
 		}
 	}()
@@ -153,14 +161,14 @@ func (a *Aggregate) linesToFields(ctx context.Context) <-chan map[string]string
 	return ch
 }
 
-func (a *Aggregate) addMoreFields(ctx context.Context, fieldsCh <-chan map[string]string) <-chan map[string]string {
+func (a *Aggregate) addFields(ctx context.Context, fieldsCh <-chan map[string]string) <-chan map[string]string {
 	ch := make(chan map[string]string)
 
 	go func() {
 		defer close(ch)
 
 		for {
-			// fieldsCh will be closed via 'linesToFields' if ctx is done
+			// fieldsCh will be closed via 'makeFields' if ctx is done
 			fields, ok := <-fieldsCh
 			if !ok {
 				return
@@ -179,7 +187,7 @@ func (a *Aggregate) addMoreFields(ctx context.Context, fieldsCh <-chan map[strin
 	return ch
 }
 
-func (a *Aggregate) fieldsToMaprLines(ctx context.Context, fieldsCh <-chan map[string]string, maprLines chan<- string) {
+func (a *Aggregate) makeMaprLines(ctx context.Context, fieldsCh <-chan map[string]string, maprLines chan<- string) {
 	group := mapr.NewGroupSet()
 
 	serialize := func() {
@@ -200,18 +208,10 @@ func (a *Aggregate) fieldsToMaprLines(ctx context.Context, fieldsCh <-chan map[s
 		case <-a.serialize:
 			serialize()
 		case <-a.flush:
-			logger.Info("Flushing mapreduce result")
 			serialize()
-			a.flush <- struct{}{}
-			logger.Info("Done flushing mapreduce result")
+			a.flushed <- struct{}{}
 		case <-ctx.Done():
 			return
-		case <-a.ctx.Done():
-			logger.Info("Flushing mapreduce result")
-			serialize()
-			a.flush <- struct{}{}
-			logger.Info("Done flushing mapreduce result")
-			return
 		}
 	}
 }
@@ -254,6 +254,8 @@ func (a *Aggregate) aggregate(group *mapr.GroupSet, fields map[string]string) {
 func (a *Aggregate) Serialize(ctx context.Context) {
 	select {
 	case a.serialize <- struct{}{}:
+	case <-time.After(time.Minute):
+		logger.Warn("Starting to serialize mapredice data takes over a minute")
 	case <-ctx.Done():
 	}
 }
@@ -261,15 +263,20 @@ func (a *Aggregate) Serialize(ctx context.Context) {
 // Flush all data.
 func (a *Aggregate) Flush() {
 	select {
-	case <-a.ctx.Done():
-		return
 	case a.flush <- struct{}{}:
+		logger.Info("Flushing mapreduce data")
 	case <-time.After(time.Minute):
+		logger.Warn("Starting to flush mapreduce data takes over a minute")
+		return
+	case <-a.done.Done():
 		return
 	}
 
 	select {
-	case <-a.flush:
+	case <-a.flushed:
+		logger.Info("Done flushing")
 	case <-time.After(time.Minute):
+		logger.Warn("Waiting for data to be flushed takes over a minute")
+	case <-a.done.Done():
 	}
 }
diff --git a/internal/regex/flag.go b/internal/regex/flag.go
index d3ff712..396bda0 100644
--- a/internal/regex/flag.go
+++ b/internal/regex/flag.go
@@ -2,6 +2,7 @@ package regex
 
 import "fmt"
 
+// Flag for regex.
 type Flag int
 
 const (
@@ -15,6 +16,7 @@ const (
 	Noop Flag = iota
 )
 
+// NewFlag returns a new regex flag.
 func NewFlag(str string) (Flag, error) {
 	switch str {
 	case "default":
diff --git a/internal/regex/regex.go b/internal/regex/regex.go
index 707cb48..2561659 100644
--- a/internal/regex/regex.go
+++ b/internal/regex/regex.go
@@ -8,6 +8,7 @@ import (
 	"github.com/mimecast/dtail/internal/io/logger"
 )
 
+// Regex for filtering lines.
 type Regex struct {
 	// The original regex string
 	regexStr string
@@ -24,6 +25,7 @@ func (r Regex) String() string {
 		r.regexStr, r.flags, r.initialized, r.re == nil)
 }
 
+// NewNoop is a noop regex (doing nothing).
 func NewNoop() Regex {
 	return Regex{
 		flags:       []Flag{Noop},
@@ -31,6 +33,7 @@ func NewNoop() Regex {
 	}
 }
 
+// New returns a new regex object.
 func New(regexStr string, flag Flag) (Regex, error) {
 	if regexStr == "" || regexStr == "." || regexStr == ".*" {
 		return NewNoop(), nil
@@ -39,6 +42,10 @@ func New(regexStr string, flag Flag) (Regex, error) {
 }
 
 func new(regexStr string, flags []Flag) (Regex, error) {
+	if len(flags) == 0 {
+		flags = append(flags, Default)
+	}
+
 	r := Regex{
 		regexStr: regexStr,
 		flags:    flags,
@@ -55,6 +62,7 @@ func new(regexStr string, flags []Flag) (Regex, error) {
 	return r, nil
 }
 
+// Match a byte string.
 func (r Regex) Match(bytes []byte) bool {
 	switch r.flags[0] {
 	case Default:
@@ -68,6 +76,7 @@ func (r Regex) Match(bytes []byte) bool {
 	}
 }
 
+// MatchString matches a string.
 func (r Regex) MatchString(str string) bool {
 	switch r.flags[0] {
 	case Default:
@@ -81,6 +90,7 @@ func (r Regex) MatchString(str string) bool {
 	}
 }
 
+// Serialize the regex.
 func (r Regex) Serialize() string {
 	var flags []string
 	for _, flag := range r.flags {
@@ -94,6 +104,7 @@ func (r Regex) Serialize() string {
 	return fmt.Sprintf("regex:%s %s", strings.Join(flags, ","), r.regexStr)
 }
 
+// Deserialize the regex.
 func Deserialize(str string) (Regex, error) {
 	// Get regex string
 	s := strings.SplitN(str, " ", 2)
diff --git a/internal/server/continuous.go b/internal/server/continuous.go
index 583d136..f75c732 100644
--- a/internal/server/continuous.go
+++ b/internal/server/continuous.go
@@ -92,7 +92,7 @@ func (c *continuous) runJob(ctx context.Context, job config.Continuous) {
 	}
 
 	logger.Info(fmt.Sprintf("Starting job %s", job.Name))
-	status := client.Start(jobCtx, make(chan struct{}))
+	status := client.Start(jobCtx, make(chan string))
 	logMessage := fmt.Sprintf("Job exited with status %d", status)
 
 	if status != 0 {
diff --git a/internal/server/handlers/controlhandler.go b/internal/server/handlers/controlhandler.go
index daa9835..8cc5a40 100644
--- a/internal/server/handlers/controlhandler.go
+++ b/internal/server/handlers/controlhandler.go
@@ -1,20 +1,19 @@
 package handlers
 
 import (
-	"context"
 	"fmt"
 	"io"
 	"os"
 	"strings"
 
+	"github.com/mimecast/dtail/internal"
 	"github.com/mimecast/dtail/internal/io/logger"
 	user "github.com/mimecast/dtail/internal/user/server"
 )
 
 // ControlHandler is used for control functions and health monitoring.
 type ControlHandler struct {
-	ctx            context.Context
-	done           chan struct{}
+	done           *internal.Done
 	hostname       string
 	payload        []byte
 	serverMessages chan string
@@ -22,12 +21,11 @@ type ControlHandler struct {
 }
 
 // NewControlHandler returns a new control handler.
-func NewControlHandler(ctx context.Context, user *user.User) (*ControlHandler, <-chan struct{}) {
+func NewControlHandler(user *user.User) *ControlHandler {
 	logger.Debug(user, "Creating control handler")
 
 	h := ControlHandler{
-		ctx:            ctx,
-		done:           make(chan struct{}),
+		done:           internal.NewDone(),
 		serverMessages: make(chan string, 10),
 		user:           user,
 	}
@@ -40,7 +38,17 @@ func NewControlHandler(ctx context.Context, user *user.User) (*ControlHandler, <
 	s := strings.Split(fqdn, ".")
 	h.hostname = s[0]
 
-	return &h, h.done
+	return &h
+}
+
+// Shutdown the handler.
+func (h *ControlHandler) Shutdown() {
+	h.done.Shutdown()
+}
+
+// Done channel of the handler.
+func (h *ControlHandler) Done() <-chan struct{} {
+	return h.done.Done()
 }
 
 // Read is to send data to the client via the Reader interface.
@@ -51,7 +59,7 @@ func (h *ControlHandler) Read(p []byte) (n int, err error) {
 			wholePayload := []byte(fmt.Sprintf("SERVER|%s|%s\n", h.hostname, message))
 			n = copy(p, wholePayload)
 			return
-		case <-h.ctx.Done():
+		case <-h.done.Done():
 			return 0, io.EOF
 		}
 	}
@@ -63,7 +71,7 @@ func (h *ControlHandler) Write(p []byte) (n int, err error) {
 		switch c {
 		case ';':
 			wholePayload := strings.TrimSpace(string(h.payload))
-			h.handleCommand(h.ctx, wholePayload)
+			h.handleCommand(wholePayload)
 			h.payload = nil
 
 		default:
@@ -75,7 +83,7 @@ func (h *ControlHandler) Write(p []byte) (n int, err error) {
 	return
 }
 
-func (h *ControlHandler) handleCommand(ctx context.Context, command string) {
+func (h *ControlHandler) handleCommand(command string) {
 	logger.Info(h.user, command)
 	s := strings.Split(command, " ")
 	logger.Debug(h.user, "Receiving command", command, s)
diff --git a/internal/server/handlers/handler.go b/internal/server/handlers/handler.go
index c42ceb9..b04e854 100644
--- a/internal/server/handlers/handler.go
+++ b/internal/server/handlers/handler.go
@@ -5,4 +5,6 @@ import "io"
 // Handler interface for server side functionality.
 type Handler interface {
 	io.ReadWriter
+	Shutdown()
+	Done() <-chan struct{}
 }
diff --git a/internal/server/handlers/serverhandler.go b/internal/server/handlers/serverhandler.go
index