appview,spindle: handle graceful shutdown of websockets

Signed-off-by: oppiliappan <[e-mail]>

appview/pipelines/pipelines.go

···
 		l.Error("websocket upgrade failed", "err", err)
 		return
 	}
-	defer clientConn.Close()
+	defer func() {
+		_ = clientConn.WriteControl(
+			websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "log stream complete"),
+			time.Now().Add(time.Second),
+		)
+		clientConn.Close()
+	}()
 
 	ctx, cancel := context.WithCancel(r.Context())
 	defer cancel()
···
 	// start a goroutine to read from spindle
 	go func() {
 		defer close(msgChan)
+		defer close(errChan)
+
 		for {
 			_, msg, err := spindleConn.ReadMessage()
 			if err != nil {
-				errChan <- err
+				if websocket.IsCloseError(err,
+					websocket.CloseNormalClosure,
+					websocket.CloseGoingAway,
+					websocket.CloseAbnormalClosure) {
+					errChan <- nil // signal graceful end
+				} else {
+					errChan <- err
+				}
 				return
 			}
 			msgChan <- msg
···
 			l.Info("client disconnected")
 			return
 		case err := <-errChan:
-			l.Error("error reading from spindle", "err", err)
+			if err != nil {
+				l.Error("error reading from spindle", "err", err)
+			}
+
+			if err == nil {
+				l.Info("log tail complete")
+			}
+
 			return
 		case msg := <-msgChan:
 			var logLine spindlemodel.LogLine

spindle/db/events.go

···
 
 }
 
+func (d *DB) GetStatus(workflowId models.WorkflowId) (*tangled.PipelineStatus, error) {
+	pipelineAtUri := workflowId.PipelineId.AtUri()
+
+	var eventJson string
+	err := d.QueryRow(
+		`
+		select
+			event from events
+		where
+			nsid = ?
+			and json_extract(event, '$.pipeline') = ?
+			and json_extract(event, '$.workflow') = ?
+		order by
+			created desc
+		limit
+			1
+		`,
+		tangled.PipelineStatusNSID,
+		string(pipelineAtUri),
+		workflowId.Name,
+	).Scan(&eventJson)
+
+	if err != nil {
+		return nil, err
+	}
+
+	var status tangled.PipelineStatus
+	if err := json.Unmarshal([]byte(eventJson), &status); err != nil {
+		return nil, err
+	}
+
+	return &status, nil
+}
+
 func (d *DB) StatusPending(workflowId models.WorkflowId, n *notifier.Notifier) error {
 	return d.createStatusEvent(workflowId, models.StatusKindPending, nil, nil, n)
 }

spindle/engine/engine.go

···
 	}
 	defer wfLogger.Close()
 
-	_, err = stdcopy.StdCopy(wfLogger.Stdout(), wfLogger.Stderr(), logs)
+	_, err = stdcopy.StdCopy(
+		wfLogger.Writer("stdout", stepIdx),
+		wfLogger.Writer("stderr", stepIdx),
+		logs,
+	)
 	if err != nil && err != io.EOF && !errors.Is(err, context.DeadlineExceeded) {
 		return fmt.Errorf("failed to copy logs: %w", err)
 	}

spindle/engine/logger.go

···
 }
 
 func NewWorkflowLogger(baseDir string, wid models.WorkflowId) (*WorkflowLogger, error) {
-	dir := filepath.Join(baseDir, wid.String())
-	if err := os.MkdirAll(dir, 0755); err != nil {
-		return nil, fmt.Errorf("creating log dir: %w", err)
-	}
-
 	path := LogFilePath(baseDir, wid)
 
-	file, err := os.Create(path)
+	file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
 	if err != nil {
 		return nil, fmt.Errorf("creating log file: %w", err)
 	}
···
 	return l.file.Close()
 }
 
-func OpenLogFile(baseDir string, workflowID models.WorkflowId) (*os.File, error) {
-	logPath := LogFilePath(baseDir, workflowID)
-
-	file, err := os.Open(logPath)
-	if err != nil {
-		return nil, fmt.Errorf("error opening log file: %w", err)
-	}
-
-	return file, nil
-}
-
 func LogFilePath(baseDir string, workflowID models.WorkflowId) string {
 	logFilePath := filepath.Join(baseDir, fmt.Sprintf("%s.log", workflowID.String()))
 	return logFilePath
 }
 
-func (l *WorkflowLogger) Stdout() io.Writer {
-	return &jsonWriter{logger: l, stream: "stdout"}
-}
-
-func (l *WorkflowLogger) Stderr() io.Writer {
-	return &jsonWriter{logger: l, stream: "stderr"}
+func (l *WorkflowLogger) Writer(stream string, stepId int) io.Writer {
+	return &jsonWriter{logger: l, stream: stream, stepId: stepId}
 }
 
 type jsonWriter struct {
 	logger *WorkflowLogger
 	stream string
+	stepId int
 }
 
 func (w *jsonWriter) Write(p []byte) (int, error) {
···
 	entry := models.LogLine{
 		Stream: w.stream,
 		Data:   line,
+		StepId: w.stepId,
 	}
 
 	if err := w.logger.encoder.Encode(entry); err != nil {

spindle/models/models.go

···
 type LogLine struct {
 	Stream string `json:"s"`
 	Data   string `json:"d"`
+	StepId int    `json:"i"`
 }

spindle/stream.go

···
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net/http"
 	"strconv"
 	"time"
···
 		http.Error(w, "failed to upgrade", http.StatusInternalServerError)
 		return
 	}
-	defer conn.Close()
+	defer func() {
+		_ = conn.WriteControl(
+			websocket.CloseMessage,
+			websocket.FormatCloseMessage(websocket.CloseNormalClosure, "log stream complete"),
+			time.Now().Add(time.Second),
+		)
+		conn.Close()
+	}()
 	l.Debug("upgraded http to wss")
 
 	ctx, cancel := context.WithCancel(r.Context())
···
 	}()
 
 	if err := s.streamLogsFromDisk(ctx, conn, wid); err != nil {
-		l.Error("log stream failed", "err", err)
+		l.Info("log stream ended", "err", err)
 	}
-	l.Debug("logs connection closed")
+
+	l.Info("logs connection closed")
 }
 
 func (s *Spindle) streamLogsFromDisk(ctx context.Context, conn *websocket.Conn, wid models.WorkflowId) error {
+	status, err := s.db.GetStatus(wid)
+	if err != nil {
+		return err
+	}
+	isFinished := models.StatusKind(status.Status).IsFinish()
+
 	filePath := engine.LogFilePath(s.cfg.Pipelines.LogDir, wid)
 
 	config := tail.Config{
-		Follow:    true,
-		ReOpen:    true,
+		Follow:    !isFinished,
+		ReOpen:    !isFinished,
 		MustExist: false,
-		Location:  &tail.SeekInfo{Offset: 0, Whence: 0},
-		Logger:    tail.DiscardingLogger,
+		Location: &tail.SeekInfo{
+			Offset: 0,
+			Whence: io.SeekStart,
+		},
+		// Logger: tail.DiscardingLogger,
 	}
 
 	t, err := tail.TailFile(filePath, config)
···
 		case <-ctx.Done():
 			return ctx.Err()
 		case line := <-t.Lines:
+			if line == nil && isFinished {
+				return fmt.Errorf("tail completed")
+			}
+
 			if line == nil {
 				return fmt.Errorf("tail channel closed unexpectedly")
 			}