fix(fieldagent): stop log tail on WebSocket disconnect

internal/fieldagent/log_session_manager.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/eclipse-iofog/edgelet/internal/config"
@@ -87,6 +88,7 @@ type LogSessionInfo struct {
 	Session     *LogSession
 	ContainerID string
 	IsStreaming bool
+	tailCancel  context.CancelFunc
 }
 
 // LogSessionManager manages active log sessions
@@ -228,6 +230,11 @@ func (lsm *LogSessionManager) HandleLogSessions(fetchedSessions []*LogSession) {
 		if info, exists := lsm.activeSessions[sessionID]; exists {
 			// Update existing session if needed
 			info.Session = session
+			wsHandler := lsm.webSocketHandlers[sessionID]
+			if wsHandler != nil && !wsHandler.IsConnected() {
+				logging.LogInfo(logSessionManagerModuleName, fmt.Sprintf("Reconnecting log WebSocket for session: %s", sessionID))
+				lsm.reconnectLogSessionLocked(sessionID)
+			}
 			if !info.IsStreaming && session.Status == "ACTIVE" {
 				// Session was pending, now active - start streaming
 				logging.LogInfo(logSessionManagerModuleName, fmt.Sprintf("Session became active, starting stream: %s", sessionID))
@@ -359,6 +366,8 @@ func (lsm *LogSessionManager) startMicroserviceLogStreamingLocked(sessionID, mic
 		return
 	}
 
+	lsm.stopLogStreamingLocked(sessionID)
+
 	// Look up container via the engine abstraction (works for Docker, iofog, and all others)
 	container, err := lsm.containerEngine.GetContainer(microserviceUUID)
 	if err != nil {
@@ -382,7 +391,26 @@ func (lsm *LogSessionManager) startMicroserviceLogStreamingLocked(sessionID, mic
 
 	if tailConfig == nil {
 		tailConfig = parseTailConfig(nil)
+	} else {
+		cfgCopy := *tailConfig
+		tailConfig = &cfgCopy
+	}
+
+	tailCtx, tailCancel := context.WithCancel(context.Background())
+	handedOff := false
+	defer func() {
+		if !handedOff {
+			tailCancel()
+		}
+	}()
+
+	tailConfig.Ctx = tailCtx
+	info, ok := lsm.activeSessions[sessionID]
+	if !ok {
+		return
 	}
+	info.tailCancel = tailCancel
+	handedOff = true
 
 	handler := &logTailHandler{
 		sessionID:        sessionID,
@@ -393,10 +421,8 @@ func (lsm *LogSessionManager) startMicroserviceLogStreamingLocked(sessionID, mic
 
 	lsm.tailCallbacks[sessionID] = handler
 
-	if info, ok := lsm.activeSessions[sessionID]; ok {
-		info.ContainerID = containerID
-		info.IsStreaming = true
-	}
+	info.ContainerID = containerID
+	info.IsStreaming = true
 
 	// Run TailContainerLogs in a goroutine — it blocks in follow mode until the session ends.
 	// If called synchronously, lsm.mu would be held for the entire stream, deadlocking
@@ -430,16 +456,15 @@ type logTailHandler struct {
 	microserviceUUID string
 	wsHandler        *LogSessionWebSocketHandler
 	lsm              *LogSessionManager
+	stopped          atomic.Bool
 }
 
 func (h *logTailHandler) OnLogLine(_, _ string, lineBytes []byte, _ engine.StreamType) {
-	if h.wsHandler != nil && len(lineBytes) > 0 {
-		// Always call SendMessage — it buffers when not active and sends when active
-		msgType := byte(6) // LogTypeLogLine
-		if err := h.wsHandler.SendMessage(msgType, lineBytes); err != nil {
-			logging.LogError(logSessionManagerModuleName, "Error sending log line", err)
-		}
+	if h.stopped.Load() || h.wsHandler == nil || len(lineBytes) == 0 {
+		return
 	}
+	msgType := byte(6) // LogTypeLogLine
+	_ = h.wsHandler.SendMessage(msgType, lineBytes)
 }
 
 func (h *logTailHandler) OnComplete(sessionID string) {
@@ -468,6 +493,8 @@ func (h *logTailHandler) OnError(sessionID string, err error) {
 func (lsm *LogSessionManager) startFogLogStreamingLocked(sessionID, iofogUUID string, tailConfig *utils.TailConfig) {
 	logging.LogInfo(logSessionManagerModuleName, fmt.Sprintf("Starting fog log streaming: sessionId=%s, iofogUuid=%s", sessionID, iofogUUID))
 
+	lsm.stopLogStreamingLocked(sessionID)
+
 	// Get WebSocket handler
 	wsHandler := lsm.webSocketHandlers[sessionID]
 	if wsHandler == nil {
@@ -505,17 +532,16 @@ type fogLogHandler struct {
 	iofogUUID string
 	wsHandler *LogSessionWebSocketHandler
 	lsm       *LogSessionManager
+	stopped   atomic.Bool
 }
 
 func (h *fogLogHandler) OnLogLine(_, _, line string) {
-	if h.wsHandler != nil && line != "" {
-		// Always call SendMessage — it buffers when not active and sends when active
-		lineBytes := []byte(line)
-		msgType := byte(6) // LogTypeLogLine
-		if err := h.wsHandler.SendMessage(msgType, lineBytes); err != nil {
-			logging.LogError(logSessionManagerModuleName, "Error sending log line", err)
-		}
+	if h.stopped.Load() || h.wsHandler == nil || line == "" {
+		return
 	}
+	lineBytes := []byte(line)
+	msgType := byte(6) // LogTypeLogLine
+	_ = h.wsHandler.SendMessage(msgType, lineBytes)
 }
 
 func (h *fogLogHandler) OnComplete(sessionID string) {
@@ -544,16 +570,58 @@ func (h *fogLogHandler) OnError(sessionID string, err error) {
 	delete(h.lsm.localLogReaders, sessionID)
 }
 
-// stopLogSessionLocked stops a log session (must be called with lock held)
-func (lsm *LogSessionManager) stopLogSessionLocked(sessionID string) {
-	// Stop LocalLogReader if it exists
+// stopLogStreamingLocked stops active tail readers without removing the session (must be called with lock held).
+func (lsm *LogSessionManager) stopLogStreamingLocked(sessionID string) {
 	if reader, ok := lsm.localLogReaders[sessionID]; ok {
 		reader.Stop()
 		delete(lsm.localLogReaders, sessionID)
 	}
 
-	// Remove engine tail callback if it exists
-	delete(lsm.tailCallbacks, sessionID)
+	if handler, ok := lsm.tailCallbacks[sessionID]; ok {
+		handler.stopped.Store(true)
+		delete(lsm.tailCallbacks, sessionID)
+	}
+
+	if info, ok := lsm.activeSessions[sessionID]; ok {
+		if info.tailCancel != nil {
+			info.tailCancel()
+			info.tailCancel = nil
+		}
+		info.IsStreaming = false
+	}
+}
+
+// HandleWebSocketTransportClose stops log tailing when the controller WebSocket drops.
+// The session row is retained so HandleLogSessions can reconnect on the next poll.
+func (lsm *LogSessionManager) HandleWebSocketTransportClose(sessionID string) {
+	lsm.mu.Lock()
+	defer lsm.mu.Unlock()
+	logging.LogInfo(logSessionManagerModuleName, fmt.Sprintf("Stopping log streaming after WebSocket transport close: %s", sessionID))
+	lsm.stopLogStreamingLocked(sessionID)
+}
+
+func (lsm *LogSessionManager) reconnectLogSessionLocked(sessionID string) {
+	wsHandler := lsm.webSocketHandlers[sessionID]
+	if wsHandler == nil {
+		return
+	}
+	if wsHandler.IsConnected() {
+		return
+	}
+	lsm.mu.Unlock()
+	if wsHandler.needsResetBeforeConnect() {
+		wsHandler.Reset()
+	}
+	err := wsHandler.Connect()
+	lsm.mu.Lock()
+	if err != nil {
+		logging.LogError(logSessionManagerModuleName, fmt.Sprintf("Failed to reconnect log WebSocket: %s", sessionID), err)
+	}
+}
+
+// stopLogSessionLocked stops a log session (must be called with lock held)
+func (lsm *LogSessionManager) stopLogSessionLocked(sessionID string) {
+	lsm.stopLogStreamingLocked(sessionID)
 
 	logging.LogInfo(logSessionManagerModuleName, fmt.Sprintf("Stopping log session: %s", sessionID))
 

internal/fieldagent/log_session_manager_test.go

@@ -0,0 +1,55 @@
+package fieldagent
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+func TestLogSessionManager_HandleWebSocketTransportCloseStopsTail(t *testing.T) {
+	lsm := &LogSessionManager{
+		activeSessions: make(map[string]*LogSessionInfo),
+		tailCallbacks:  make(map[string]*logTailHandler),
+	}
+
+	sessionID := "log-session-stop"
+	tailCtx, cancel := context.WithCancel(context.Background())
+
+	info := &LogSessionInfo{
+		Session:     &LogSession{SessionID: sessionID},
+		IsStreaming: true,
+		tailCancel:  cancel,
+	}
+	lsm.activeSessions[sessionID] = info
+
+	handler := &logTailHandler{sessionID: sessionID, lsm: lsm}
+	lsm.tailCallbacks[sessionID] = handler
+
+	lsm.HandleWebSocketTransportClose(sessionID)
+
+	if info.IsStreaming {
+		t.Fatal("expected IsStreaming=false after transport close")
+	}
+	if !handler.stopped.Load() {
+		t.Fatal("expected tail handler marked stopped")
+	}
+
+	select {
+	case <-tailCtx.Done():
+	case <-time.After(time.Second):
+		t.Fatal("expected tail context canceled")
+	}
+}
+
+func TestLogHandler_SendMessageBuffersWhenDisconnected(t *testing.T) {
+	h := newTestLogHandler("ws://unused", "log-buffer")
+	h.isConnected.Store(false)
+	h.isActive.Store(false)
+
+	if err := h.SendMessage(LogTypeLogLine, []byte("hello")); err != nil {
+		t.Fatalf("SendMessage: %v", err)
+	}
+	if h.bufferedFrames.Load() != 1 {
+		t.Fatalf("bufferedFrames = %d, want 1", h.bufferedFrames.Load())
+	}
+}

internal/fieldagent/log_websocket.go

@@ -266,43 +266,42 @@ func (h *LogSessionWebSocketHandler) transitionState(from, to LogConnectionState
 	return false
 }
 
-// SendMessage sends a log message to the controller
+// SendMessage sends a log message to the controller.
+// Output is buffered while disconnected or pending activation (exec-session parity).
 func (h *LogSessionWebSocketHandler) SendMessage(msgType byte, data []byte) error {
-	if !h.isConnected.Load() {
-		logging.LogWarn(logWebSocketModuleName, "Cannot send message - not connected")
-		return errors.New("not connected")
+	if !h.isConnected.Load() || !h.isActive.Load() {
+		logging.LogDebug(logWebSocketModuleName, fmt.Sprintf("Buffering log output while connection is not active: connected=%v active=%v type=%d length=%d",
+			h.isConnected.Load(), h.isActive.Load(), msgType, len(data)))
+		return h.bufferOutput(msgType, data)
 	}
+	return h.writeOutboundMessage(msgType, data)
+}
 
-	// If not active, buffer the output
-	if !h.isActive.Load() {
-		logging.LogDebug(logWebSocketModuleName, fmt.Sprintf("Buffering log output while connection is not active: type=%d, length=%d", msgType, len(data)))
-
-		// Check buffer limits
-		if h.bufferedFrames.Load() >= logMaxBufferedFrames {
-			logging.LogWarn(logWebSocketModuleName, "Buffer full, dropping frame")
-			return errors.New("buffer full")
-		}
-		if h.bufferedSize.Load()+int64(len(data)) > logMaxBufferSize {
-			logging.LogWarn(logWebSocketModuleName, "Buffer size limit reached, dropping frame")
-			return errors.New("buffer size limit reached")
-		}
+func (h *LogSessionWebSocketHandler) bufferOutput(msgType byte, data []byte) error {
+	if h.bufferedFrames.Load() >= logMaxBufferedFrames {
+		logging.LogWarn(logWebSocketModuleName, "Buffer full, dropping frame")
+		return errors.New("buffer full")
+	}
+	if h.bufferedSize.Load()+int64(len(data)) > logMaxBufferSize {
+		logging.LogWarn(logWebSocketModuleName, "Buffer size limit reached, dropping frame")
+		return errors.New("buffer size limit reached")
+	}
 
-		// Store both the type and data so flushBuffer can preserve the original msgType.
-		bufferedData := make([]byte, len(data))
-		copy(bufferedData, data)
+	bufferedData := make([]byte, len(data))
+	copy(bufferedData, data)
 
-		select {
-		case h.outputBuffer <- logFrame{msgType: msgType, data: bufferedData}:
-			h.bufferedFrames.Add(1)
-			h.bufferedSize.Add(int64(len(data)))
-		default:
-			logging.LogWarn(logWebSocketModuleName, "Buffer channel full, dropping frame")
-			return errors.New("buffer channel full")
-		}
-		return nil
+	select {
+	case h.outputBuffer <- logFrame{msgType: msgType, data: bufferedData}:
+		h.bufferedFrames.Add(1)
+		h.bufferedSize.Add(int64(len(data)))
+	default:
+		logging.LogWarn(logWebSocketModuleName, "Buffer channel full, dropping frame")
+		return errors.New("buffer channel full")
 	}
+	return nil
+}
 
-	// Pack message using MessagePack
+func (h *LogSessionWebSocketHandler) writeOutboundMessage(msgType byte, data []byte) error {
 	var buf bytes.Buffer
 	enc := msgpack.NewEncoder(&buf)
 
@@ -576,7 +575,7 @@ func (h *LogSessionWebSocketHandler) handleMessage(data []byte) {
 		h.handleLogStart(msgData, sessionID)
 	case LogTypeLogStop:
 		logging.LogInfo(logWebSocketModuleName, fmt.Sprintf("Received LOG_STOP message for session: %s", sessionID))
-		h.handleClose()
+		h.handleSessionClose()
 	case LogTypeLogError:
 		h.handleLogError(msgData)
 	default:
@@ -644,11 +643,10 @@ func (h *LogSessionWebSocketHandler) flushBuffer() {
 	for {
 		select {
 		case frame := <-h.outputBuffer:
-			if err := h.SendMessage(frame.msgType, frame.data); err != nil {
+			h.bufferedFrames.Add(-1)
+			h.bufferedSize.Add(-int64(len(frame.data)))
+			if err := h.writeOutboundMessage(frame.msgType, frame.data); err != nil {
 				logging.LogError(logWebSocketModuleName, "Failed to send buffered message", err)
-			} else {
-				h.bufferedFrames.Add(-1)
-				h.bufferedSize.Add(-int64(len(frame.data)))
 			}
 		default:
 			return
@@ -672,20 +670,44 @@ func (h *LogSessionWebSocketHandler) IsActive() bool {
 	return h.isActive.Load()
 }
 
-// handleClose handles WebSocket close
-func (h *LogSessionWebSocketHandler) handleClose() {
+// needsResetBeforeConnect reports whether stale transport state must be torn down before Connect().
+func (h *LogSessionWebSocketHandler) needsResetBeforeConnect() bool {
+	if h == nil {
+		return false
+	}
+	if h.IsConnected() {
+		return true
+	}
+	return h.GetConnectionState() != LogStateDisconnected
+}
+
+func (h *LogSessionWebSocketHandler) handleTransportClose() {
 	if !h.isConnected.Load() {
 		logging.LogDebug(logWebSocketModuleName, fmt.Sprintf("Already disconnected for session: %s", h.sessionID))
 		return
 	}
 
-	logging.LogInfo(logWebSocketModuleName, fmt.Sprintf("Handling close for session: %s, connectionState=%v",
-		h.sessionID, h.state.Load()))
+	logging.LogInfo(logWebSocketModuleName, fmt.Sprintf("Transport closed for log session: %s (tail stopped; session retained for reconnect)", h.sessionID))
 
 	h.isConnected.Store(false)
 	h.isActive.Store(false)
 	h.state.Store(LogStateDisconnected)
+	if h.logSessionManager != nil {
+		h.logSessionManager.HandleWebSocketTransportClose(h.sessionID)
+	}
 	h.Reset()
+}
 
-	logging.LogInfo(logWebSocketModuleName, fmt.Sprintf("Close handling completed for session: %s", h.sessionID))
+func (h *LogSessionWebSocketHandler) handleSessionClose() {
+	sessionID := h.sessionID
+	if h.logSessionManager != nil {
+		go h.logSessionManager.StopLogSession(sessionID)
+		return
+	}
+	h.handleTransportClose()
+}
+
+// handleClose handles unexpected WebSocket transport loss.
+func (h *LogSessionWebSocketHandler) handleClose() {
+	h.handleTransportClose()
 }