Fix serverless output draining regressions

2 files changed, 35 insertions, 1 deletions

diff --git a/internal/clients/handlers/basehandler.go b/internal/clients/handlers/basehandler.go
index 8da4556..2979091 100644
--- a/internal/clients/handlers/basehandler.go
+++ b/internal/clients/handlers/basehandler.go
@@ -192,7 +192,9 @@ func (h *baseHandler) handleHiddenMessage(message string) {
 		strings.HasPrefix(message, protocol.HiddenSessionErrorPrefix):
 		h.handleSessionAckMessage(message)
 	case strings.HasPrefix(message, ".syn close connection"):
-		go h.SendMessage(".ack close connection")
+		if err := h.SendMessage(".ack close connection"); err != nil {
+			dlog.Client.Debug(h.server, "Unable to acknowledge close connection", err)
+		}
 		h.Shutdown()
 	}
 }
diff --git a/internal/clients/handlers/basehandler_test.go b/internal/clients/handlers/basehandler_test.go
index 7db2bb8..3e8aaa1 100644
--- a/internal/clients/handlers/basehandler_test.go
+++ b/internal/clients/handlers/basehandler_test.go
@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/mimecast/dtail/internal"
+	"github.com/mimecast/dtail/internal/io/dlog"
 	"github.com/mimecast/dtail/internal/protocol"
 )
 
@@ -172,3 +173,34 @@ func TestHandleSessionAckMessage(t *testing.T) {
 		t.Fatalf("unexpected session ack: %#v", ack)
 	}
 }
+
+func TestHandleCloseConnectionAcknowledgesBeforeShutdown(t *testing.T) {
+	originalLogger := dlog.Client
+	dlog.Client = &dlog.DLog{}
+	t.Cleanup(func() {
+		dlog.Client = originalLogger
+	})
+
+	handler := baseHandler{
+		done:     internal.NewDone(),
+		server:   "server-under-test",
+		commands: make(chan string, 1),
+	}
+
+	handler.handleHiddenMessage(".syn close connection")
+
+	select {
+	case command := <-handler.commands:
+		if command == "" {
+			t.Fatal("expected close acknowledgement command")
+		}
+	case <-time.After(10 * time.Millisecond):
+		t.Fatal("expected close acknowledgement command to be queued")
+	}
+
+	select {
+	case <-handler.Done():
+	default:
+		t.Fatal("expected handler to be shut down after close acknowledgement")
+	}
+}