From 7f7d89ac561f40993247a710c8ce3b5892df4cc0 Mon Sep 17 00:00:00 2001 From: Anirudh Oppiliappan Date: Thu, 12 Jun 2025 12:54:23 +0300 Subject: [PATCH] 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 --- spindle/engine/engine.go | 103 +++++++++++++++++++++++++++++++++++++-- spindle/server.go | 5 +- spindle/stream.go | 100 +++++++++++++++++++++++++++++++++++++ 3 files changed, 201 insertions(+), 7 deletions(-) diff --git a/spindle/engine/engine.go b/spindle/engine/engine.go index 29e2316b..8ba534ca 100644 --- a/spindle/engine/engine.go +++ b/spindle/engine/engine.go @@ -1,6 +1,7 @@ package engine import ( + "bufio" "context" "fmt" "io" @@ -32,6 +33,10 @@ type Engine struct { 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) { @@ -42,7 +47,17 @@ 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. @@ -136,6 +151,22 @@ func (e *Engine) StartWorkflows(ctx context.Context, pipeline *tangled.Pipeline, // 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{ @@ -166,7 +197,7 @@ func (e *Engine) StartSteps(ctx context.Context, steps []*tangled.Pipeline_Step, 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 @@ -211,7 +242,7 @@ func (e *Engine) WaitStep(ctx context.Context, containerID string) (*container.S 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, @@ -223,13 +254,75 @@ func (e *Engine) TailStep(ctx context.Context, containerID string) error { 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() { + 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() { - _, _ = stdcopy.StdCopy(os.Stdout, os.Stdout, logs) - _ = logs.Close() + 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 { return "workspace-" + id } diff --git a/spindle/server.go b/spindle/server.go index f2beb4d7..81158e00 100644 --- a/spindle/server.go +++ b/spindle/server.go @@ -6,6 +6,7 @@ import ( "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" @@ -52,7 +53,6 @@ func Run(ctx context.Context) error { } n := notifier.New() - eng, err := engine.New(ctx, d, &n) if err != nil { return err @@ -89,9 +89,10 @@ func Run(ctx context.Context) error { } 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 } diff --git a/spindle/stream.go b/spindle/stream.go index 9167d47f..ffc70268 100644 --- a/spindle/stream.go +++ b/spindle/stream.go @@ -1,11 +1,13 @@ package spindle import ( + "fmt" "net/http" "time" "context" + "github.com/go-chi/chi/v5" "github.com/gorilla/websocket" ) @@ -74,6 +76,104 @@ func (s *Spindle) Events(w http.ResponseWriter, r *http.Request) { } } +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 { ops, err := s.db.GetPipelineStatusAsRecords(*cursor) if err != nil { -- 2.51.2