spindle/engine: stream logs to websocket

Logs are streamed for each running pipeline on a websocket at
/logs/{pipelineID}. engine.TailStep demuxes stdout and stderr from the
container's logs and pipes that out to corresponding stdout and stderr
channels.

These channels are maintained inside engine's container
struct, key'd by the pipeline ID, and protected by a read/write mutex.
engine.LogChannels fetches the stdout/stderr chans as recieve-only if
the pipeline is known to exist.

Signed-off-by: Anirudh Oppiliappan <[e-mail]>

spindle/engine/engine.go

···
 package engine
 
 import (
+	"bufio"
 	"context"
 	"fmt"
 	"io"
···
 	l      *slog.Logger
 	db     *db.DB
 	n      *notifier.Notifier
+
+	chanMu      sync.RWMutex
+	stdoutChans map[string]chan string
+	stderrChans map[string]chan string
 }
 
 func New(ctx context.Context, db *db.DB, n *notifier.Notifier) (*Engine, error) {
···
 
 	l := log.FromContext(ctx).With("component", "spindle")
 
-	return &Engine{docker: dcli, l: l, db: db, n: n}, nil
+	e := &Engine{
+		docker: dcli,
+		l:      l,
+		db:     db,
+		n:      n,
+	}
+
+	e.stdoutChans = make(map[string]chan string, 100)
+	e.stderrChans = make(map[string]chan string, 100)
+
+	return e, nil
 }
 
 // SetupPipeline sets up a new network for the pipeline, and possibly volumes etc.
