merge develop

1 files changed, 37 insertions, 69 deletions

diff --git a/internal/clients/handlers/healthhandler.go b/internal/clients/handlers/healthhandler.go
index 0440706..47b594e 100644
--- a/internal/clients/handlers/healthhandler.go
+++ b/internal/clients/handlers/healthhandler.go
@@ -1,88 +1,56 @@
 package handlers
 
 import (
-	"errors"
-	"fmt"
-	"time"
+	"strings"
 
 	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/io/dlog"
+	"github.com/mimecast/dtail/internal/protocol"
 )
 
-// HealthHandler implements the handler required for health checks.
+// HealthHandler is the handler used on the client side for running mapreduce
+// aggregations.
 type HealthHandler struct {
-	done *internal.Done
-	// Buffer of incoming data from server.
-	receiveBuf []byte
-	// To send commands to the server.
-	commands chan string
-	// To receive messages from the server.
-	receive chan<- string
-	// The remote server address
-	server string
-	// The return status.
-	status int
-}
-
-// NewHealthHandler returns a new health check handler.
-func NewHealthHandler(server string, receive chan<- string) *HealthHandler {
-	h := HealthHandler{
-		server:   server,
-		receive:  receive,
-		commands: make(chan string),
-		status:   -1,
-		done:     internal.NewDone(),
-	}
-
-	return &h
-}
-
-// Server returns the remote server name.
-func (h *HealthHandler) Server() string {
-	return h.server
-}
-
-// Status of the handler.
-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():
+	baseHandler
+}
+
+// NewHealthHandler returns a new health client handler.
+func NewHealthHandler(server string) *HealthHandler {
+	dlog.Client.Debug(server, "Creating new health handler")
+	return &HealthHandler{
+		baseHandler: baseHandler{
+			server:       server,
+			shellStarted: false,
+			commands:     make(chan string),
+			status:       2, // Assume CRITICAL status by default.
+			done:         internal.NewDone(),
+		},
 	}
-
-	return nil
 }
 
-// Server writes byte stream to client.
+// Read data from the dtail server via Writer interface.
 func (h *HealthHandler) Write(p []byte) (n int, err error) {
 	for _, b := range p {
-		h.receiveBuf = append(h.receiveBuf, b)
-		if b == '\n' {
-			h.receive <- string(h.receiveBuf)
-			h.receiveBuf = h.receiveBuf[:0]
+		switch b {
+		case '\n', protocol.MessageDelimiter:
+			message := h.baseHandler.receiveBuf.String()
+			h.handleMessage(message)
+			h.baseHandler.receiveBuf.Reset()
+		default:
+			h.baseHandler.receiveBuf.WriteByte(b)
 		}
 	}
-
 	return len(p), nil
 }
 
-// Server reads byte stream from client.
-func (h *HealthHandler) Read(p []byte) (n int, err error) {
-	n = copy(p, []byte(<-h.commands))
-	return
+func (h *HealthHandler) handleMessage(message string) {
+	if len(message) > 0 && message[0] == '.' {
+		h.baseHandler.handleHiddenMessage(message)
+		return
+	}
+	s := strings.Split(message, protocol.FieldDelimiter)
+	message = s[len(s)-1]
+	if message == "OK" {
+		h.baseHandler.status = 0
+	}
 }