diff --git a/eventstream/eventstream_test.go b/eventstream/eventstream_test.go --- a/eventstream/eventstream_test.go +++ b/eventstream/eventstream_test.go @@ -293,6 +293,55 @@ t.Fatalf("got %d rows, want %d", count, total) } } +func TestHighWaterSeedsClockFromStoredEvents(t *testing.T) { + db, err := sql.Open("sqlite3", t.TempDir()+"/events.db") + if err != nil { + t.Fatalf("open: %v", err) + } + t.Cleanup(func() { db.Close() }) + if _, err := db.Exec(`create table events ( + rkey text not null, + nsid text not null, + event text not null, + created integer not null, + primary key (rkey, nsid) + )`); err != nil { + t.Fatalf("schema: %v", err) + } + + stored := time.Now().Add(time.Hour).UnixNano() + if _, err := db.Exec( + `insert into events (rkey, nsid, event, created) values (?, ?, ?, ?)`, + "stored", "sh.tangled.test", "{}", stored, + ); err != nil { + t.Fatalf("seed event: %v", err) + } + + cut, err := HighWater(db) + if err != nil { + t.Fatalf("HighWater() error = %v", err) + } + if cut < stored { + t.Fatalf("HighWater() = %d, want at least stored cursor %d", cut, stored) + } + + n := notifier.New() + if err := Insert(db, Event{ + Rkey: "new", + Nsid: "sh.tangled.test", + EventJson: json.RawMessage("{}"), + }, &n); err != nil { + t.Fatalf("Insert() error = %v", err) + } + events, err := List(db, cut, 10) + if err != nil { + t.Fatalf("List() error = %v", err) + } + if len(events) != 1 || events[0].Rkey != "new" || events[0].Created <= cut { + t.Fatalf("events after cut = %+v, want only new event above %d", events, cut) + } +} + func isCloseErr(err error) bool { if err == nil { return false diff --git a/eventstream/store.go b/eventstream/store.go --- a/eventstream/store.go +++ b/eventstream/store.go @@ -19,6 +19,35 @@ clockMu sync.Mutex lastNanos int64 ) +func HighWater(s Store) (int64, error) { + clockMu.Lock() + defer clockMu.Unlock() + + rows, err := s.Query(`select coalesce(max(created), 0) from events`) + if err != nil { + return 0, err + } + defer rows.Close() + + var created int64 + if !rows.Next() { + if err := rows.Err(); err != nil { + return 0, err + } + return 0, sql.ErrNoRows + } + if err := rows.Scan(&created); err != nil { + return 0, err + } + if err := rows.Err(); err != nil { + return 0, err + } + if created > lastNanos { + lastNanos = created + } + return lastNanos, nil +} + func Insert(s Store, ev Event, n *notifier.Notifier) error { clockMu.Lock() defer clockMu.Unlock() diff --git a/spindle/db/events.go b/spindle/db/events.go --- a/spindle/db/events.go +++ b/spindle/db/events.go @@ -19,6 +19,10 @@ func (d *DB) GetEvents(cursor int64, limit int) ([]eventstream.Event, error) { return eventstream.List(d, cursor, limit) } +func (d *DB) EventHighWater() (int64, error) { + return eventstream.HighWater(d) +} + func (d *DB) CreatePipelineEvent(rkey string, pipeline tangled.Pipeline, n *notifier.Notifier) error { eventJson, err := json.Marshal(pipeline) if err != nil { diff --git a/spindle/mill/executor/executor.go b/spindle/mill/executor/executor.go new file mode 100644 --- /dev/null +++ b/spindle/mill/executor/executor.go @@ -0,0 +1,532 @@ +package executor + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "maps" + "net/http" + "runtime" + "sync" + "time" + + "github.com/bluesky-social/indigo/atproto/syntax" + "github.com/gorilla/websocket" + + "tangled.org/core/api/tangled" + "tangled.org/core/notifier" + "tangled.org/core/spindle/config" + "tangled.org/core/spindle/db" + "tangled.org/core/spindle/engine" + "tangled.org/core/spindle/models" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +const ( + dialBackoffMin = 1 * time.Second + dialBackoffMax = 30 * time.Second + snapshotEvery = 15 * time.Second + defaultSeats = 4 +) + +type Executor struct { + millURL string + secret string + nodeID string + seats int + + engines map[string]models.Engine + db *db.DB + n *notifier.Notifier + cfg *config.Config + l *slog.Logger + + relay *relayLog + relayMu sync.Mutex // offset assign + send must be serialized so wire order == offset order + + connMu sync.Mutex + enc messageEncoder + + mu sync.Mutex + active map[string]*reservation + draining bool + activeCount int +} + +type reservation struct { + leaseID string + wid models.WorkflowId + realEngine models.Engine + slot engine.WorkflowSlot + wf *models.Workflow + repoDid syntax.DID + + committed bool + cancelled bool + cancel context.CancelFunc + ttlTimer *time.Timer + stopTail func() +} + +type messageEncoder interface { + Encode(*millproto.Message) error +} + +type reservationCleanup struct { + stopTail func() + slot engine.WorkflowSlot +} + +func (c reservationCleanup) run() { + if c.stopTail != nil { + c.stopTail() + } + if c.slot != nil { + c.slot.Release() + } +} + +func New(cfg *config.Config, engines map[string]models.Engine, d *db.DB, n *notifier.Notifier, l *slog.Logger) *Executor { + seats := cfg.Mill.Seats + if seats <= 0 { + seats = defaultSeats + } + return &Executor{ + millURL: cfg.Mill.URL, + secret: cfg.Mill.SharedSecret, + nodeID: cfg.Server.Hostname, + seats: seats, + engines: engines, + db: d, + n: n, + cfg: cfg, + l: l.With("component", "mill.executor"), + relay: newRelayLog(), + active: make(map[string]*reservation), + } +} + +func (e *Executor) Connect(ctx context.Context) { + sub := e.n.Subscribe() + cursor, err := e.db.EventHighWater() + if err != nil { + e.n.Unsubscribe(sub) + e.l.Error("establish event cursor failed", "err", err) + return + } + e.drainEvents(&cursor) + go e.observeLoop(ctx, sub, cursor) + defer e.n.Unsubscribe(sub) + + backoff := dialBackoffMin + for { + if ctx.Err() != nil { + return + } + err := e.runSession(ctx) + if ctx.Err() != nil { + return + } + e.l.Warn("mill session ended; reconnecting", "err", err, "backoff", backoff) + select { + case <-ctx.Done(): + return + case <-time.After(backoff): + } + backoff = min(backoff*2, dialBackoffMax) + } +} + +func (e *Executor) runSession(ctx context.Context) error { + header := http.Header{} + if e.secret != "" { + header.Set("Authorization", "Bearer "+e.secret) + } + conn, _, err := websocket.DefaultDialer.DialContext(ctx, e.millURL, header) + if err != nil { + return fmt.Errorf("dial mill: %w", err) + } + defer conn.Close() + + sessionCtx, cancelSession := context.WithCancel(ctx) + defer cancelSession() + stopClose := context.AfterFunc(sessionCtx, func() { _ = conn.Close() }) + defer stopClose() + + stream := millproto.NewWSStream(conn) + enc := millproto.NewEncoder(stream) + dec := millproto.NewDecoder(stream) + + hello := &millproto.Message{Hello: &millv1.Hello{ + ProtocolVersion: millproto.ProtocolVersion, + NodeId: e.nodeID, + Engines: e.engineNames(), + Arch: runtime.GOARCH, + LastOffset: e.relay.lastOffset(), + }} + if err := enc.Encode(hello); err != nil { + return fmt.Errorf("send hello: %w", err) + } + + resumeMsg, err := dec.Decode() + if err != nil { + return fmt.Errorf("read resume: %w", err) + } + resume := resumeMsg.GetResume() + if resume == nil { + return fmt.Errorf("expected resume, got something else") + } + + // install encoder and replay anything the mill missed. + e.connMu.Lock() + e.enc = enc + e.connMu.Unlock() + defer func() { + e.connMu.Lock() + e.enc = nil + e.connMu.Unlock() + }() + + readErr := make(chan error, 1) + go func() { + for { + msg, err := dec.Decode() + if err != nil { + readErr <- fmt.Errorf("read: %w", err) + return + } + e.dispatch(sessionCtx, msg) + } + }() + + e.replay(resume.GetAckOffset()) + e.pushSnapshot() + e.l.Info("connected to mill", "node", e.nodeID, "resumeFrom", resume.GetAckOffset()) + + go e.snapshotLoop(sessionCtx, enc) + return <-readErr +} + +func (e *Executor) replay(ackOffset uint64) { + e.relayMu.Lock() + defer e.relayMu.Unlock() + e.relay.ack(ackOffset) + for _, msg := range e.relay.since(ackOffset) { + e.send(msg) + } +} + +func (e *Executor) send(msg *millproto.Message) { + e.connMu.Lock() + enc := e.enc + e.connMu.Unlock() + if enc != nil { + _ = enc.Encode(msg) + } +} + +func (e *Executor) dispatch(ctx context.Context, msg *millproto.Message) { + switch { + case msg.GetReserveSeat() != nil: + e.handleReserve(ctx, msg.GetReserveSeat()) + case msg.GetCommitLease() != nil: + e.handleCommit(ctx, msg.GetCommitLease()) + case msg.GetReleaseLease() != nil: + e.handleRelease(msg.GetReleaseLease().GetLeaseId()) + case msg.GetCancelAttempt() != nil: + e.handleCancel(msg.GetCancelAttempt().GetLeaseId()) + case msg.GetAck() != nil: + e.relay.ack(msg.GetAck().GetUpToOffset()) + default: + e.l.Warn("executor received unexpected message") + } +} + +// --- reserve / commit / release / cancel ---------------------------------- + +func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { + reject := func(reason string) { + e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ + LeaseId: rs.GetLeaseId(), + Accepted: false, + RejectReason: reason, + }}) + } + + e.mu.Lock() + draining := e.draining + e.mu.Unlock() + if draining { + reject("draining") + return + } + + realEngine, ok := e.engines[rs.GetTargetEngine()] + if !ok { + reject("unknown engine " + rs.GetTargetEngine()) + return + } + slotter, ok := realEngine.(engine.WorkflowSlotter) + if !ok { + reject("engine does not support workflow slots") + return + } + + var twf tangled.Pipeline_Workflow + if err := json.Unmarshal([]byte(rs.GetRawWorkflowJson()), &twf); err != nil { + reject("bad workflow json") + return + } + var tpl tangled.Pipeline + if err := json.Unmarshal([]byte(rs.GetRawPipelineJson()), &tpl); err != nil { + reject("bad pipeline json") + return + } + + pipelineId := models.PipelineId{Knot: rs.GetKnot(), Rkey: rs.GetRkey()} + wid := models.WorkflowId{PipelineId: pipelineId, Name: twf.Name} + + wf, err := realEngine.InitWorkflow(twf, tpl) + if err != nil { + reject("init workflow: " + err.Error()) + return + } + // the job skipped processPipeline, so inject TANGLED_* env here. + if wf.Environment == nil { + wf.Environment = make(map[string]string) + } + maps.Copy(wf.Environment, models.PipelineEnvVars(tpl.TriggerMetadata, pipelineId)) + + // NoWait: the executor doesn't queue locally; the mill owns the backlog. + slot, err := slotter.AcquireWorkflowSlot(ctx, wid, wf, engine.NoWait) + if err != nil { + reject(err.Error()) + return + } + + var repoDid syntax.DID + if tpl.TriggerMetadata != nil && tpl.TriggerMetadata.Repo != nil && tpl.TriggerMetadata.Repo.RepoDid != nil { + repoDid, _ = syntax.ParseDID(*tpl.TriggerMetadata.Repo.RepoDid) + } + + res := &reservation{ + leaseID: rs.GetLeaseId(), + wid: wid, + realEngine: realEngine, + slot: slot, + wf: wf, + repoDid: repoDid, + } + res.ttlTimer = time.AfterFunc(ttlDuration(rs.GetTtlSeconds()), func() { e.expireReservation(res.leaseID) }) + + e.mu.Lock() + e.active[res.leaseID] = res + e.activeCount++ + e.mu.Unlock() + + e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ + LeaseId: rs.GetLeaseId(), + Accepted: true, + Score: e.bidScore(), + }}) + e.pushSnapshot() +} + +func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) { + e.mu.Lock() + res := e.active[cl.GetLeaseId()] + if res == nil { + e.mu.Unlock() + e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ + LeaseId: cl.GetLeaseId(), + Accepted: false, + RejectReason: "reservation missing or expired", + }}) + return + } + if res.committed { + e.mu.Unlock() + e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: cl.GetLeaseId()}}) + return + } + res.committed = true + if res.ttlTimer != nil { + res.ttlTimer.Stop() + } + jobCtx, cancel := context.WithCancel(ctx) + res.cancel = cancel + e.mu.Unlock() + + vault := newMemVault(cl.GetSecrets()) + re := newReservedEngine(res.realEngine, res.slot) + pipeline := &models.Pipeline{ + RepoDid: res.repoDid, + Workflows: map[models.Engine][]models.Workflow{re: {*res.wf}}, + TrustedSource: true, + } + + e.startTail(res) + + go engine.StartWorkflows(e.l, vault, e.cfg, e.db, e.n, jobCtx, pipeline, res.wid.PipelineId) + + e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: cl.GetLeaseId()}}) +} + +func (e *Executor) handleRelease(leaseID string) { + cleanup, ok := e.takeUncommittedReservation(leaseID, true) + if !ok { + return + } + cleanup.run() + e.pushSnapshot() +} + +func (e *Executor) handleCancel(leaseID string) { + e.mu.Lock() + res := e.active[leaseID] + if res == nil { + e.mu.Unlock() + return + } + res.cancelled = true + cancel := res.cancel + committed := res.committed + var cleanup reservationCleanup + if !committed { + cleanup = e.removeReservationLocked(res, true) + } + e.mu.Unlock() + + if cancel != nil { + cancel() + } + // committed jobs clean up when the observe loop sees the terminal row; an + // uncommitted cancel still needs the slot released. + if !committed { + cleanup.run() + e.pushSnapshot() + } +} + +func (e *Executor) expireReservation(leaseID string) { + cleanup, ok := e.takeUncommittedReservation(leaseID, true) + if !ok { + return + } + e.l.Warn("reservation expired before commit", "lease", leaseID) + cleanup.run() + e.pushSnapshot() +} + +func (e *Executor) takeUncommittedReservation(leaseID string, releaseSlot bool) (reservationCleanup, bool) { + e.mu.Lock() + defer e.mu.Unlock() + res := e.active[leaseID] + if res == nil || res.committed { + return reservationCleanup{}, false + } + return e.removeReservationLocked(res, releaseSlot), true +} + +func (e *Executor) finishReservation(res *reservation) (reservationCleanup, bool, bool) { + e.mu.Lock() + defer e.mu.Unlock() + if e.active[res.leaseID] != res { + return reservationCleanup{}, false, false + } + cancelled := res.cancelled + return e.removeReservationLocked(res, false), cancelled, true +} + +func (e *Executor) removeReservationLocked(res *reservation, releaseSlot bool) reservationCleanup { + delete(e.active, res.leaseID) + if e.activeCount > 0 { + e.activeCount-- + } + if res.ttlTimer != nil { + res.ttlTimer.Stop() + res.ttlTimer = nil + } + cleanup := reservationCleanup{stopTail: res.stopTail} + res.stopTail = nil + if releaseSlot { + cleanup.slot = res.slot + res.slot = nil + } + return cleanup +} + +// --- snapshots ------------------------------------------------------------- + +func (e *Executor) snapshotLoop(ctx context.Context, enc *millproto.Encoder) { + t := time.NewTicker(snapshotEvery) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + // stop pushing once this connection is replaced. + e.connMu.Lock() + cur := e.enc + e.connMu.Unlock() + if cur != enc { + return + } + e.pushSnapshot() + } + } +} + +func (e *Executor) pushSnapshot() { + e.mu.Lock() + free := uint32(0) + if !e.draining { + if n := e.seats - e.activeCount; n > 0 { + free = uint32(n) + } + } + e.mu.Unlock() + + engines := make(map[string]*millv1.EngineSnapshot, len(e.engines)) + for name := range e.engines { + engines[name] = &millv1.EngineSnapshot{FreeSeats: free} + } + e.send(&millproto.Message{NodeSnapshot: &millv1.NodeSnapshot{ + NodeId: e.nodeID, + Engines: engines, + }}) +} + +func (e *Executor) Drain() { + e.mu.Lock() + e.draining = true + e.mu.Unlock() + e.pushSnapshot() +} + +func (e *Executor) bidScore() float64 { + e.mu.Lock() + defer e.mu.Unlock() + return float64(e.seats - e.activeCount) +} + +func (e *Executor) engineNames() []string { + names := make([]string, 0, len(e.engines)) + for name := range e.engines { + names = append(names, name) + } + return names +} + +func ttlDuration(secs uint32) time.Duration { + if secs == 0 { + return defaultReservationTTL + } + return time.Duration(secs) * time.Second +} + +const defaultReservationTTL = 60 * time.Second diff --git a/spindle/mill/executor/observe.go b/spindle/mill/executor/observe.go new file mode 100644 --- /dev/null +++ b/spindle/mill/executor/observe.go @@ -0,0 +1,184 @@ +package executor + +import ( + "context" + "io" + "sync" + "time" + + "github.com/hpcloud/tail" + + "tangled.org/core/api/tangled" + "tangled.org/core/spindle/models" + "tangled.org/core/spindle/secrets" + + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +// observeLoop watches the executor's own eventstream and relays the status rows +// belonging to active leases. This is the "observe-by-cursor" relay: the jobs +// write status to the real db/notifier exactly as a standalone spindle would, +// and we forward what is relevant without touching the execution path. +func (e *Executor) observeLoop(ctx context.Context, sub <-chan struct{}, cursor int64) { + ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-sub: + case <-ticker.C: + } + e.drainEvents(&cursor) + } +} + +func (e *Executor) drainEvents(cursor *int64) { + for { + evs, err := e.db.GetEvents(*cursor, 100) + if err != nil { + e.l.Error("observe GetEvents failed", "err", err) + return + } + for _, ev := range evs { + *cursor = ev.Created + if ev.Nsid != tangled.PipelineStatusNSID { + continue + } + st, ok := parseStatus(ev.EventJson) + if !ok { + continue + } + e.onStatusRow(st) + } + if len(evs) < 100 { + return + } + } +} + +func (e *Executor) onStatusRow(st *tangled.PipelineStatus) { + res := e.reservationFor(st.Pipeline, st.Workflow) + if res == nil { + return + } + + kind := models.StatusKind(st.Status) + switch { + case kind.IsFinish(): + e.finishJob(res, st) + case kind == models.StatusKindRunning: + e.relayMu.Lock() + e.send(e.relay.appendStatus(res.leaseID, st, st.Pipeline)) + e.relayMu.Unlock() + } +} + +// finishJob flushes the log tail, then relays the terminal as an AttemptResult +// (which the mill turns into the terminal status row), then forgets the lease. +// Tail flush happens first so all log offsets precede the terminal offset. +func (e *Executor) finishJob(res *reservation, st *tangled.PipelineStatus) { + cleanup, cancelled, ok := e.finishReservation(res) + if !ok { + return + } + cleanup.run() + + e.relayMu.Lock() + terminalStatus := st.Status + if cancelled { + terminalStatus = string(models.StatusKindCancelled) + } + e.send(e.relay.appendTerminal(res.leaseID, terminalStatus, st)) + e.relayMu.Unlock() + + e.pushSnapshot() +} + +// reservationFor finds the active reservation whose workflow matches a status +// row, by pipeline aturi + workflow name. Active reservations are few (bounded +// by seats), so a scan is cheaper than maintaining a parallel index. +func (e *Executor) reservationFor(pipelineAturi, workflow string) *reservation { + e.mu.Lock() + defer e.mu.Unlock() + for _, res := range e.active { + if string(res.wid.PipelineId.AtUri()) == pipelineAturi && res.wid.Name == workflow { + return res + } + } + return nil +} + +func (e *Executor) relayLogLine(leaseID, text string) { + e.relayMu.Lock() + defer e.relayMu.Unlock() + e.send(e.relay.appendLog(leaseID, []byte(text))) +} + +// startTail follows the job's local log file and relays each line. The lines are +// already-encoded models.LogLine JSON, so we forward them verbatim. +func (e *Executor) startTail(res *reservation) { + path := models.LogFilePath(e.cfg.Server.LogDir, res.wid) + t, err := tail.TailFile(path, tail.Config{ + Follow: true, + ReOpen: true, + MustExist: false, + Location: &tail.SeekInfo{Offset: 0, Whence: io.SeekStart}, + Logger: tail.DiscardingLogger, + }) + if err != nil { + e.l.Error("tail log file failed", "wid", res.wid, "err", err) + return + } + + done := make(chan struct{}) + go func() { + defer close(done) + for line := range t.Lines { + if line == nil || line.Err != nil { + continue + } + e.relayLogLine(res.leaseID, line.Text) + } + }() + + var once sync.Once + res.stopTail = func() { + once.Do(func() { + // Stop() blocks until tailing ends and t.Lines is closed, so the + // consumer goroutine then drains the remaining (finite) buffered + // lines and closes done. Waiting on done in full guarantees every + // log line is relayed before finishJob sends the terminal — without + // it a slow drain lets the terminal overtake a log line, which the + // mill then drops against its sealed log file. + _ = t.Stop() + <-done + }) + } +} + +// memVault is a temporary in-memory secrets manager holding the secrets the +// mill handed over at commit. They never touch disk on the executor. +type memVault struct { + secrets []secrets.UnlockedSecret +} + +func newMemVault(pb []*millv1.Secret) *memVault { + s := make([]secrets.UnlockedSecret, len(pb)) + for i, x := range pb { + s[i] = secrets.UnlockedSecret{Key: x.GetKey(), Value: x.GetValue()} + } + return &memVault{secrets: s} +} + +func (v *memVault) GetSecretsUnlocked(ctx context.Context, repo secrets.RepoIdentifier) ([]secrets.UnlockedSecret, error) { + return v.secrets, nil +} +func (v *memVault) GetSecretsLocked(ctx context.Context, repo secrets.RepoIdentifier) ([]secrets.LockedSecret, error) { + return nil, nil +} +func (v *memVault) AddSecret(ctx context.Context, s secrets.UnlockedSecret) error { return nil } +func (v *memVault) RemoveSecret(ctx context.Context, s secrets.Secret[any]) error { return nil } + +var _ secrets.Manager = (*memVault)(nil) diff --git a/spindle/mill/executor/relay.go b/spindle/mill/executor/relay.go new file mode 100644 --- /dev/null +++ b/spindle/mill/executor/relay.go @@ -0,0 +1,152 @@ +package executor + +import ( + "encoding/json" + "sync" + + "tangled.org/core/api/tangled" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +// relayLog is the per-session sequential event log the executor keeps so it can +// replay anything the mill missed across a reconnect. Only mill-facing relay +// messages (StatusEvent / LogLine / AttemptResult) get offsets and are buffered; +// control messages do not. +type relayLog struct { + mu sync.Mutex + entries []*millproto.Message + next uint64 // next offset to assign + acked uint64 // highest offset the mill has acked +} + +func newRelayLog() *relayLog { + return &relayLog{next: 1} +} + +func (rl *relayLog) appendStatus(leaseID string, st *tangled.PipelineStatus, pipelineAturi string) *millproto.Message { + rl.mu.Lock() + defer rl.mu.Unlock() + offset := rl.next + rl.next++ + var exit int64 + if st.ExitCode != nil { + exit = *st.ExitCode + } + var errStr string + if st.Error != nil { + errStr = *st.Error + } + msg := &millproto.Message{StatusEvent: &millv1.StatusEvent{ + Offset: offset, + LeaseId: leaseID, + Status: st.Status, + Error: errStr, + ExitCode: exit, + Workflow: st.Workflow, + PipelineAturi: pipelineAturi, + }} + rl.entries = append(rl.entries, msg) + return msg +} + +func (rl *relayLog) appendLog(leaseID string, raw []byte) *millproto.Message { + rl.mu.Lock() + defer rl.mu.Unlock() + offset := rl.next + rl.next++ + msg := &millproto.Message{LogLine: &millv1.LogLine{ + Offset: offset, + LeaseId: leaseID, + RawJson: raw, + }} + rl.entries = append(rl.entries, msg) + return msg +} + +func (rl *relayLog) appendTerminal(leaseID, status string, st *tangled.PipelineStatus) *millproto.Message { + rl.mu.Lock() + defer rl.mu.Unlock() + offset := rl.next + rl.next++ + var exit int64 + var errStr string + if st != nil { + if st.ExitCode != nil { + exit = *st.ExitCode + } + if st.Error != nil { + errStr = *st.Error + } + } + msg := &millproto.Message{AttemptResult: &millv1.AttemptResult{ + Offset: offset, + LeaseId: leaseID, + TerminalStatus: status, + Error: errStr, + ExitCode: exit, + }} + rl.entries = append(rl.entries, msg) + return msg +} + +func (rl *relayLog) ack(upTo uint64) { + rl.mu.Lock() + defer rl.mu.Unlock() + if upTo > rl.acked { + rl.acked = upTo + } + kept := rl.entries[:0] + for _, e := range rl.entries { + if offsetOf(e) > rl.acked { + kept = append(kept, e) + } + } + // nil the drained tail so acked messages can be GC'd (kept aliases the same + // backing array as entries). + for i := len(kept); i < len(rl.entries); i++ { + rl.entries[i] = nil + } + rl.entries = kept +} + +// since returns buffered entries with offset strictly greater than offset, for +// replay on reconnect. +func (rl *relayLog) since(offset uint64) []*millproto.Message { + rl.mu.Lock() + defer rl.mu.Unlock() + var out []*millproto.Message + for _, e := range rl.entries { + if offsetOf(e) > offset { + out = append(out, e) + } + } + return out +} + +func (rl *relayLog) lastOffset() uint64 { + rl.mu.Lock() + defer rl.mu.Unlock() + return rl.next - 1 +} + +func offsetOf(m *millproto.Message) uint64 { + switch { + case m.GetStatusEvent() != nil: + return m.GetStatusEvent().GetOffset() + case m.GetLogLine() != nil: + return m.GetLogLine().GetOffset() + case m.GetAttemptResult() != nil: + return m.GetAttemptResult().GetOffset() + } + return 0 +} + +func parseStatus(raw json.RawMessage) (*tangled.PipelineStatus, bool) { + var st tangled.PipelineStatus + if err := json.Unmarshal(raw, &st); err != nil { + return nil, false + } + return &st, true +} diff --git a/spindle/mill/executor/reserved.go b/spindle/mill/executor/reserved.go new file mode 100644 --- /dev/null +++ b/spindle/mill/executor/reserved.go @@ -0,0 +1,66 @@ +package executor + +import ( + "context" + "fmt" + "sync" + "time" + + "tangled.org/core/api/tangled" + "tangled.org/core/spindle/engine" + "tangled.org/core/spindle/models" + "tangled.org/core/spindle/secrets" +) + +// reservedEngine wraps a real engine so that the slot acquired up-front during +// ReserveSeat is the slot StartWorkflows gets, instead of acquiring a second +// one. Every other method delegates straight to the real engine, so the +// execution path runs byte-for-byte as it would standalone. +// +// This is the only change to the execution path on an executor. +type reservedEngine struct { + inner models.Engine + slot engine.WorkflowSlot + once sync.Once +} + +// newReservedEngine returns a wrapper around inner that hands back slot exactly +// once from AcquireWorkflowSlot. +func newReservedEngine(inner models.Engine, slot engine.WorkflowSlot) models.Engine { + return &reservedEngine{inner: inner, slot: slot} +} + +func (e *reservedEngine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) { + return e.inner.InitWorkflow(twf, tpl) +} + +func (e *reservedEngine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, wfLogger models.WorkflowLogger) error { + return e.inner.SetupWorkflow(ctx, wid, wf, wfLogger) +} + +func (e *reservedEngine) WorkflowTimeout() time.Duration { + return e.inner.WorkflowTimeout() +} + +func (e *reservedEngine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error { + return e.inner.DestroyWorkflow(ctx, wid) +} + +func (e *reservedEngine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, secrets []secrets.UnlockedSecret, wfLogger models.WorkflowLogger) error { + return e.inner.RunStep(ctx, wid, w, idx, secrets, wfLogger) +} + +// AcquireWorkflowSlot hands back the pre-acquired slot exactly once. The slot +// was already obtained (and resource-accounted) during ReserveSeat, so a second +// acquire would double-count. +func (e *reservedEngine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, _ engine.AcquireMode) (engine.WorkflowSlot, error) { + var slot engine.WorkflowSlot + e.once.Do(func() { + slot = e.slot + e.slot = nil + }) + if slot == nil { + return nil, fmt.Errorf("reserved slot already consumed") + } + return slot, nil +} diff --git a/spindle/mill/executor/reserved_test.go b/spindle/mill/executor/reserved_test.go new file mode 100644 --- /dev/null +++ b/spindle/mill/executor/reserved_test.go @@ -0,0 +1,371 @@ +package executor + +import ( + "context" + "fmt" + "io" + "log/slog" + "net/http" + "net/http/httptest" + "path/filepath" + "strings" + "testing" + "time" + + "github.com/bluesky-social/indigo/atproto/syntax" + "github.com/gorilla/websocket" + + "tangled.org/core/api/tangled" + "tangled.org/core/notifier" + "tangled.org/core/spindle/config" + "tangled.org/core/spindle/db" + "tangled.org/core/spindle/engine" + "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" + "tangled.org/core/spindle/models" + "tangled.org/core/spindle/secrets" +) + +type captureEncoder struct { + messages chan *millproto.Message +} + +func newCaptureEncoder() *captureEncoder { + return &captureEncoder{messages: make(chan *millproto.Message, 4)} +} + +func (e *captureEncoder) Encode(msg *millproto.Message) error { + e.messages <- msg + return nil +} + +type fakeSlot struct{ released int } + +func (s *fakeSlot) Release() { s.released++ } + +// fakeEngine records which methods got called and never acquires a slot of its +// own; the wrapper must supply the slot. +type fakeEngine struct { + setupCalled bool + runCalled bool + destroyCalled bool + acquireCalled bool + secrets chan []secrets.UnlockedSecret + done chan struct{} +} + +func (e *fakeEngine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) { + return &models.Workflow{Name: twf.Name}, nil +} +func (e *fakeEngine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, l models.WorkflowLogger) error { + e.setupCalled = true + return nil +} +func (e *fakeEngine) WorkflowTimeout() time.Duration { return 7 * time.Minute } +func (e *fakeEngine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error { + e.destroyCalled = true + if e.done != nil { + close(e.done) + } + return nil +} +func (e *fakeEngine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, s []secrets.UnlockedSecret, l models.WorkflowLogger) error { + e.runCalled = true + if e.secrets != nil { + e.secrets <- s + } + return nil +} + +func (e *fakeEngine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, mode engine.AcquireMode) (engine.WorkflowSlot, error) { + e.acquireCalled = true + return engine.NoopSlot{}, nil +} + +type fakeStep struct{} + +func (fakeStep) Name() string { return "test" } +func (fakeStep) Command() string { return "true" } +func (fakeStep) Kind() models.StepKind { return models.StepKindUser } +func TestReservedEngineHandsBackHeldSlotOnce(t *testing.T) { + inner := &fakeEngine{} + slot := &fakeSlot{} + re := newReservedEngine(inner, slot) + + got, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait) + if err != nil { + t.Fatalf("AcquireWorkflowSlot() error = %v", err) + } + if got != engine.WorkflowSlot(slot) { + t.Fatal("AcquireWorkflowSlot() did not return the held slot") + } + if inner.acquireCalled { + t.Fatal("wrapper must not call the inner engine's AcquireWorkflowSlot") + } + + if _, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait); err == nil { + t.Fatal("second AcquireWorkflowSlot() should error") + } +} + +func TestReservedEngineForwardsLifecycle(t *testing.T) { + inner := &fakeEngine{} + re := newReservedEngine(inner, &fakeSlot{}) + + if d := re.WorkflowTimeout(); d != 7*time.Minute { + t.Fatalf("WorkflowTimeout() = %v, want 7m", d) + } + if err := re.SetupWorkflow(context.Background(), models.WorkflowId{}, &models.Workflow{}, models.NullLogger{}); err != nil { + t.Fatal(err) + } + if err := re.RunStep(context.Background(), models.WorkflowId{}, &models.Workflow{Steps: []models.Step{}}, 0, nil, models.NullLogger{}); err != nil { + t.Fatal(err) + } + if err := re.DestroyWorkflow(context.Background(), models.WorkflowId{}); err != nil { + t.Fatal(err) + } + if !inner.setupCalled || !inner.runCalled || !inner.destroyCalled { + t.Fatalf("lifecycle not forwarded: %+v", inner) + } +} + +func TestHandleCommitIsIdempotent(t *testing.T) { + enc := newCaptureEncoder() + + e := &Executor{ + l: slog.New(slog.NewTextHandler(io.Discard, nil)), + enc: enc, + active: map[string]*reservation{"lease-1": {leaseID: "lease-1", committed: true}}, + } + + e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "lease-1"}) + msg := <-enc.messages + if got := msg.GetCommitted().GetLeaseId(); got != "lease-1" { + t.Fatalf("Committed lease = %q, want lease-1", got) + } +} + +func TestHandleCommitRejectsMissingReservation(t *testing.T) { + enc := newCaptureEncoder() + e := &Executor{ + l: slog.New(slog.NewTextHandler(io.Discard, nil)), + enc: enc, + active: make(map[string]*reservation), + } + + e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "expired"}) + result := (<-enc.messages).GetReserveResult() + if result == nil { + t.Fatal("missing reservation commit did not receive a ReserveResult") + } + if result.GetLeaseId() != "expired" || result.GetAccepted() { + t.Fatalf("ReserveResult = %+v, want correlated rejection", result) + } +} + +func TestHandleCommitPreservesPreauthorizedSecrets(t *testing.T) { + d, err := db.Make(context.Background(), filepath.Join(t.TempDir(), "spindle.db")) + if err != nil { + t.Fatalf("db.Make() error = %v", err) + } + t.Cleanup(func() { d.Close() }) + n := notifier.New() + enc := newCaptureEncoder() + inner := &fakeEngine{ + secrets: make(chan []secrets.UnlockedSecret, 1), + done: make(chan struct{}), + } + repoDid, err := syntax.ParseDID("did:plc:abcdefghijklmnopqrstuvwxyz") + if err != nil { + t.Fatalf("ParseDID() error = %v", err) + } + res := &reservation{ + leaseID: "lease-1", + wid: models.WorkflowId{PipelineId: models.PipelineId{Knot: "did:plc:test", Rkey: "pipeline"}, Name: "build"}, + realEngine: inner, + slot: engine.NoopSlot{}, + wf: &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}}, + repoDid: repoDid, + } + e := &Executor{ + cfg: &config.Config{Server: config.Server{LogDir: t.TempDir()}}, + db: d, + n: &n, + l: slog.New(slog.NewTextHandler(io.Discard, nil)), + enc: enc, + relay: newRelayLog(), + active: map[string]*reservation{res.leaseID: res}, + } + + e.handleCommit(context.Background(), &millv1.CommitLease{ + LeaseId: res.leaseID, + Secrets: []*millv1.Secret{{Key: "TOKEN", Value: "secret-value"}}, + }) + if got := (<-enc.messages).GetCommitted().GetLeaseId(); got != res.leaseID { + t.Fatalf("Committed lease = %q, want %q", got, res.leaseID) + } + select { + case got := <-inner.secrets: + if len(got) != 1 || got[0].Key != "TOKEN" || got[0].Value != "secret-value" { + t.Fatalf("RunStep secrets = %+v", got) + } + case <-time.After(2 * time.Second): + t.Fatal("RunStep did not receive CommitLease secrets") + } + select { + case <-inner.done: + case <-time.After(2 * time.Second): + t.Fatal("workflow did not finish") + } + if res.stopTail != nil { + res.stopTail() + } +} + +func TestRunSessionCancellationClosesStalledWebsocket(t *testing.T) { + connected := make(chan struct{}) + release := make(chan struct{}) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Upgrade(w, r, nil, 1024, 1024) + if err != nil { + return + } + defer conn.Close() + close(connected) + <-release + })) + t.Cleanup(func() { + close(release) + srv.Close() + }) + + e := testSessionExecutor("ws" + strings.TrimPrefix(srv.URL, "http")) + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { done <- e.runSession(ctx) }() + <-connected + cancel() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("runSession did not return after context cancellation") + } +} + +func TestRunSessionConsumesAcksWhileReplaying(t *testing.T) { + e := testSessionExecutor("") + e.relay.appendLog("lease", []byte("first")) + e.relay.appendLog("lease", make([]byte, 7<<20)) + + serverErr := make(chan error, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Upgrade(w, r, nil, 1024, 1024) + if err != nil { + serverErr <- err + return + } + defer conn.Close() + if tcp, ok := conn.UnderlyingConn().(interface{ SetReadBuffer(int) error }); ok { + _ = tcp.SetReadBuffer(1024) + } + stream := millproto.NewWSStream(conn) + enc := millproto.NewEncoder(stream) + dec := millproto.NewDecoder(stream) + if _, err := dec.Decode(); err != nil { + serverErr <- err + return + } + if err := enc.Encode(&millproto.Message{Resume: &millv1.Resume{}}); err != nil { + serverErr <- err + return + } + first, err := dec.Decode() + if err != nil { + serverErr <- err + return + } + if got := offsetOf(first); got != 1 { + serverErr <- fmt.Errorf("first replay offset = %d, want 1", got) + return + } + if err := enc.Encode(&millproto.Message{Ack: &millv1.Ack{UpToOffset: 1}}); err != nil { + serverErr <- err + return + } + + deadline := time.Now().Add(10 * time.Second) + for { + e.relay.mu.Lock() + acked := e.relay.acked + e.relay.mu.Unlock() + if acked >= 1 { + break + } + if time.Now().After(deadline) { + serverErr <- context.DeadlineExceeded + return + } + time.Sleep(time.Millisecond) + } + + serverErr <- nil + })) + t.Cleanup(srv.Close) + e.millURL = "ws" + strings.TrimPrefix(srv.URL, "http") + + done := make(chan error, 1) + go func() { done <- e.runSession(context.Background()) }() + select { + case err := <-serverErr: + if err != nil { + t.Fatalf("server: %v", err) + } + case <-time.After(20 * time.Second): + t.Fatal("replay stalled while the peer waited for its acknowledgement to be consumed") + } + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("runSession did not stop after peer closed") + } +} + +func testSessionExecutor(url string) *Executor { + return &Executor{ + millURL: url, + nodeID: "node-1", + seats: 1, + engines: map[string]models.Engine{"dummy": &fakeEngine{}}, + cfg: &config.Config{Server: config.Server{Dev: true}}, + l: slog.New(slog.NewTextHandler(io.Discard, nil)), + relay: newRelayLog(), + active: make(map[string]*reservation), + } +} + +func TestFinishJobReportsCancelledReservationAsCancelled(t *testing.T) { + res := &reservation{ + leaseID: "lease-1", + wid: models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}, + cancelled: true, + } + e := &Executor{ + relay: newRelayLog(), + active: map[string]*reservation{res.leaseID: res}, + } + + e.finishJob(res, &tangled.PipelineStatus{ + Pipeline: string(res.wid.PipelineId.AtUri()), + Workflow: res.wid.Name, + Status: string(models.StatusKindFailed), + }) + + if len(e.relay.entries) != 1 { + t.Fatalf("relay entries = %d, want 1", len(e.relay.entries)) + } + got := e.relay.entries[0].GetAttemptResult().GetTerminalStatus() + if got != string(models.StatusKindCancelled) { + t.Fatalf("terminal status = %q, want cancelled", got) + } +}