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 {