···
 // ONLY marks pipeline as failed if container's exit code is non-zero.
 // All other errors are bubbled up.
 func (e *Engine) StartSteps(ctx context.Context, steps []*tangled.Pipeline_Step, id, image string) error {
+	// set up logging channels
+	e.chanMu.Lock()
+	if _, exists := e.stdoutChans[id]; !exists {
+		e.stdoutChans[id] = make(chan string, 100)
+	}
+	if _, exists := e.stderrChans[id]; !exists {
+		e.stderrChans[id] = make(chan string, 100)
+	}
+	e.chanMu.Unlock()
+
+	// close channels after all steps are complete
+	defer func() {
+		close(e.stdoutChans[id])
+		close(e.stderrChans[id])
+	}()
+
 	for _, step := range steps {
 		hostConfig := hostConfig(id)
 		resp, err := e.docker.ContainerCreate(ctx, &container.Config{
···
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			err := e.TailStep(ctx, resp.ID)
+			err := e.TailStep(ctx, resp.ID, id)
 			if err != nil {
 				e.l.Error("failed to tail container", "container", resp.ID)
 				return
···
 	return info.State, nil
 }
 
-func (e *Engine) TailStep(ctx context.Context, containerID string) error {
+func (e *Engine) TailStep(ctx context.Context, containerID, pipelineID string) error {
 	logs, err := e.docker.ContainerLogs(ctx, containerID, container.LogsOptions{
 		Follow:     true,
 		ShowStdout: true,
···
 		return err
 	}
 
+	// using StdCopy we demux logs and stream stdout and stderr to different
+	// channels.
+	//
+	//    stdout w||r stdoutCh
+	//    stderr w||r stderrCh
+	//
+
+	rpipeOut, wpipeOut := io.Pipe()
+	rpipeErr, wpipeErr := io.Pipe()
+
 	go func() {
-		_, _ = stdcopy.StdCopy(os.Stdout, os.Stdout, logs)
-		_ = logs.Close()
+		defer wpipeOut.Close()
+		defer wpipeErr.Close()
+		_, err := stdcopy.StdCopy(wpipeOut, wpipeErr, logs)
+		if err != nil && err != io.EOF {
+			e.l.Error("failed to copy logs", "error", err)
+		}
+	}()
+
+	// read from stdout and send to stdout pipe
+	// NOTE: the stdoutCh channnel is closed further up in StartSteps
+	// once all steps are done.
+	go func() {
+		e.chanMu.RLock()
+		stdoutCh := e.stdoutChans[pipelineID]
+		e.chanMu.RUnlock()
+
+		scanner := bufio.NewScanner(rpipeOut)
+		for scanner.Scan() {
+			stdoutCh <- scanner.Text()
+		}
+		if err := scanner.Err(); err != nil {
+			e.l.Error("failed to scan stdout", "error", err)
+		}
+	}()
+
+	// read from stderr and send to stderr pipe
+	// NOTE: the stderrCh channnel is closed further up in StartSteps
+	// once all steps are done.
+	go func() {
+		e.chanMu.RLock()
+		stderrCh := e.stderrChans[pipelineID]
+		e.chanMu.RUnlock()
+
+		scanner := bufio.NewScanner(rpipeErr)
+		for scanner.Scan() {
+			stderrCh <- scanner.Text()
+		}
+		if err := scanner.Err(); err != nil {
+			e.l.Error("failed to scan stderr", "error", err)
+		}
 	}()
+
 	return nil
+}
+
+func (e *Engine) LogChannels(pipelineID string) (stdout <-chan string, stderr <-chan string, ok bool) {
+	e.chanMu.RLock()
+	defer e.chanMu.RUnlock()
+
+	stdoutCh, ok1 := e.stdoutChans[pipelineID]
+	stderrCh, ok2 := e.stderrChans[pipelineID]
+
+	if !ok1 || !ok2 {
+		return nil, nil, false
+	}
+	return stdoutCh, stderrCh, true
 }
 
 func workspaceVolume(id string) string {

spindle/server.go

···
 	"log/slog"
 	"net/http"
 
+	"github.com/go-chi/chi/v5"
 	"golang.org/x/net/context"
 	"tangled.sh/tangled.sh/core/api/tangled"
 	"tangled.sh/tangled.sh/core/jetstream"
···
 	}
 
 	n := notifier.New()
-
 	eng, err := engine.New(ctx, d, &n)
 	if err != nil {
 		return err
···
 }
 
 func (s *Spindle) Router() http.Handler {
-	mux := &http.ServeMux{}
+	mux := chi.NewRouter()
 
 	mux.HandleFunc("/events", s.Events)
+	mux.HandleFunc("/logs/{pipelineID}", s.Logs)
 	return mux
 }
 

spindle/stream.go

···
 package spindle
 
 import (
+	"fmt"
 	"net/http"
 	"time"
 
 	"context"
 
+	"github.com/go-chi/chi/v5"
 	"github.com/gorilla/websocket"
 )
 
···
 			}
 		}
 	}
+}
+
+func (s *Spindle) Logs(w http.ResponseWriter, r *http.Request) {
+	l := s.l.With("handler", "Logs")
+
+	pipelineID := chi.URLParam(r, "pipelineID")
+	if pipelineID == "" {
+		http.Error(w, "pipelineID required", http.StatusBadRequest)
+		return
+	}
+	l = l.With("pipelineID", pipelineID)
+
+	conn, err := upgrader.Upgrade(w, r, nil)
+	if err != nil {
+		l.Error("websocket upgrade failed", "err", err)
+		http.Error(w, "failed to upgrade", http.StatusInternalServerError)
+		return
+	}
+	defer conn.Close()
+	l.Info("upgraded http to wss")
+
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
+
+	go func() {
+		for {
+			if _, _, err := conn.NextReader(); err != nil {
+				l.Info("client disconnected", "err", err)
+				cancel()
+				return
+			}
+		}
+	}()
+
+	if err := s.streamLogs(ctx, conn, pipelineID); err != nil {
+		l.Error("streamLogs failed", "err", err)
+	}
+	l.Info("logs connection closed")
+}
+
+func (s *Spindle) streamLogs(ctx context.Context, conn *websocket.Conn, pipelineID string) error {
+	l := s.l.With("pipelineID", pipelineID)
+
+	stdoutCh, stderrCh, ok := s.eng.LogChannels(pipelineID)
+	if !ok {
+		return fmt.Errorf("pipelineID %q not found", pipelineID)
+	}
+
+	done := make(chan struct{})
+
+	go func() {
+		for {
+			select {
+			case line, ok := <-stdoutCh:
+				if !ok {
+					done <- struct{}{}
+					return
+				}
+				msg := map[string]string{"type": "stdout", "data": line}
+				if err := conn.WriteJSON(msg); err != nil {
+					l.Error("write stdout failed", "err", err)
+					done <- struct{}{}
+					return
+				}
+			case <-ctx.Done():
+				done <- struct{}{}
+				return
+			}
+		}
+	}()
+
+	go func() {
+		for {
+			select {
+			case line, ok := <-stderrCh:
+				if !ok {
+					done <- struct{}{}
+					return
+				}
+				msg := map[string]string{"type": "stderr", "data": line}
+				if err := conn.WriteJSON(msg); err != nil {
+					l.Error("write stderr failed", "err", err)
+					done <- struct{}{}
+					return
+				}
+			case <-ctx.Done():
+				done <- struct{}{}
+				return
+			}
+		}
+	}()
+
+	select {
+	case <-done:
+	case <-ctx.Done():
+	}
+
+	return nil
 }
 
 func (s *Spindle) streamPipelines(conn *websocket.Conn, cursor *string) error {