spindle/mill: implement mill host · mokkenstorm.dev/core@3b25bda

+41 -15

spindle/db/events.go

··· 32 32 return d.insertEvent(event, n) 33 33 } 34 34 35 - func (d *DB) createStatusEvent( 36 - workflowId models.WorkflowId, 37 - statusKind models.StatusKind, 38 - workflowError *string, 39 - exitCode *int64, 40 - n *notifier.Notifier, 41 - ) error { 42 - now := time.Now() 43 - pipelineAtUri := workflowId.PipelineId.AtUri() 35 + // statusEvent builds the eventstream row for a pipeline status. Created is left 36 + // zero so insertEvent stamps the local clock. 37 + func statusEvent(pipelineAtUri, workflow, status string, workflowError *string, exitCode *int64) (eventstream.Event, error) { 44 38 s := tangled.PipelineStatus{ 45 - CreatedAt: now.Format(time.RFC3339), 39 + CreatedAt: time.Now().Format(time.RFC3339), 46 40 Error: workflowError, 47 41 ExitCode: exitCode, 48 - Pipeline: string(pipelineAtUri), 49 - Workflow: workflowId.Name, 50 - Status: string(statusKind), 42 + Pipeline: pipelineAtUri, 43 + Workflow: workflow, 44 + Status: status, 51 45 } 52 46 53 47 eventJson, err := json.Marshal(s) 54 48 if err != nil { 55 - return err 49 + return eventstream.Event{}, err 56 50 } 57 51 58 - event := eventstream.Event{ 52 + return eventstream.Event{ 59 53 Rkey: tid.TID(), 60 54 Nsid: tangled.PipelineStatusNSID, 61 55 EventJson: eventJson, 56 + }, nil 57 + } 58 + 59 + func (d *DB) createStatusEvent( 60 + workflowId models.WorkflowId, 61 + statusKind models.StatusKind, 62 + workflowError *string, 63 + exitCode *int64, 64 + n *notifier.Notifier, 65 + ) error { 66 + event, err := statusEvent(string(workflowId.PipelineId.AtUri()), workflowId.Name, string(statusKind), workflowError, exitCode) 67 + if err != nil { 68 + return err 62 69 } 70 + return d.insertEvent(event, n) 71 + } 63 72 73 + // InsertRelayedStatus re-authors a status row relayed from an executor into the 74 + // mill's own eventstream. The mill assigns its own Created timestamp (via 75 + // insertEvent passing Created=0) so relayed rows order correctly against the 76 + // mill's cursor, and NotifyAll wakes the mill's /events subscribers exactly 77 + // as a local status write would. 78 + func (d *DB) InsertRelayedStatus( 79 + pipelineAtUri string, 80 + workflow string, 81 + status string, 82 + workflowError *string, 83 + exitCode *int64, 84 + n *notifier.Notifier, 85 + ) error { 86 + event, err := statusEvent(pipelineAtUri, workflow, status, workflowError, exitCode) 87 + if err != nil { 88 + return err 89 + } 64 90 return d.insertEvent(event, n) 65 91 } 66 92

+240

spindle/mill/auth_test.go

··· 1 + package mill 2 + 3 + import ( 4 + "context" 5 + "io" 6 + "log/slog" 7 + "os" 8 + "path/filepath" 9 + "testing" 10 + "time" 11 + 12 + "tangled.org/core/notifier" 13 + "tangled.org/core/spindle/db" 14 + "tangled.org/core/spindle/models" 15 + 16 + millproto "tangled.org/core/spindle/mill/proto" 17 + millv1 "tangled.org/core/spindle/mill/proto/gen" 18 + ) 19 + 20 + func discardLogger() *slog.Logger { 21 + return slog.New(slog.NewTextHandler(io.Discard, nil)) 22 + } 23 + 24 + func nopEncoder() scriptedEncoder { 25 + return scriptedEncoder(func(*millproto.Message) error { return nil }) 26 + } 27 + 28 + func TestAttachSessionRejectsSecondLiveSession(t *testing.T) { 29 + l := discardLogger() 30 + m := New(l, Config{ReconnectGrace: time.Minute}) 31 + 32 + sessionOf := func(node string) *millSession { 33 + m.mu.Lock() 34 + defer m.mu.Unlock() 35 + return m.sessions[node] 36 + } 37 + 38 + sess1 := newSession("node-1", nopEncoder(), l) 39 + if _, ok := m.attachSession(sess1); !ok { 40 + t.Fatal("first attach of a node was rejected; want accept") 41 + } 42 + if sessionOf("node-1") != sess1 { 43 + t.Fatal("first session was not registered as the live session") 44 + } 45 + 46 + sess2 := newSession("node-1", nopEncoder(), l) 47 + if _, ok := m.attachSession(sess2); ok { 48 + t.Fatal("second live attach for an already-live node was accepted; a valid token hijacked the executor") 49 + } 50 + if sessionOf("node-1") != sess1 { 51 + t.Fatal("rejected newcomer evicted the incumbent session") 52 + } 53 + 54 + m.detachSession(sess1) 55 + sess3 := newSession("node-1", nopEncoder(), l) 56 + if _, ok := m.attachSession(sess3); !ok { 57 + t.Fatal("attach during the incumbent's reconnect grace was rejected; want adopt") 58 + } 59 + if sessionOf("node-1") != sess3 { 60 + t.Fatal("adopted session was not installed as the live session") 61 + } 62 + } 63 + 64 + func TestOnAttemptResultIgnoresForeignLease(t *testing.T) { 65 + l := discardLogger() 66 + m := New(l, Config{ReconnectGrace: time.Minute}) 67 + 68 + foreign := newLease("lease-foreign", "node-a", "dummy") 69 + m.mu.Lock() 70 + m.leases[foreign.id] = foreign 71 + m.mu.Unlock() 72 + 73 + sessB := newSession("node-b", nopEncoder(), l) 74 + m.attachSession(sessB) 75 + m.onAttemptResult(sessB, &millv1.AttemptResult{ 76 + Offset: 1, 77 + LeaseId: foreign.id, 78 + TerminalStatus: string(models.StatusKindSuccess), 79 + }) 80 + 81 + if _, ok := pollTerminal(foreign); ok { 82 + t.Fatal("attempt-result on a foreign lease delivered a terminal; an executor forged another node's job result") 83 + } 84 + if foreign.getState() == leaseDone { 85 + t.Fatal("attempt-result on a foreign lease sealed the lease") 86 + } 87 + } 88 + 89 + func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) { 90 + l := discardLogger() 91 + m := New(l, Config{ReconnectGrace: time.Minute}) 92 + 93 + // a bystander lease the reporting node does own, to prove an absent-lease 94 + // relay does not spill onto some other lease. 95 + bystander := newLease("lease-bystander", "node-b", "dummy") 96 + m.mu.Lock() 97 + m.leases[bystander.id] = bystander 98 + m.mu.Unlock() 99 + 100 + sessB := newSession("node-b", nopEncoder(), l) 101 + m.attachSession(sessB) 102 + 103 + m.onAttemptResult(sessB, &millv1.AttemptResult{ 104 + Offset: 1, 105 + LeaseId: "lease-nonexistent", 106 + TerminalStatus: string(models.StatusKindSuccess), 107 + }) 108 + 109 + if _, ok := pollTerminal(bystander); ok { 110 + t.Fatal("attempt-result for an absent lease delivered a terminal to a bystander lease") 111 + } 112 + if bystander.getState() == leaseDone { 113 + t.Fatal("attempt-result for an absent lease sealed a bystander lease") 114 + } 115 + } 116 + 117 + // The owned-lease path proves the ignore tests above are not passing merely 118 + // because delivery is broken: a correctly owned terminal IS delivered. 119 + func TestOnAttemptResultDeliversOwnedLease(t *testing.T) { 120 + l := discardLogger() 121 + m := New(l, Config{ReconnectGrace: time.Minute}) 122 + 123 + owned := newLease("lease-owned", "node-b", "dummy") 124 + m.mu.Lock() 125 + m.leases[owned.id] = owned 126 + m.mu.Unlock() 127 + 128 + sessB := newSession("node-b", nopEncoder(), l) 129 + m.attachSession(sessB) 130 + m.onAttemptResult(sessB, &millv1.AttemptResult{ 131 + Offset: 1, 132 + LeaseId: owned.id, 133 + TerminalStatus: string(models.StatusKindSuccess), 134 + }) 135 + 136 + res, ok := pollTerminal(owned) 137 + if !ok { 138 + t.Fatal("attempt-result on an owned lease was not delivered") 139 + } 140 + if got := res.GetTerminalStatus(); got != string(models.StatusKindSuccess) { 141 + t.Fatalf("delivered terminal status = %q, want %q", got, string(models.StatusKindSuccess)) 142 + } 143 + if owned.getState() != leaseDone { 144 + t.Fatal("owned lease was not sealed after its terminal was delivered") 145 + } 146 + } 147 + 148 + func TestOnLogRelayRelayOwnership(t *testing.T) { 149 + l := discardLogger() 150 + dir := t.TempDir() 151 + m := New(l, Config{LogDir: dir, ReconnectGrace: time.Minute}) 152 + 153 + foreign := newLease("lease-foreign", "node-x", "dummy") 154 + foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"} 155 + owned := newLease("lease-owned", "node-z", "dummy") 156 + owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"} 157 + m.mu.Lock() 158 + m.leases[foreign.id] = foreign 159 + m.leases[owned.id] = owned 160 + m.mu.Unlock() 161 + 162 + sessY := newSession("node-y", nopEncoder(), l) 163 + m.attachSession(sessY) 164 + sessZ := newSession("node-z", nopEncoder(), l) 165 + m.attachSession(sessZ) 166 + 167 + m.onLogRelay(sessY, &millv1.LogLine{ 168 + Offset: 1, 169 + LeaseId: foreign.id, 170 + RawJson: []byte(`{"line":"forged"}`), 171 + }) 172 + if _, err := os.Stat(models.LogFilePath(dir, foreign.wid)); !os.IsNotExist(err) { 173 + t.Fatalf("log relay for a foreign lease wrote a file (stat err=%v); an executor forged another node's logs", err) 174 + } 175 + 176 + line := []byte(`{"line":"legit"}`) 177 + m.onLogRelay(sessZ, &millv1.LogLine{ 178 + Offset: 1, 179 + LeaseId: owned.id, 180 + RawJson: line, 181 + }) 182 + got, err := os.ReadFile(models.LogFilePath(dir, owned.wid)) 183 + if err != nil { 184 + t.Fatalf("owned log relay did not write the log file: %v", err) 185 + } 186 + if want := string(line) + "\n"; string(got) != want { 187 + t.Fatalf("owned log line = %q, want %q", string(got), want) 188 + } 189 + } 190 + 191 + func TestOnStatusRelayRelayOwnership(t *testing.T) { 192 + ctx := context.Background() 193 + l := discardLogger() 194 + 195 + bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db")) 196 + if err != nil { 197 + t.Fatalf("db.Make: %v", err) 198 + } 199 + t.Cleanup(func() { bdb.Close() }) 200 + n := notifier.New() 201 + 202 + m := New(l, Config{ReconnectGrace: time.Minute}) 203 + m.Attach(bdb, &n) 204 + 205 + foreign := newLease("lease-foreign", "node-x", "dummy") 206 + foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"} 207 + owned := newLease("lease-owned", "node-z", "dummy") 208 + owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"} 209 + m.mu.Lock() 210 + m.leases[foreign.id] = foreign 211 + m.leases[owned.id] = owned 212 + m.mu.Unlock() 213 + 214 + sessY := newSession("node-y", nopEncoder(), l) 215 + m.attachSession(sessY) 216 + sessZ := newSession("node-z", nopEncoder(), l) 217 + m.attachSession(sessZ) 218 + 219 + m.onStatusRelay(sessY, &millv1.StatusEvent{ 220 + Offset: 1, 221 + LeaseId: foreign.id, 222 + Status: "running", 223 + }) 224 + if _, err := bdb.GetStatus(foreign.wid); err == nil { 225 + t.Fatal("status relay for a foreign lease authored a status row; an executor forged another pipeline's status") 226 + } 227 + 228 + m.onStatusRelay(sessZ, &millv1.StatusEvent{ 229 + Offset: 1, 230 + LeaseId: owned.id, 231 + Status: "running", 232 + }) 233 + st, err := bdb.GetStatus(owned.wid) 234 + if err != nil { 235 + t.Fatalf("owned status relay did not author a status row: %v", err) 236 + } 237 + if st.Status != "running" { 238 + t.Fatalf("owned status = %q, want %q", st.Status, "running") 239 + } 240 + }

+94

spindle/mill/engine.go

··· 1 + package mill 2 + 3 + import ( 4 + "context" 5 + "log/slog" 6 + "time" 7 + 8 + "tangled.org/core/api/tangled" 9 + "tangled.org/core/spindle/engine" 10 + "tangled.org/core/spindle/models" 11 + "tangled.org/core/spindle/secrets" 12 + ) 13 + 14 + // millWorkflowState is stashed in Workflow.Data (which is `any`, like both real 15 + // engines use it). InitWorkflow parses nothing here: it just carries the raw 16 + // pipeline/workflow forward so the executor can run the real InitWorkflow later. 17 + type millWorkflowState struct { 18 + TargetEngine string 19 + RawWorkflow tangled.Pipeline_Workflow 20 + RawPipeline tangled.Pipeline 21 + Wid models.WorkflowId // stamped at placement (InitWorkflow can't see it) 22 + Lease *RemoteLease 23 + } 24 + 25 + // Engine is the mill's stand-in for a real engine, registered under the real 26 + // engine names ("microvm", "nixery"). All registered names share one Mill. 27 + type Engine struct { 28 + name string 29 + mill *Mill 30 + l *slog.Logger 31 + } 32 + 33 + // NewEngine returns a mill engine view for one engine name. 34 + func NewEngine(name string, mill *Mill) *Engine { 35 + return &Engine{name: name, mill: mill, l: mill.l.With("engine", "mill:"+name)} 36 + } 37 + 38 + // InitWorkflow returns a synthetic one-step workflow so processPipeline injects 39 + // TANGLED_* env and marks pending normally. The real engine's InitWorkflow runs 40 + // exactly once, on the executor inside ReserveSeat; the commit path reuses that 41 + // workflow rather than re-initializing. 42 + func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) { 43 + return &models.Workflow{ 44 + Name: twf.Name, 45 + Environment: map[string]string{}, 46 + Steps: []models.Step{remoteStep{}}, 47 + Data: &millWorkflowState{ 48 + TargetEngine: e.name, 49 + RawWorkflow: twf, 50 + RawPipeline: tpl, 51 + }, 52 + }, nil 53 + } 54 + 55 + // WorkflowLogger gives StartWorkflows a no-op logger for the mill's synthetic 56 + // workflow: the executor's real log lines are relayed into this wid's log file 57 + // directly (see Loom.onLogRelay), so a local logger here would only write 58 + // competing, step-id-colliding lines. This is engine's optional logger-provider 59 + // seam. 60 + func (e *Engine) WorkflowLogger(wid models.WorkflowId) models.WorkflowLogger { 61 + return models.NullLogger{} 62 + } 63 + 64 + // AcquireWorkflowSlot is the placement seam: it blocks doing remote placement, 65 + // which the user sees as "pending". mode is irrelevant here (only StartWorkflows 66 + // calls this, with Wait); the mill always places. 67 + func (e *Engine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, _ engine.AcquireMode) (engine.WorkflowSlot, error) { 68 + return e.mill.place(ctx, e.name, wid, wf) 69 + } 70 + 71 + // SetupWorkflow is a no-op info log; the real setup happens on the executor. 72 + func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, wfLogger models.WorkflowLogger) error { 73 + e.l.Info("remote job placed, awaiting commit", "wid", wid) 74 + return nil 75 + } 76 + 77 + // RunStep is commit-and-wait: hand over secrets and block on the terminal 78 + // result relayed over the session. 79 + func (e *Engine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, unlocked []secrets.UnlockedSecret, wfLogger models.WorkflowLogger) error { 80 + return e.mill.commitAndWait(ctx, w, unlocked) 81 + } 82 + 83 + // WorkflowTimeout is deliberately generous: the executor enforces the real 84 + // per-engine timeout and relays it. The mill only caps a hung/silent 85 + // executor, with a true death caught separately by reconnect grace. 86 + func (e *Engine) WorkflowTimeout() time.Duration { 87 + return e.mill.cfg.JobTimeout 88 + } 89 + 90 + // DestroyWorkflow cancels a still-running attempt; no-op if already terminal. 91 + func (e *Engine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error { 92 + e.mill.destroy(wid) 93 + return nil 94 + }

+72

spindle/mill/handler.go

··· 1 + package mill 2 + 3 + import ( 4 + "net/http" 5 + 6 + "github.com/gorilla/websocket" 7 + 8 + millproto "tangled.org/core/spindle/mill/proto" 9 + millv1 "tangled.org/core/spindle/mill/proto/gen" 10 + ) 11 + 12 + var upgrader = websocket.Upgrader{ 13 + ReadBufferSize: 1024, 14 + WriteBufferSize: 1024, 15 + } 16 + 17 + // HandleExecutorConn terminates one executor websocket. Auth is a pre-shared 18 + // secret in the Authorization header, checked before the upgrade so a bad token 19 + // never opens a socket. 20 + func (m *Mill) HandleExecutorConn(w http.ResponseWriter, r *http.Request) { 21 + if m.cfg.SharedSecret != "" { 22 + if r.Header.Get("Authorization") != "Bearer "+m.cfg.SharedSecret { 23 + http.Error(w, "unauthorized", http.StatusUnauthorized) 24 + return 25 + } 26 + } 27 + 28 + conn, err := upgrader.Upgrade(w, r, nil) 29 + if err != nil { 30 + m.l.Error("fleet ws upgrade failed", "err", err) 31 + return 32 + } 33 + defer conn.Close() 34 + 35 + stream := millproto.NewWSStream(conn) 36 + enc := millproto.NewEncoder(stream) 37 + dec := millproto.NewDecoder(stream) 38 + 39 + // first frame must be Hello. 40 + hello, err := dec.Decode() 41 + if err != nil { 42 + m.l.Error("fleet read hello failed", "err", err) 43 + return 44 + } 45 + h := hello.GetHello() 46 + if h == nil { 47 + m.l.Error("fleet first frame was not hello") 48 + return 49 + } 50 + if h.GetProtocolVersion() != millproto.ProtocolVersion { 51 + m.l.Error("fleet protocol version mismatch", "got", h.GetProtocolVersion(), "want", millproto.ProtocolVersion) 52 + return 53 + } 54 + 55 + sess := newSession(h.GetNodeId(), enc, m.l) 56 + resume, ok := m.attachSession(sess) 57 + if !ok { 58 + m.l.Warn("rejecting duplicate live executor session", "node", sess.nodeID) 59 + return 60 + } 61 + m.l.Info("executor connected", "node", sess.nodeID, "engines", h.GetEngines(), "arch", h.GetArch(), "resume", resume) 62 + 63 + if err := sess.send(&millproto.Message{Resume: &millv1.Resume{AckOffset: resume}}); err != nil { 64 + m.l.Error("fleet send resume failed", "err", err) 65 + m.detachSession(sess) 66 + return 67 + } 68 + m.sessionReady(sess) 69 + 70 + sess.readLoop(m, dec) 71 + m.detachSession(sess) 72 + }

+205

spindle/mill/lease.go

··· 1 + package mill 2 + 3 + import ( 4 + "os" 5 + "sync" 6 + 7 + "tangled.org/core/spindle/models" 8 + 9 + millv1 "tangled.org/core/spindle/mill/proto/gen" 10 + ) 11 + 12 + // leaseState is the mill's view of a remote attempt. It collapses the doc's 13 + // full state diagram to what the mill actually needs to drive cleanup and the 14 + // commit-and-wait in RunStep. 15 + type leaseState int32 16 + 17 + const ( 18 + // won a bid, executor is holding a seat, not yet committed. 19 + leaseReserved leaseState = iota 20 + // CommitLease has been sent; the executor may already be running the job, 21 + // but the mill has not necessarily received Committed yet. 22 + leaseCommitting 23 + // CommitLease sent and acked; the job is running on the executor. 24 + leaseRunning 25 + // a terminal result arrived (or we gave up). no further action. 26 + leaseDone 27 + ) 28 + 29 + type cancelAction int 30 + 31 + const ( 32 + cancelNoop cancelAction = iota 33 + cancelLocal 34 + cancelRemote 35 + ) 36 + 37 + // RemoteLease is the mill-side fencing token for one placed job. 38 + type RemoteLease struct { 39 + id string 40 + nodeID string 41 + engine string 42 + wid models.WorkflowId // the job this lease carries; set once placed 43 + 44 + mu sync.Mutex 45 + state leaseState 46 + cancel bool 47 + reason string 48 + terminal chan *millv1.AttemptResult // buffered(1); RunStep waits here 49 + dead chan struct{} // closed when the executor is lost past grace 50 + deadOnce sync.Once 51 + 52 + // log file for relayed lines, opened lazily and held open for the lease's 53 + // life so we don't reopen per line. Guarded by logMu, closed on cleanup. 54 + logMu sync.Mutex 55 + logFile *os.File 56 + logClosed bool 57 + } 58 + 59 + func newLease(id, nodeID, engine string) *RemoteLease { 60 + return &RemoteLease{ 61 + id: id, 62 + nodeID: nodeID, 63 + engine: engine, 64 + state: leaseReserved, 65 + terminal: make(chan *millv1.AttemptResult, 1), 66 + dead: make(chan struct{}), 67 + } 68 + } 69 + 70 + func (l *RemoteLease) setState(s leaseState) { 71 + l.mu.Lock() 72 + l.state = s 73 + l.mu.Unlock() 74 + } 75 + 76 + func (l *RemoteLease) markCommitting() bool { 77 + l.mu.Lock() 78 + defer l.mu.Unlock() 79 + if l.state == leaseDone { 80 + return false 81 + } 82 + if l.state == leaseReserved { 83 + l.state = leaseCommitting 84 + } 85 + return true 86 + } 87 + 88 + func (l *RemoteLease) markRunning() { 89 + l.mu.Lock() 90 + if l.state != leaseDone { 91 + l.state = leaseRunning 92 + } 93 + l.mu.Unlock() 94 + } 95 + 96 + func (l *RemoteLease) getState() leaseState { 97 + l.mu.Lock() 98 + defer l.mu.Unlock() 99 + return l.state 100 + } 101 + 102 + func (l *RemoteLease) requestCancel(reason string) cancelAction { 103 + l.mu.Lock() 104 + defer l.mu.Unlock() 105 + if l.state == leaseDone { 106 + return cancelNoop 107 + } 108 + l.cancel = true 109 + l.reason = reason 110 + if l.state == leaseReserved { 111 + return cancelLocal 112 + } 113 + return cancelRemote 114 + } 115 + 116 + func (l *RemoteLease) cancelRequested() (bool, string) { 117 + l.mu.Lock() 118 + defer l.mu.Unlock() 119 + return l.cancel, l.reason 120 + } 121 + 122 + func (l *RemoteLease) deliverCancelled(reason string) { 123 + l.deliverTerminal(&millv1.AttemptResult{ 124 + LeaseId: l.id, 125 + TerminalStatus: string(models.StatusKindCancelled), 126 + Error: reason, 127 + }) 128 + } 129 + 130 + // deliverTerminal hands the terminal result to a waiting RunStep without 131 + // blocking. Idempotent: a duplicate (e.g. replayed on reconnect) is dropped 132 + // because the channel is buffered to one and the lease is already done. 133 + func (l *RemoteLease) deliverTerminal(res *millv1.AttemptResult) { 134 + l.mu.Lock() 135 + if l.state == leaseDone { 136 + l.mu.Unlock() 137 + return 138 + } 139 + l.state = leaseDone 140 + l.mu.Unlock() 141 + 142 + select { 143 + case l.terminal <- res: 144 + default: 145 + } 146 + } 147 + 148 + // markDead signals that the executor running this lease is gone for good. 149 + func (l *RemoteLease) markDead() { 150 + l.deadOnce.Do(func() { close(l.dead) }) 151 + } 152 + 153 + // writeLog appends a relayed log line to the lease's log file, opening it once 154 + // and holding the handle open for the lease's life (relayed lines for a lease 155 + // all arrive on one session reader goroutine, so this is serialized). A write 156 + // after closeLog is dropped so a replayed line can't resurrect the file. 157 + func (l *RemoteLease) writeLog(path string, line []byte) error { 158 + l.logMu.Lock() 159 + defer l.logMu.Unlock() 160 + if l.logClosed { 161 + return nil 162 + } 163 + if l.logFile == nil { 164 + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) 165 + if err != nil { 166 + return err 167 + } 168 + l.logFile = f 169 + } 170 + _, err := l.logFile.Write(line) 171 + return err 172 + } 173 + 174 + // Idempotent. 175 + func (l *RemoteLease) closeLog() { 176 + l.logMu.Lock() 177 + defer l.logMu.Unlock() 178 + l.logClosed = true 179 + if l.logFile != nil { 180 + _ = l.logFile.Close() 181 + l.logFile = nil 182 + } 183 + } 184 + 185 + // remoteStep is the mill's synthetic single step. The real steps run on the 186 + // executor; the mill never mirrors them. 187 + type remoteStep struct{} 188 + 189 + func (remoteStep) Name() string { return "remote execution" } 190 + func (remoteStep) Command() string { return "" } 191 + func (remoteStep) Kind() models.StepKind { return models.StepKindSystem } 192 + 193 + // millSlot is what AcquireWorkflowSlot returns; its Release unwinds placement. 194 + type millSlot struct { 195 + fleet *Mill 196 + lease *RemoteLease 197 + once sync.Once 198 + } 199 + 200 + func (s *millSlot) Release() { 201 + if s == nil { 202 + return 203 + } 204 + s.once.Do(func() { s.fleet.releaseSlot(s) }) 205 + }

+693

spindle/mill/mill.go

··· 1 + package mill 2 + 3 + import ( 4 + "context" 5 + "encoding/json" 6 + "errors" 7 + "fmt" 8 + "log/slog" 9 + "sort" 10 + "sync" 11 + "time" 12 + 13 + "tangled.org/core/notifier" 14 + "tangled.org/core/spindle/db" 15 + "tangled.org/core/spindle/engine" 16 + "tangled.org/core/spindle/models" 17 + "tangled.org/core/spindle/secrets" 18 + "tangled.org/core/tid" 19 + 20 + millproto "tangled.org/core/spindle/mill/proto" 21 + millv1 "tangled.org/core/spindle/mill/proto/gen" 22 + ) 23 + 24 + const ( 25 + defaultReconnectGrace = 45 * time.Second 26 + defaultJobTimeout = 24 * time.Hour 27 + defaultBidTimeout = 5 * time.Second 28 + defaultTopK = 3 29 + // re-poll in case a change notification is missed. 30 + placementPoll = 2 * time.Second 31 + ) 32 + 33 + type Config struct { 34 + SharedSecret string 35 + MaxPending int 36 + ReconnectGrace time.Duration 37 + JobTimeout time.Duration 38 + BidTimeout time.Duration 39 + TopK int 40 + LogDir string 41 + } 42 + 43 + type Mill struct { 44 + l *slog.Logger 45 + cfg Config 46 + 47 + // re-author relayed status rows into the mill's stream so appview sees 48 + // them unchanged. 49 + db *db.DB 50 + n *notifier.Notifier 51 + 52 + mu sync.Mutex 53 + sessions map[string]*millSession 54 + leases map[string]*RemoteLease 55 + nodeOffset map[string]uint64 56 + pending int 57 + changeCh chan struct{} // closed + replaced to wake placement waiters 58 + 59 + leaseSeq uint64 60 + } 61 + 62 + func New(l *slog.Logger, cfg Config) *Mill { 63 + if cfg.ReconnectGrace <= 0 { 64 + cfg.ReconnectGrace = defaultReconnectGrace 65 + } 66 + if cfg.JobTimeout <= 0 { 67 + cfg.JobTimeout = defaultJobTimeout 68 + } 69 + if cfg.BidTimeout <= 0 { 70 + cfg.BidTimeout = defaultBidTimeout 71 + } 72 + if cfg.TopK <= 0 { 73 + cfg.TopK = defaultTopK 74 + } 75 + return &Mill{ 76 + l: l, 77 + cfg: cfg, 78 + sessions: make(map[string]*millSession), 79 + leases: make(map[string]*RemoteLease), 80 + nodeOffset: make(map[string]uint64), 81 + changeCh: make(chan struct{}), 82 + } 83 + } 84 + 85 + func (m *Mill) Attach(d *db.DB, n *notifier.Notifier) { 86 + m.mu.Lock() 87 + m.db = d 88 + m.n = n 89 + m.mu.Unlock() 90 + } 91 + 92 + func (m *Mill) nextLeaseID() string { 93 + m.mu.Lock() 94 + m.leaseSeq++ 95 + seq := m.leaseSeq 96 + m.mu.Unlock() 97 + return fmt.Sprintf("%s-%d", tid.TID(), seq) 98 + } 99 + 100 + func (m *Mill) notifyChange() { 101 + m.mu.Lock() 102 + m.notifyChangeLocked() 103 + m.mu.Unlock() 104 + } 105 + 106 + func (m *Mill) notifyChangeLocked() { 107 + close(m.changeCh) 108 + m.changeCh = make(chan struct{}) 109 + } 110 + 111 + func (m *Mill) currentChangeCh() <-chan struct{} { 112 + m.mu.Lock() 113 + defer m.mu.Unlock() 114 + return m.changeCh 115 + } 116 + 117 + // --- session lifecycle ----------------------------------------------------- 118 + 119 + func (m *Mill) attachSession(sess *millSession) (uint64, bool) { 120 + m.mu.Lock() 121 + defer m.mu.Unlock() 122 + 123 + if old := m.sessions[sess.nodeID]; old != nil { 124 + if !old.disconnected { 125 + // reject a second live session for the same identity so a valid token 126 + // can't hijack an in-flight executor. 127 + return 0, false 128 + } 129 + // the old session is in its reconnect grace window: adopt its leases. 130 + if old.graceTimer != nil { 131 + old.graceTimer.Stop() 132 + } 133 + old.close() 134 + m.l.Info("executor reconnected", "node", sess.nodeID) 135 + } 136 + m.sessions[sess.nodeID] = sess 137 + // wake commit retries waiting out this node's reconnect grace. 138 + m.notifyChangeLocked() 139 + return m.nodeOffset[sess.nodeID], true 140 + } 141 + 142 + func (m *Mill) detachSession(sess *millSession) { 143 + m.mu.Lock() 144 + if m.sessions[sess.nodeID] != sess { 145 + // already replaced by a reconnect; nothing to do. 146 + m.mu.Unlock() 147 + sess.close() 148 + return 149 + } 150 + sess.disconnected = true 151 + sess.graceTimer = time.AfterFunc(m.cfg.ReconnectGrace, func() { m.onGraceExpired(sess) }) 152 + m.mu.Unlock() 153 + 154 + sess.close() 155 + m.l.Warn("executor session lost; entering reconnect grace", "node", sess.nodeID, "grace", m.cfg.ReconnectGrace) 156 + m.notifyChange() 157 + } 158 + 159 + func (m *Mill) sessionReady(sess *millSession) { 160 + for _, lease := range m.cancelledLeasesForNode(sess.nodeID) { 161 + _, reason := lease.cancelRequested() 162 + m.sendCancel(sess, lease, reason) 163 + } 164 + m.notifyChange() 165 + } 166 + 167 + func (m *Mill) cancelledLeasesForNode(nodeID string) []*RemoteLease { 168 + m.mu.Lock() 169 + var candidates []*RemoteLease 170 + for _, lease := range m.leases { 171 + if lease.nodeID == nodeID { 172 + candidates = append(candidates, lease) 173 + } 174 + } 175 + m.mu.Unlock() 176 + 177 + var leases []*RemoteLease 178 + for _, lease := range candidates { 179 + if cancelled, _ := lease.cancelRequested(); cancelled && lease.getState() != leaseDone { 180 + leases = append(leases, lease) 181 + } 182 + } 183 + return leases 184 + } 185 + 186 + func (m *Mill) onGraceExpired(sess *millSession) { 187 + m.mu.Lock() 188 + if m.sessions[sess.nodeID] != sess || !sess.disconnected { 189 + // reconnected in the meantime. 190 + m.mu.Unlock() 191 + return 192 + } 193 + var dead []*RemoteLease 194 + for _, lease := range m.leases { 195 + if lease.nodeID == sess.nodeID { 196 + dead = append(dead, lease) 197 + } 198 + } 199 + delete(m.sessions, sess.nodeID) 200 + delete(m.nodeOffset, sess.nodeID) 201 + m.mu.Unlock() 202 + 203 + m.l.Warn("executor declared dead; failing its in-flight jobs", "node", sess.nodeID, "jobs", len(dead)) 204 + for _, lease := range dead { 205 + if cancelled, reason := lease.cancelRequested(); cancelled { 206 + lease.deliverCancelled(reason) 207 + } else { 208 + lease.markDead() 209 + } 210 + } 211 + m.notifyChange() 212 + } 213 + 214 + // --- placement (the WorkflowSlotter seam) --------------------------------- 215 + 216 + func (m *Mill) place(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow) (engine.WorkflowSlot, error) { 217 + m.mu.Lock() 218 + if m.cfg.MaxPending > 0 && m.pending >= m.cfg.MaxPending { 219 + max := m.cfg.MaxPending 220 + cur := m.pending 221 + m.mu.Unlock() 222 + return nil, fmt.Errorf("%w: mill has %d pending jobs (max %d)", engine.ErrNoWorkflowSlots, cur, max) 223 + } 224 + m.pending++ 225 + m.mu.Unlock() 226 + defer func() { 227 + m.mu.Lock() 228 + m.pending-- 229 + m.mu.Unlock() 230 + }() 231 + 232 + // stamp wid early so the executor reconstructs logs/status under the same 233 + // pipeline id the mill used for "pending". 234 + if st, ok := wf.Data.(*millWorkflowState); ok && st != nil { 235 + st.Wid = wid 236 + } 237 + 238 + for { 239 + if err := ctx.Err(); err != nil { 240 + return nil, err 241 + } 242 + 243 + lease, err := m.bid(ctx, engineName, wf) 244 + if err != nil { 245 + return nil, err 246 + } 247 + if lease != nil { 248 + lease.wid = wid 249 + m.mu.Lock() 250 + m.leases[lease.id] = lease 251 + if st, ok := wf.Data.(*millWorkflowState); ok && st != nil { 252 + st.Lease = lease 253 + } 254 + m.mu.Unlock() 255 + return &millSlot{fleet: m, lease: lease}, nil 256 + } 257 + 258 + // nobody to place on right now: wait for a change, ctx, or a poll tick. 259 + m.mu.Lock() 260 + ch := m.changeCh 261 + m.mu.Unlock() 262 + select { 263 + case <-ctx.Done(): 264 + return nil, ctx.Err() 265 + case <-ch: 266 + case <-time.After(placementPoll): 267 + } 268 + } 269 + } 270 + 271 + func (m *Mill) bid(ctx context.Context, engineName string, wf *models.Workflow) (*RemoteLease, error) { 272 + rawPipeline, rawWorkflow, err := marshalJob(wf) 273 + if err != nil { 274 + return nil, err 275 + } 276 + 277 + candidates := m.rankCandidates(engineName) 278 + if len(candidates) == 0 { 279 + return nil, nil 280 + } 281 + if len(candidates) > m.cfg.TopK { 282 + candidates = candidates[:m.cfg.TopK] 283 + } 284 + 285 + bidCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout) 286 + defer cancel() 287 + 288 + type bidResult struct { 289 + sess *millSession 290 + lease *RemoteLease 291 + score float64 292 + } 293 + results := make(chan bidResult, len(candidates)) 294 + var wg sync.WaitGroup 295 + for _, sess := range candidates { 296 + wg.Add(1) 297 + go func(sess *millSession) { 298 + defer wg.Done() 299 + leaseID := m.nextLeaseID() 300 + lease := newLease(leaseID, sess.nodeID, engineName) 301 + wid := m.jobWid(wf) 302 + msg := &millproto.Message{ReserveSeat: &millv1.ReserveSeat{ 303 + LeaseId: leaseID, 304 + TargetEngine: engineName, 305 + RawPipelineJson: rawPipeline, 306 + RawWorkflowJson: rawWorkflow, 307 + Knot: wid.Knot, 308 + Rkey: wid.Rkey, 309 + TtlSeconds: uint32(m.cfg.ReconnectGrace / time.Second), 310 + }} 311 + resp, err := sess.request(bidCtx, leaseID, msg) 312 + if err != nil { 313 + return 314 + } 315 + rr := resp.GetReserveResult() 316 + if rr == nil || !rr.GetAccepted() { 317 + return 318 + } 319 + results <- bidResult{sess: sess, lease: lease, score: rr.GetScore()} 320 + }(sess) 321 + } 322 + wg.Wait() 323 + close(results) 324 + 325 + var winner *bidResult 326 + var losers []*RemoteLease 327 + for r := range results { 328 + if winner == nil || r.score > winner.score { 329 + if winner != nil { 330 + losers = append(losers, winner.lease) 331 + } 332 + winner = &r 333 + } else { 334 + losers = append(losers, r.lease) 335 + } 336 + } 337 + 338 + // release the also-rans so they free their held seats immediately. 339 + for _, l := range losers { 340 + m.releaseRemote(l) 341 + } 342 + 343 + if winner == nil { 344 + return nil, nil 345 + } 346 + return winner.lease, nil 347 + } 348 + 349 + func (m *Mill) rankCandidates(engineName string) []*millSession { 350 + m.mu.Lock() 351 + defer m.mu.Unlock() 352 + 353 + type ranked struct { 354 + sess *millSession 355 + seats uint32 356 + } 357 + var rs []ranked 358 + for _, s := range m.sessions { 359 + if s.disconnected { 360 + continue 361 + } 362 + if s.snapshot == nil { 363 + continue 364 + } 365 + es, ok := s.snapshot.GetEngines()[engineName] 366 + if !ok || es.GetFreeSeats() == 0 { 367 + continue 368 + } 369 + rs = append(rs, ranked{sess: s, seats: es.GetFreeSeats()}) 370 + } 371 + sort.SliceStable(rs, func(i, j int) bool { return rs[i].seats > rs[j].seats }) 372 + 373 + out := make([]*millSession, len(rs)) 374 + for i := range rs { 375 + out[i] = rs[i].sess 376 + } 377 + return out 378 + } 379 + 380 + // --- commit + wait (RunStep) ---------------------------------------------- 381 + 382 + func (m *Mill) commitAndWait(ctx context.Context, wf *models.Workflow, unlocked []secrets.UnlockedSecret) error { 383 + st, ok := wf.Data.(*millWorkflowState) 384 + if !ok || st == nil || st.Lease == nil { 385 + return fmt.Errorf("mill workflow state missing lease") 386 + } 387 + lease := st.Lease 388 + 389 + pbSecrets := make([]*millv1.Secret, len(unlocked)) 390 + for i, s := range unlocked { 391 + pbSecrets[i] = &millv1.Secret{Key: s.Key, Value: s.Value} 392 + } 393 + 394 + commit := &millproto.Message{CommitLease: &millv1.CommitLease{ 395 + LeaseId: lease.id, 396 + Secrets: pbSecrets, 397 + }} 398 + 399 + for { 400 + if res, ok := pollTerminal(lease); ok { 401 + return terminalError(res.GetTerminalStatus()) 402 + } 403 + if !lease.markCommitting() { 404 + return engine.ErrWorkflowFailed 405 + } 406 + 407 + sess := m.sessionForNode(lease.nodeID) 408 + if sess == nil { 409 + if done, err := m.waitCommitRetry(ctx, lease); done || err != nil { 410 + return err 411 + } 412 + continue 413 + } 414 + 415 + reqCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout) 416 + resp, err := sess.request(reqCtx, lease.id, commit) 417 + cancel() 418 + if err != nil { 419 + switch { 420 + case errors.Is(err, errSessionClosed): 421 + if done, err := m.waitCommitRetry(ctx, lease); done || err != nil { 422 + return err 423 + } 424 + continue 425 + case errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil: 426 + continue 427 + case errors.Is(err, context.DeadlineExceeded): 428 + return engine.ErrTimedOut 429 + case errors.Is(err, context.Canceled): 430 + return err 431 + default: 432 + m.l.Warn("commit lease send failed; waiting for reconnect", "lease", lease.id, "node", lease.nodeID, "err", err) 433 + if done, err := m.waitCommitRetry(ctx, lease); done || err != nil { 434 + return err 435 + } 436 + continue 437 + } 438 + } 439 + if resp.GetCommitted() == nil { 440 + return engine.ErrWorkflowFailed 441 + } 442 + lease.markRunning() 443 + if cancelled, reason := lease.cancelRequested(); cancelled { 444 + m.sendCancel(sess, lease, reason) 445 + } 446 + break 447 + } 448 + 449 + select { 450 + case res := <-lease.terminal: 451 + return terminalError(res.GetTerminalStatus()) 452 + case <-lease.dead: 453 + m.l.Warn("executor lost while running job", "lease", lease.id, "node", lease.nodeID) 454 + return engine.ErrWorkflowFailed 455 + case <-ctx.Done(): 456 + if ctx.Err() == context.DeadlineExceeded { 457 + return engine.ErrTimedOut 458 + } 459 + return ctx.Err() 460 + } 461 + } 462 + 463 + func (m *Mill) waitCommitRetry(ctx context.Context, lease *RemoteLease) (bool, error) { 464 + // grab the channel before re-checking for a live session so a reconnect 465 + // landing in between still closes the channel we select on. 466 + ch := m.currentChangeCh() 467 + if m.sessionForNode(lease.nodeID) != nil { 468 + return false, nil 469 + } 470 + select { 471 + case res := <-lease.terminal: 472 + return true, terminalError(res.GetTerminalStatus()) 473 + case <-lease.dead: 474 + return true, engine.ErrWorkflowFailed 475 + case <-ctx.Done(): 476 + if ctx.Err() == context.DeadlineExceeded { 477 + return true, engine.ErrTimedOut 478 + } 479 + return true, ctx.Err() 480 + case <-ch: 481 + return false, nil 482 + case <-time.After(placementPoll): 483 + return false, nil 484 + } 485 + } 486 + 487 + func terminalError(status string) error { 488 + switch models.StatusKind(status) { 489 + case models.StatusKindSuccess: 490 + return nil 491 + case models.StatusKindTimeout: 492 + return engine.ErrTimedOut 493 + case models.StatusKindCancelled: 494 + return engine.ErrCancelled 495 + default: 496 + return engine.ErrWorkflowFailed 497 + } 498 + } 499 + 500 + func pollTerminal(lease *RemoteLease) (*millv1.AttemptResult, bool) { 501 + select { 502 + case res := <-lease.terminal: 503 + return res, true 504 + default: 505 + return nil, false 506 + } 507 + } 508 + 509 + // --- cleanup --------------------------------------------------------------- 510 + 511 + func (m *Mill) destroy(wid models.WorkflowId) { 512 + m.mu.Lock() 513 + var lease *RemoteLease 514 + for _, l := range m.leases { 515 + if l.wid == wid { 516 + lease = l 517 + break 518 + } 519 + } 520 + m.mu.Unlock() 521 + if lease == nil { 522 + return 523 + } 524 + reason := "workflow destroyed" 525 + switch lease.requestCancel(reason) { 526 + case cancelLocal: 527 + if sess := m.sessionForNode(lease.nodeID); sess != nil { 528 + _ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}}) 529 + } 530 + lease.deliverCancelled(reason) 531 + case cancelRemote: 532 + if sess := m.sessionForNode(lease.nodeID); sess != nil { 533 + m.sendCancel(sess, lease, reason) 534 + } 535 + } 536 + } 537 + 538 + func (m *Mill) releaseSlot(s *millSlot) { 539 + lease := s.lease 540 + if lease.getState() == leaseReserved { 541 + m.releaseRemote(lease) 542 + } 543 + lease.closeLog() 544 + 545 + m.mu.Lock() 546 + delete(m.leases, lease.id) 547 + m.mu.Unlock() 548 + 549 + m.notifyChange() 550 + } 551 + 552 + func (m *Mill) releaseRemote(lease *RemoteLease) { 553 + lease.setState(leaseDone) 554 + if sess := m.sessionForNode(lease.nodeID); sess != nil { 555 + _ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}}) 556 + } 557 + } 558 + 559 + func (m *Mill) sendCancel(sess *millSession, lease *RemoteLease, reason string) { 560 + _ = sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{ 561 + LeaseId: lease.id, 562 + Reason: reason, 563 + }}) 564 + } 565 + 566 + func (m *Mill) sessionForNode(nodeID string) *millSession { 567 + m.mu.Lock() 568 + defer m.mu.Unlock() 569 + s := m.sessions[nodeID] 570 + if s == nil || s.disconnected { 571 + return nil 572 + } 573 + return s 574 + } 575 + 576 + // --- relay handlers -------------------------------------------------------- 577 + 578 + func (m *Mill) shouldAcceptOffset(sess *millSession, offset uint64) bool { 579 + m.mu.Lock() 580 + defer m.mu.Unlock() 581 + current := m.nodeOffset[sess.nodeID] 582 + return offset == current+1 583 + } 584 + 585 + func (m *Mill) ackOffset(sess *millSession, offset uint64) { 586 + m.mu.Lock() 587 + if offset > m.nodeOffset[sess.nodeID] { 588 + m.nodeOffset[sess.nodeID] = offset 589 + } 590 + m.mu.Unlock() 591 + 592 + _ = sess.send(&millproto.Message{Ack: &millv1.Ack{UpToOffset: offset}}) 593 + } 594 + 595 + func (m *Mill) processRelay(sess *millSession, offset uint64, kind string, apply func() error) { 596 + if !m.shouldAcceptOffset(sess, offset) { 597 + return 598 + } 599 + // ack only after the side effect lands, or a transient write error makes 600 + // the message unreplayable. 601 + if err := apply(); err != nil { 602 + m.l.Error("process relayed message failed", "kind", kind, "node", sess.nodeID, "offset", offset, "err", err) 603 + return 604 + } 605 + m.ackOffset(sess, offset) 606 + } 607 + 608 + func (m *Mill) onSnapshot(sess *millSession, snap *millv1.NodeSnapshot) { 609 + m.mu.Lock() 610 + sess.snapshot = snap 611 + m.mu.Unlock() 612 + m.notifyChange() 613 + } 614 + 615 + func (m *Mill) onStatusRelay(sess *millSession, ev *millv1.StatusEvent) { 616 + m.processRelay(sess, ev.GetOffset(), "status", func() error { 617 + if m.db == nil { 618 + return fmt.Errorf("mill db not attached") 619 + } 620 + m.mu.Lock() 621 + lease := m.leases[ev.GetLeaseId()] 622 + m.mu.Unlock() 623 + if lease == nil || lease.nodeID != sess.nodeID { 624 + // not a lease this executor holds; don't let it author status rows for 625 + // arbitrary pipelines. 626 + return nil 627 + } 628 + var errMsg *string 629 + if e := ev.GetError(); e != "" { 630 + errMsg = &e 631 + } 632 + var exit *int64 633 + if c := ev.GetExitCode(); c != 0 { 634 + exit = &c 635 + } 636 + // authoritative pipeline/workflow come from the lease, not whatever the 637 + // executor sent. 638 + return m.db.InsertRelayedStatus(string(lease.wid.PipelineId.AtUri()), lease.wid.Name, ev.GetStatus(), errMsg, exit, m.n) 639 + }) 640 + } 641 + 642 + func (m *Mill) onLogRelay(sess *millSession, ll *millv1.LogLine) { 643 + m.processRelay(sess, ll.GetOffset(), "log", func() error { 644 + m.mu.Lock() 645 + lease := m.leases[ll.GetLeaseId()] 646 + m.mu.Unlock() 647 + if lease == nil || lease.nodeID != sess.nodeID { 648 + return nil 649 + } 650 + path := models.LogFilePath(m.cfg.LogDir, lease.wid) 651 + line := append(ll.GetRawJson(), '\n') 652 + return lease.writeLog(path, line) 653 + }) 654 + } 655 + 656 + func (m *Mill) onAttemptResult(sess *millSession, ar *millv1.AttemptResult) { 657 + m.processRelay(sess, ar.GetOffset(), "attempt-result", func() error { 658 + m.mu.Lock() 659 + lease := m.leases[ar.GetLeaseId()] 660 + m.mu.Unlock() 661 + if lease == nil || lease.nodeID != sess.nodeID { 662 + return nil 663 + } 664 + lease.deliverTerminal(ar) 665 + return nil 666 + }) 667 + } 668 + 669 + // --- helpers --------------------------------------------------------------- 670 + 671 + func marshalJob(wf *models.Workflow) (pipeline string, workflow string, err error) { 672 + st, ok := wf.Data.(*millWorkflowState) 673 + if !ok || st == nil { 674 + return "", "", fmt.Errorf("mill workflow state missing") 675 + } 676 + p, err := json.Marshal(st.RawPipeline) 677 + if err != nil { 678 + return "", "", fmt.Errorf("marshal pipeline: %w", err) 679 + } 680 + w, err := json.Marshal(st.RawWorkflow) 681 + if err != nil { 682 + return "", "", fmt.Errorf("marshal workflow: %w", err) 683 + } 684 + return string(p), string(w), nil 685 + } 686 + 687 + func (m *Mill) jobWid(wf *models.Workflow) models.WorkflowId { 688 + st, ok := wf.Data.(*millWorkflowState) 689 + if !ok || st == nil { 690 + return models.WorkflowId{} 691 + } 692 + return st.Wid 693 + }

+177

spindle/mill/mill_test.go

··· 1 + package mill 2 + 3 + import ( 4 + "context" 5 + "errors" 6 + "io" 7 + "log/slog" 8 + "testing" 9 + "time" 10 + 11 + "tangled.org/core/api/tangled" 12 + "tangled.org/core/spindle/engine" 13 + "tangled.org/core/spindle/models" 14 + 15 + millproto "tangled.org/core/spindle/mill/proto" 16 + millv1 "tangled.org/core/spindle/mill/proto/gen" 17 + ) 18 + 19 + type scriptedEncoder func(*millproto.Message) error 20 + 21 + func (e scriptedEncoder) Encode(msg *millproto.Message) error { return e(msg) } 22 + 23 + func testWorkflow(name string) *models.Workflow { 24 + return &models.Workflow{ 25 + Name: name, 26 + Environment: map[string]string{}, 27 + Steps: []models.Step{remoteStep{}}, 28 + Data: &millWorkflowState{ 29 + TargetEngine: "dummy", 30 + RawWorkflow: tangled.Pipeline_Workflow{Name: name}, 31 + RawPipeline: tangled.Pipeline{}, 32 + }, 33 + } 34 + } 35 + 36 + func TestCommitRetriesAfterSessionCloseBeforeCommitted(t *testing.T) { 37 + l := slog.New(slog.NewTextHandler(io.Discard, nil)) 38 + m := New(l, Config{BidTimeout: 25 * time.Millisecond, ReconnectGrace: time.Second}) 39 + wf := testWorkflow("build") 40 + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"} 41 + lease := newLease("lease-1", "node-1", "dummy") 42 + lease.wid = wid 43 + wf.Data.(*millWorkflowState).Lease = lease 44 + 45 + m.mu.Lock() 46 + m.leases[lease.id] = lease 47 + m.mu.Unlock() 48 + 49 + var sess1 *millSession 50 + firstCommit := make(chan struct{}) 51 + sess1 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error { 52 + if msg.GetCommitLease() != nil { 53 + close(firstCommit) 54 + m.detachSession(sess1) 55 + } 56 + return nil 57 + }), l) 58 + m.attachSession(sess1) 59 + 60 + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) 61 + defer cancel() 62 + done := make(chan error, 1) 63 + go func() { done <- m.commitAndWait(ctx, wf, nil) }() 64 + 65 + select { 66 + case <-firstCommit: 67 + case <-ctx.Done(): 68 + t.Fatal("first commit was not sent") 69 + } 70 + 71 + var sess2 *millSession 72 + sess2 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error { 73 + if msg.GetCommitLease() == nil { 74 + return nil 75 + } 76 + leaseID := msg.GetCommitLease().GetLeaseId() 77 + sess2.deliver(leaseID, &millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}}) 78 + m.onAttemptResult(sess2, &millv1.AttemptResult{ 79 + Offset: 1, 80 + LeaseId: leaseID, 81 + TerminalStatus: string(models.StatusKindSuccess), 82 + }) 83 + return nil 84 + }), l) 85 + m.attachSession(sess2) 86 + m.sessionReady(sess2) 87 + 88 + select { 89 + case err := <-done: 90 + if err != nil { 91 + t.Fatalf("commitAndWait() error = %v, want success after reconnect", err) 92 + } 93 + case <-ctx.Done(): 94 + t.Fatal("commitAndWait() did not finish after reconnect") 95 + } 96 + } 97 + 98 + func TestDestroyRunningLeaseDoesNotDropCancelledTerminal(t *testing.T) { 99 + l := slog.New(slog.NewTextHandler(io.Discard, nil)) 100 + m := New(l, Config{}) 101 + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"} 102 + lease := newLease("lease-1", "node-1", "dummy") 103 + lease.wid = wid 104 + lease.setState(leaseRunning) 105 + 106 + m.mu.Lock() 107 + m.leases[lease.id] = lease 108 + m.mu.Unlock() 109 + 110 + m.destroy(wid) 111 + if lease.getState() == leaseDone { 112 + t.Fatal("destroy sealed the lease before the terminal result") 113 + } 114 + 115 + lease.deliverTerminal(&millv1.AttemptResult{ 116 + LeaseId: lease.id, 117 + TerminalStatus: string(models.StatusKindCancelled), 118 + }) 119 + res := <-lease.terminal 120 + if err := terminalError(res.GetTerminalStatus()); !errors.Is(err, engine.ErrCancelled) { 121 + t.Fatalf("terminalError() = %v, want ErrCancelled", err) 122 + } 123 + } 124 + 125 + func TestPlaceBlocksWhenNoCapacity(t *testing.T) { 126 + l := slog.New(slog.NewTextHandler(io.Discard, nil)) 127 + m := New(l, Config{}) 128 + wf := testWorkflow("build") 129 + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"} 130 + 131 + // no executors at all: place must block until ctx expires (user sees pending). 132 + ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond) 133 + defer cancel() 134 + 135 + _, err := m.place(ctx, "dummy", wid, wf) 136 + if err != context.DeadlineExceeded { 137 + t.Fatalf("place() error = %v, want DeadlineExceeded", err) 138 + } 139 + } 140 + 141 + func TestMaxPendingRejects(t *testing.T) { 142 + l := slog.New(slog.NewTextHandler(io.Discard, nil)) 143 + m := New(l, Config{MaxPending: 1}) 144 + 145 + m.mu.Lock() 146 + m.pending = 1 147 + m.mu.Unlock() 148 + 149 + wf2 := testWorkflow("b") 150 + _, err := m.place(context.Background(), "dummy", models.WorkflowId{Name: "b"}, wf2) 151 + if err == nil { 152 + t.Fatal("place() past maxPending should error") 153 + } 154 + } 155 + 156 + func TestSessionRequestCancelledBeforeRegistrationOrSend(t *testing.T) { 157 + sent := 0 158 + sess := newSession("node-1", scriptedEncoder(func(*millproto.Message) error { 159 + sent++ 160 + return nil 161 + }), discardLogger()) 162 + ctx, cancel := context.WithCancel(context.Background()) 163 + cancel() 164 + _, err := sess.request(ctx, "lease-1", &millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: "lease-1"}}) 165 + if !errors.Is(err, context.Canceled) { 166 + t.Fatalf("request error = %v, want context.Canceled", err) 167 + } 168 + if sent != 0 { 169 + t.Fatalf("request sent %d messages for an already-cancelled context, want 0", sent) 170 + } 171 + sess.mu.Lock() 172 + pending := len(sess.pending) 173 + sess.mu.Unlock() 174 + if pending != 0 { 175 + t.Fatalf("request left %d pending waiters, want 0", pending) 176 + } 177 + }

+142

spindle/mill/session.go

··· 1 + package mill 2 + 3 + import ( 4 + "context" 5 + "errors" 6 + "log/slog" 7 + "sync" 8 + "time" 9 + 10 + millproto "tangled.org/core/spindle/mill/proto" 11 + millv1 "tangled.org/core/spindle/mill/proto/gen" 12 + ) 13 + 14 + var errSessionClosed = errors.New("mill: executor session closed") 15 + 16 + // millSession is one live websocket to an executor. Unlike the microVM 17 + // AgentSession it is multiplexed: many leases plus async status/log relay share 18 + // it, so a single reader goroutine demuxes by message type and correlates 19 + // request/response by lease id. We never hold a lock across a decode. 20 + type millSession struct { 21 + nodeID string 22 + enc messageEncoder 23 + l *slog.Logger 24 + 25 + // snapshot, disconnected and graceTimer are guarded by Mill.mu, not the 26 + // session mutex below (the fleet ranks across sessions under its own lock). 27 + snapshot *millv1.NodeSnapshot 28 + disconnected bool 29 + graceTimer *time.Timer 30 + 31 + mu sync.Mutex 32 + pending map[string]chan *millproto.Message // lease_id -> response waiter 33 + 34 + closeOnce sync.Once 35 + closed chan struct{} 36 + } 37 + 38 + type messageEncoder interface { 39 + Encode(*millproto.Message) error 40 + } 41 + 42 + func newSession(nodeID string, enc messageEncoder, l *slog.Logger) *millSession { 43 + return &millSession{ 44 + nodeID: nodeID, 45 + enc: enc, 46 + l: l, 47 + pending: make(map[string]chan *millproto.Message), 48 + closed: make(chan struct{}), 49 + } 50 + } 51 + 52 + func (s *millSession) send(msg *millproto.Message) error { 53 + return s.enc.Encode(msg) 54 + } 55 + 56 + func (s *millSession) close() { 57 + s.closeOnce.Do(func() { close(s.closed) }) 58 + } 59 + 60 + // await registers a one-shot waiter for the next response addressed to leaseID, 61 + // returning the channel and a cancel that unregisters it. 62 + func (s *millSession) await(leaseID string) (<-chan *millproto.Message, func()) { 63 + ch := make(chan *millproto.Message, 1) 64 + s.mu.Lock() 65 + s.pending[leaseID] = ch 66 + s.mu.Unlock() 67 + return ch, func() { 68 + s.mu.Lock() 69 + if s.pending[leaseID] == ch { 70 + delete(s.pending, leaseID) 71 + } 72 + s.mu.Unlock() 73 + } 74 + } 75 + 76 + // deliver routes a correlated response to a waiter, if any is registered. 77 + func (s *millSession) deliver(leaseID string, msg *millproto.Message) { 78 + s.mu.Lock() 79 + ch := s.pending[leaseID] 80 + delete(s.pending, leaseID) 81 + s.mu.Unlock() 82 + if ch != nil { 83 + select { 84 + case ch <- msg: 85 + default: 86 + } 87 + } 88 + } 89 + 90 + // request sends msg and waits for the correlated response addressed to leaseID, 91 + // honouring ctx and session close. 92 + func (s *millSession) request(ctx context.Context, leaseID string, msg *millproto.Message) (*millproto.Message, error) { 93 + if err := ctx.Err(); err != nil { 94 + return nil, err 95 + } 96 + ch, cancel := s.await(leaseID) 97 + defer cancel() 98 + 99 + if err := s.send(msg); err != nil { 100 + return nil, err 101 + } 102 + 103 + select { 104 + case resp := <-ch: 105 + return resp, nil 106 + case <-ctx.Done(): 107 + return nil, ctx.Err() 108 + case <-s.closed: 109 + return nil, errSessionClosed 110 + } 111 + } 112 + 113 + // readLoop demuxes incoming frames until the decoder errors (connection gone). 114 + func (s *millSession) readLoop(m *Mill, dec *millproto.Decoder) { 115 + for { 116 + msg, err := dec.Decode() 117 + if err != nil { 118 + s.l.Debug("session read ended", "node", s.nodeID, "err", err) 119 + return 120 + } 121 + s.dispatch(m, msg) 122 + } 123 + } 124 + 125 + func (s *millSession) dispatch(m *Mill, msg *millproto.Message) { 126 + switch { 127 + case msg.GetNodeSnapshot() != nil: 128 + m.onSnapshot(s, msg.GetNodeSnapshot()) 129 + case msg.GetReserveResult() != nil: 130 + s.deliver(msg.GetReserveResult().GetLeaseId(), msg) 131 + case msg.GetCommitted() != nil: 132 + s.deliver(msg.GetCommitted().GetLeaseId(), msg) 133 + case msg.GetStatusEvent() != nil: 134 + m.onStatusRelay(s, msg.GetStatusEvent()) 135 + case msg.GetLogLine() != nil: 136 + m.onLogRelay(s, msg.GetLogLine()) 137 + case msg.GetAttemptResult() != nil: 138 + m.onAttemptResult(s, msg.GetAttemptResult()) 139 + default: 140 + s.l.Warn("session received unexpected message", "node", s.nodeID) 141 + } 142 + }

Configure Feed

Configure Feed