diff --git a/spindle/db/events.go b/spindle/db/events.go index fa5b0b9b..c36612d1 100644 --- a/spindle/db/events.go +++ b/spindle/db/events.go @@ -32,35 +32,61 @@ func (d *DB) CreatePipelineEvent(rkey string, pipeline tangled.Pipeline, n *noti return d.insertEvent(event, n) } -func (d *DB) createStatusEvent( - workflowId models.WorkflowId, - statusKind models.StatusKind, - workflowError *string, - exitCode *int64, - n *notifier.Notifier, -) error { - now := time.Now() - pipelineAtUri := workflowId.PipelineId.AtUri() +// statusEvent builds the eventstream row for a pipeline status. Created is left +// zero so insertEvent stamps the local clock. +func statusEvent(pipelineAtUri, workflow, status string, workflowError *string, exitCode *int64) (eventstream.Event, error) { s := tangled.PipelineStatus{ - CreatedAt: now.Format(time.RFC3339), + CreatedAt: time.Now().Format(time.RFC3339), Error: workflowError, ExitCode: exitCode, - Pipeline: string(pipelineAtUri), - Workflow: workflowId.Name, - Status: string(statusKind), + Pipeline: pipelineAtUri, + Workflow: workflow, + Status: status, } eventJson, err := json.Marshal(s) if err != nil { - return err + return eventstream.Event{}, err } - event := eventstream.Event{ + return eventstream.Event{ Rkey: tid.TID(), Nsid: tangled.PipelineStatusNSID, EventJson: eventJson, + }, nil +} + +func (d *DB) createStatusEvent( + workflowId models.WorkflowId, + statusKind models.StatusKind, + workflowError *string, + exitCode *int64, + n *notifier.Notifier, +) error { + event, err := statusEvent(string(workflowId.PipelineId.AtUri()), workflowId.Name, string(statusKind), workflowError, exitCode) + if err != nil { + return err } + return d.insertEvent(event, n) +} +// InsertRelayedStatus re-authors a status row relayed from an executor into the +// mill's own eventstream. The mill assigns its own Created timestamp (via +// insertEvent passing Created=0) so relayed rows order correctly against the +// mill's cursor, and NotifyAll wakes the mill's /events subscribers exactly +// as a local status write would. +func (d *DB) InsertRelayedStatus( + pipelineAtUri string, + workflow string, + status string, + workflowError *string, + exitCode *int64, + n *notifier.Notifier, +) error { + event, err := statusEvent(pipelineAtUri, workflow, status, workflowError, exitCode) + if err != nil { + return err + } return d.insertEvent(event, n) } diff --git a/spindle/mill/auth_test.go b/spindle/mill/auth_test.go new file mode 100644 index 00000000..8f381e9e --- /dev/null +++ b/spindle/mill/auth_test.go @@ -0,0 +1,263 @@ +package mill + +import ( + "context" + "io" + "log/slog" + "os" + "path/filepath" + "testing" + "time" + + "tangled.org/core/notifier" + "tangled.org/core/spindle/db" + "tangled.org/core/spindle/models" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +func discardLogger() *slog.Logger { + return slog.New(slog.NewTextHandler(io.Discard, nil)) +} + +func nopEncoder() scriptedEncoder { + return scriptedEncoder(func(*millproto.Message) error { return nil }) +} + +// attachSession is the anti-hijack gate: a valid token must not be able to +// displace a live executor that already holds that node identity. Only once the +// incumbent is disconnected (in its reconnect grace window) may a reconnect +// adopt the identity. +func TestAttachSessionRejectsSecondLiveSession(t *testing.T) { + l := discardLogger() + m := New(l, Config{ReconnectGrace: time.Minute}) + + sessionOf := func(node string) *millSession { + m.mu.Lock() + defer m.mu.Unlock() + return m.sessions[node] + } + + sess1 := newSession("node-1", nopEncoder(), l) + if _, ok := m.attachSession(sess1); !ok { + t.Fatal("first attach of a node was rejected; want accept") + } + if sessionOf("node-1") != sess1 { + t.Fatal("first session was not registered as the live session") + } + + // a second attach for the same node while sess1 is still live must be + // rejected and must not evict the incumbent. + sess2 := newSession("node-1", nopEncoder(), l) + if _, ok := m.attachSession(sess2); ok { + t.Fatal("second live attach for an already-live node was accepted; a valid token hijacked the executor") + } + if sessionOf("node-1") != sess1 { + t.Fatal("rejected newcomer evicted the incumbent session") + } + + // once the incumbent is disconnected (reconnect grace), a reconnect adopts + // the identity. + m.detachSession(sess1) + sess3 := newSession("node-1", nopEncoder(), l) + if _, ok := m.attachSession(sess3); !ok { + t.Fatal("attach during the incumbent's reconnect grace was rejected; want adopt") + } + if sessionOf("node-1") != sess3 { + t.Fatal("adopted session was not installed as the live session") + } +} + +// onAttemptResult must only wake the RunStep of a lease the reporting session +// actually owns. A terminal relayed for another node's lease, or for a lease +// that does not exist, must not be delivered. +func TestOnAttemptResultIgnoresForeignLease(t *testing.T) { + l := discardLogger() + m := New(l, Config{ReconnectGrace: time.Minute}) + + foreign := newLease("lease-foreign", "node-a", "dummy") + m.mu.Lock() + m.leases[foreign.id] = foreign + m.mu.Unlock() + + // node-b reports a terminal for a lease owned by node-a. + sessB := newSession("node-b", nopEncoder(), l) + m.attachSession(sessB) + m.onAttemptResult(sessB, &millv1.AttemptResult{ + Offset: 1, + LeaseId: foreign.id, + TerminalStatus: string(models.StatusKindSuccess), + }) + + if _, ok := pollTerminal(foreign); ok { + t.Fatal("attempt-result on a foreign lease delivered a terminal; an executor forged another node's job result") + } + if foreign.getState() == leaseDone { + t.Fatal("attempt-result on a foreign lease sealed the lease") + } +} + +func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) { + l := discardLogger() + m := New(l, Config{ReconnectGrace: time.Minute}) + + // a bystander lease the reporting node does own, to prove an absent-lease + // relay does not spill onto some other lease. + bystander := newLease("lease-bystander", "node-b", "dummy") + m.mu.Lock() + m.leases[bystander.id] = bystander + m.mu.Unlock() + + sessB := newSession("node-b", nopEncoder(), l) + m.attachSession(sessB) + + // no lease with this id exists: must be a no-op (and must not panic). + m.onAttemptResult(sessB, &millv1.AttemptResult{ + Offset: 1, + LeaseId: "lease-nonexistent", + TerminalStatus: string(models.StatusKindSuccess), + }) + + if _, ok := pollTerminal(bystander); ok { + t.Fatal("attempt-result for an absent lease delivered a terminal to a bystander lease") + } + if bystander.getState() == leaseDone { + t.Fatal("attempt-result for an absent lease sealed a bystander lease") + } +} + +// The owned-lease path proves the ignore tests above are not passing merely +// because delivery is broken: a correctly owned terminal IS delivered. +func TestOnAttemptResultDeliversOwnedLease(t *testing.T) { + l := discardLogger() + m := New(l, Config{ReconnectGrace: time.Minute}) + + owned := newLease("lease-owned", "node-b", "dummy") + m.mu.Lock() + m.leases[owned.id] = owned + m.mu.Unlock() + + sessB := newSession("node-b", nopEncoder(), l) + m.attachSession(sessB) + m.onAttemptResult(sessB, &millv1.AttemptResult{ + Offset: 1, + LeaseId: owned.id, + TerminalStatus: string(models.StatusKindSuccess), + }) + + res, ok := pollTerminal(owned) + if !ok { + t.Fatal("attempt-result on an owned lease was not delivered") + } + if got := res.GetTerminalStatus(); got != string(models.StatusKindSuccess) { + t.Fatalf("delivered terminal status = %q, want %q", got, string(models.StatusKindSuccess)) + } + if owned.getState() != leaseDone { + t.Fatal("owned lease was not sealed after its terminal was delivered") + } +} + +// onLogRelay must only write a log line for a lease the reporting session owns. +// A line relayed for another node's lease is dropped; a line for an owned lease +// is written verbatim to that lease's log file. +func TestOnLogRelayRelayOwnership(t *testing.T) { + l := discardLogger() + dir := t.TempDir() + m := New(l, Config{LogDir: dir, ReconnectGrace: time.Minute}) + + foreign := newLease("lease-foreign", "node-x", "dummy") + foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"} + owned := newLease("lease-owned", "node-z", "dummy") + owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"} + m.mu.Lock() + m.leases[foreign.id] = foreign + m.leases[owned.id] = owned + m.mu.Unlock() + + sessY := newSession("node-y", nopEncoder(), l) + m.attachSession(sessY) + sessZ := newSession("node-z", nopEncoder(), l) + m.attachSession(sessZ) + + // node-y relays a log line for a lease owned by node-x: dropped, no file. + m.onLogRelay(sessY, &millv1.LogLine{ + Offset: 1, + LeaseId: foreign.id, + RawJson: []byte(`{"line":"forged"}`), + }) + if _, err := os.Stat(models.LogFilePath(dir, foreign.wid)); !os.IsNotExist(err) { + t.Fatalf("log relay for a foreign lease wrote a file (stat err=%v); an executor forged another node's logs", err) + } + + // node-z relays a log line for its own lease: written verbatim. + line := []byte(`{"line":"legit"}`) + m.onLogRelay(sessZ, &millv1.LogLine{ + Offset: 1, + LeaseId: owned.id, + RawJson: line, + }) + got, err := os.ReadFile(models.LogFilePath(dir, owned.wid)) + if err != nil { + t.Fatalf("owned log relay did not write the log file: %v", err) + } + if want := string(line) + "\n"; string(got) != want { + t.Fatalf("owned log line = %q, want %q", string(got), want) + } +} + +// onStatusRelay must only author a status row for a lease the reporting session +// owns. A status relayed for another node's lease is ignored; a status for an +// owned lease is re-authored under the lease's (authoritative) workflow id. +func TestOnStatusRelayRelayOwnership(t *testing.T) { + ctx := context.Background() + l := discardLogger() + + bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db")) + if err != nil { + t.Fatalf("db.Make: %v", err) + } + t.Cleanup(func() { bdb.Close() }) + n := notifier.New() + + m := New(l, Config{ReconnectGrace: time.Minute}) + m.Attach(bdb, &n) + + foreign := newLease("lease-foreign", "node-x", "dummy") + foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"} + owned := newLease("lease-owned", "node-z", "dummy") + owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"} + m.mu.Lock() + m.leases[foreign.id] = foreign + m.leases[owned.id] = owned + m.mu.Unlock() + + sessY := newSession("node-y", nopEncoder(), l) + m.attachSession(sessY) + sessZ := newSession("node-z", nopEncoder(), l) + m.attachSession(sessZ) + + // node-y relays a status for a lease owned by node-x: no row authored. + m.onStatusRelay(sessY, &millv1.StatusEvent{ + Offset: 1, + LeaseId: foreign.id, + Status: "running", + }) + if _, err := bdb.GetStatus(foreign.wid); err == nil { + t.Fatal("status relay for a foreign lease authored a status row; an executor forged another pipeline's status") + } + + // node-z relays a status for its own lease: authored under the lease's wid. + m.onStatusRelay(sessZ, &millv1.StatusEvent{ + Offset: 1, + LeaseId: owned.id, + Status: "running", + }) + st, err := bdb.GetStatus(owned.wid) + if err != nil { + t.Fatalf("owned status relay did not author a status row: %v", err) + } + if st.Status != "running" { + t.Fatalf("owned status = %q, want %q", st.Status, "running") + } +} diff --git a/spindle/mill/engine.go b/spindle/mill/engine.go new file mode 100644 index 00000000..7dcb6f89 --- /dev/null +++ b/spindle/mill/engine.go @@ -0,0 +1,93 @@ +package mill + +import ( + "context" + "log/slog" + "time" + + "tangled.org/core/api/tangled" + "tangled.org/core/spindle/engine" + "tangled.org/core/spindle/models" + "tangled.org/core/spindle/secrets" +) + +// millWorkflowState is stashed in Workflow.Data (which is `any`, like both real +// engines use it). InitWorkflow parses nothing here: it just carries the raw +// pipeline/workflow forward so the executor can run the real InitWorkflow later. +type millWorkflowState struct { + TargetEngine string + RawWorkflow tangled.Pipeline_Workflow + RawPipeline tangled.Pipeline + Wid models.WorkflowId // stamped at placement (InitWorkflow can't see it) + Lease *RemoteLease +} + +// Engine is the mill's stand-in for a real engine, registered under the real +// engine names ("microvm", "nixery"). All registered names share one Mill. +type Engine struct { + name string + mill *Mill + l *slog.Logger +} + +// NewEngine returns a mill engine view for one engine name. +func NewEngine(name string, mill *Mill) *Engine { + return &Engine{name: name, mill: mill, l: mill.l.With("engine", "mill:"+name)} +} + +// InitWorkflow returns a synthetic one-step workflow so processPipeline injects +// TANGLED_* env and marks pending normally. The real InitWorkflow runs later on +// the executor inside ReserveSeat (intended: it runs twice). +func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) { + return &models.Workflow{ + Name: twf.Name, + Environment: map[string]string{}, + Steps: []models.Step{remoteStep{}}, + Data: &millWorkflowState{ + TargetEngine: e.name, + RawWorkflow: twf, + RawPipeline: tpl, + }, + }, nil +} + +// WorkflowLogger gives StartWorkflows a no-op logger for the mill's synthetic +// workflow: the executor's real log lines are relayed into this wid's log file +// directly (see Loom.onLogRelay), so a local logger here would only write +// competing, step-id-colliding lines. This is engine's optional logger-provider +// seam. +func (e *Engine) WorkflowLogger(wid models.WorkflowId) models.WorkflowLogger { + return models.NullLogger{} +} + +// AcquireWorkflowSlot is the placement seam: it blocks doing remote placement, +// which the user sees as "pending". mode is irrelevant here (only StartWorkflows +// calls this, with Wait); the mill always places. +func (e *Engine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, _ engine.AcquireMode) (engine.WorkflowSlot, error) { + return e.mill.place(ctx, e.name, wid, wf) +} + +// SetupWorkflow is a no-op info log; the real setup happens on the executor. +func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, wfLogger models.WorkflowLogger) error { + e.l.Info("remote job placed, awaiting commit", "wid", wid) + return nil +} + +// RunStep is commit-and-wait: hand over secrets and block on the terminal +// result relayed over the session. +func (e *Engine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, unlocked []secrets.UnlockedSecret, wfLogger models.WorkflowLogger) error { + return e.mill.commitAndWait(ctx, w, unlocked) +} + +// WorkflowTimeout is deliberately generous: the executor enforces the real +// per-engine timeout and relays it. The mill only caps a hung/silent +// executor, with a true death caught separately by reconnect grace. +func (e *Engine) WorkflowTimeout() time.Duration { + return e.mill.cfg.JobTimeout +} + +// DestroyWorkflow cancels a still-running attempt; no-op if already terminal. +func (e *Engine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error { + e.mill.destroy(wid) + return nil +} diff --git a/spindle/mill/handler.go b/spindle/mill/handler.go new file mode 100644 index 00000000..fea599cd --- /dev/null +++ b/spindle/mill/handler.go @@ -0,0 +1,72 @@ +package mill + +import ( + "net/http" + + "github.com/gorilla/websocket" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +var upgrader = websocket.Upgrader{ + ReadBufferSize: 1024, + WriteBufferSize: 1024, +} + +// HandleExecutorConn terminates one executor websocket. Auth is a pre-shared +// secret in the Authorization header, checked before the upgrade so a bad token +// never opens a socket. +func (m *Mill) HandleExecutorConn(w http.ResponseWriter, r *http.Request) { + if m.cfg.SharedSecret != "" { + if r.Header.Get("Authorization") != "Bearer "+m.cfg.SharedSecret { + http.Error(w, "unauthorized", http.StatusUnauthorized) + return + } + } + + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + m.l.Error("fleet ws upgrade failed", "err", err) + return + } + defer conn.Close() + + stream := millproto.NewWSStream(conn) + enc := millproto.NewEncoder(stream) + dec := millproto.NewDecoder(stream) + + // first frame must be Hello. + hello, err := dec.Decode() + if err != nil { + m.l.Error("fleet read hello failed", "err", err) + return + } + h := hello.GetHello() + if h == nil { + m.l.Error("fleet first frame was not hello") + return + } + if h.GetProtocolVersion() != millproto.ProtocolVersion { + m.l.Error("fleet protocol version mismatch", "got", h.GetProtocolVersion(), "want", millproto.ProtocolVersion) + return + } + + sess := newSession(h.GetNodeId(), enc, m.l) + resume, ok := m.attachSession(sess) + if !ok { + m.l.Warn("rejecting duplicate live executor session", "node", sess.nodeID) + return + } + m.l.Info("executor connected", "node", sess.nodeID, "engines", h.GetEngines(), "arch", h.GetArch(), "resume", resume) + + if err := sess.send(&millproto.Message{Resume: &millv1.Resume{AckOffset: resume}}); err != nil { + m.l.Error("fleet send resume failed", "err", err) + m.detachSession(sess) + return + } + m.sessionReady(sess) + + sess.readLoop(m, dec) + m.detachSession(sess) +} diff --git a/spindle/mill/lease.go b/spindle/mill/lease.go new file mode 100644 index 00000000..ffb810fe --- /dev/null +++ b/spindle/mill/lease.go @@ -0,0 +1,205 @@ +package mill + +import ( + "os" + "sync" + + "tangled.org/core/spindle/models" + + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +// leaseState is the mill's view of a remote attempt. It collapses the doc's +// full state diagram to what the mill actually needs to drive cleanup and the +// commit-and-wait in RunStep. +type leaseState int32 + +const ( + // won a bid, executor is holding a seat, not yet committed. + leaseReserved leaseState = iota + // CommitLease has been sent; the executor may already be running the job, + // but the mill has not necessarily received Committed yet. + leaseCommitting + // CommitLease sent and acked; the job is running on the executor. + leaseRunning + // a terminal result arrived (or we gave up). no further action. + leaseDone +) + +type cancelAction int + +const ( + cancelNoop cancelAction = iota + cancelLocal + cancelRemote +) + +// RemoteLease is the mill-side fencing token for one placed job. +type RemoteLease struct { + id string + nodeID string + engine string + wid models.WorkflowId // the job this lease carries; set once placed + + mu sync.Mutex + state leaseState + cancel bool + reason string + terminal chan *millv1.AttemptResult // buffered(1); RunStep waits here + dead chan struct{} // closed when the executor is lost past grace + deadOnce sync.Once + + // log file for relayed lines, opened lazily and held open for the lease's + // life so we don't reopen per line. Guarded by logMu, closed on cleanup. + logMu sync.Mutex + logFile *os.File + logClosed bool +} + +func newLease(id, nodeID, engine string) *RemoteLease { + return &RemoteLease{ + id: id, + nodeID: nodeID, + engine: engine, + state: leaseReserved, + terminal: make(chan *millv1.AttemptResult, 1), + dead: make(chan struct{}), + } +} + +func (l *RemoteLease) setState(s leaseState) { + l.mu.Lock() + l.state = s + l.mu.Unlock() +} + +func (l *RemoteLease) markCommitting() bool { + l.mu.Lock() + defer l.mu.Unlock() + if l.state == leaseDone { + return false + } + if l.state == leaseReserved { + l.state = leaseCommitting + } + return true +} + +func (l *RemoteLease) markRunning() { + l.mu.Lock() + if l.state != leaseDone { + l.state = leaseRunning + } + l.mu.Unlock() +} + +func (l *RemoteLease) getState() leaseState { + l.mu.Lock() + defer l.mu.Unlock() + return l.state +} + +func (l *RemoteLease) requestCancel(reason string) cancelAction { + l.mu.Lock() + defer l.mu.Unlock() + if l.state == leaseDone { + return cancelNoop + } + l.cancel = true + l.reason = reason + if l.state == leaseReserved { + return cancelLocal + } + return cancelRemote +} + +func (l *RemoteLease) cancelRequested() (bool, string) { + l.mu.Lock() + defer l.mu.Unlock() + return l.cancel, l.reason +} + +func (l *RemoteLease) deliverCancelled(reason string) { + l.deliverTerminal(&millv1.AttemptResult{ + LeaseId: l.id, + TerminalStatus: string(models.StatusKindCancelled), + Error: reason, + }) +} + +// deliverTerminal hands the terminal result to a waiting RunStep without +// blocking. Idempotent: a duplicate (e.g. replayed on reconnect) is dropped +// because the channel is buffered to one and the lease is already done. +func (l *RemoteLease) deliverTerminal(res *millv1.AttemptResult) { + l.mu.Lock() + if l.state == leaseDone { + l.mu.Unlock() + return + } + l.state = leaseDone + l.mu.Unlock() + + select { + case l.terminal <- res: + default: + } +} + +// markDead signals that the executor running this lease is gone for good. +func (l *RemoteLease) markDead() { + l.deadOnce.Do(func() { close(l.dead) }) +} + +// writeLog appends a relayed log line to the lease's log file, opening it once +// and holding the handle open for the lease's life (relayed lines for a lease +// all arrive on one session reader goroutine, so this is serialized). A write +// after closeLog is dropped so a replayed line can't resurrect the file. +func (l *RemoteLease) writeLog(path string, line []byte) error { + l.logMu.Lock() + defer l.logMu.Unlock() + if l.logClosed { + return nil + } + if l.logFile == nil { + f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + if err != nil { + return err + } + l.logFile = f + } + _, err := l.logFile.Write(line) + return err +} + +// closeLog closes the lease's log file. Idempotent. +func (l *RemoteLease) closeLog() { + l.logMu.Lock() + defer l.logMu.Unlock() + l.logClosed = true + if l.logFile != nil { + _ = l.logFile.Close() + l.logFile = nil + } +} + +// remoteStep is the mill's synthetic single step. The real steps run on the +// executor; the mill never mirrors them. +type remoteStep struct{} + +func (remoteStep) Name() string { return "remote execution" } +func (remoteStep) Command() string { return "" } +func (remoteStep) Kind() models.StepKind { return models.StepKindSystem } + +// millSlot is what AcquireWorkflowSlot returns; its Release unwinds placement. +type millSlot struct { + fleet *Mill + lease *RemoteLease + once sync.Once +} + +func (s *millSlot) Release() { + if s == nil { + return + } + s.once.Do(func() { s.fleet.releaseSlot(s) }) +} diff --git a/spindle/mill/mill.go b/spindle/mill/mill.go new file mode 100644 index 00000000..b902732d --- /dev/null +++ b/spindle/mill/mill.go @@ -0,0 +1,697 @@ +package mill + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "sort" + "sync" + "time" + + "tangled.org/core/notifier" + "tangled.org/core/spindle/db" + "tangled.org/core/spindle/engine" + "tangled.org/core/spindle/models" + "tangled.org/core/spindle/secrets" + "tangled.org/core/tid" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +const ( + defaultReconnectGrace = 45 * time.Second + defaultJobTimeout = 24 * time.Hour + defaultBidTimeout = 5 * time.Second + defaultTopK = 3 + // re-poll in case a change notification is missed. + placementPoll = 2 * time.Second +) + +type Config struct { + SharedSecret string + MaxPending int + ReconnectGrace time.Duration + JobTimeout time.Duration + BidTimeout time.Duration + TopK int + LogDir string +} + +type Mill struct { + l *slog.Logger + cfg Config + + // re-author relayed status rows into the mill's stream so appview sees + // them unchanged. + db *db.DB + n *notifier.Notifier + + mu sync.Mutex + sessions map[string]*millSession + leases map[string]*RemoteLease + nodeOffset map[string]uint64 + pending int + changeCh chan struct{} // closed + replaced to wake placement waiters + + leaseSeq uint64 +} + +func New(l *slog.Logger, cfg Config) *Mill { + if cfg.ReconnectGrace <= 0 { + cfg.ReconnectGrace = defaultReconnectGrace + } + if cfg.JobTimeout <= 0 { + cfg.JobTimeout = defaultJobTimeout + } + if cfg.BidTimeout <= 0 { + cfg.BidTimeout = defaultBidTimeout + } + if cfg.TopK <= 0 { + cfg.TopK = defaultTopK + } + return &Mill{ + l: l, + cfg: cfg, + sessions: make(map[string]*millSession), + leases: make(map[string]*RemoteLease), + nodeOffset: make(map[string]uint64), + changeCh: make(chan struct{}), + } +} + +func (m *Mill) Attach(d *db.DB, n *notifier.Notifier) { + m.mu.Lock() + m.db = d + m.n = n + m.mu.Unlock() +} + +func (m *Mill) nextLeaseID() string { + m.mu.Lock() + m.leaseSeq++ + seq := m.leaseSeq + m.mu.Unlock() + return fmt.Sprintf("%s-%d", tid.TID(), seq) +} + +func (m *Mill) notifyChange() { + m.mu.Lock() + m.notifyChangeLocked() + m.mu.Unlock() +} + +func (m *Mill) notifyChangeLocked() { + close(m.changeCh) + m.changeCh = make(chan struct{}) +} + +func (m *Mill) currentChangeCh() <-chan struct{} { + m.mu.Lock() + defer m.mu.Unlock() + return m.changeCh +} + +// --- session lifecycle ----------------------------------------------------- + +func (m *Mill) attachSession(sess *millSession) (uint64, bool) { + m.mu.Lock() + defer m.mu.Unlock() + + if old := m.sessions[sess.nodeID]; old != nil { + if !old.disconnected { + // reject a second live session for the same identity so a valid token + // can't hijack an in-flight executor. + return 0, false + } + // the old session is in its reconnect grace window: adopt its leases. + if old.graceTimer != nil { + old.graceTimer.Stop() + } + old.close() + m.l.Info("executor reconnected", "node", sess.nodeID) + } + m.sessions[sess.nodeID] = sess + return m.nodeOffset[sess.nodeID], true +} + +func (m *Mill) detachSession(sess *millSession) { + m.mu.Lock() + if m.sessions[sess.nodeID] != sess { + // already replaced by a reconnect; nothing to do. + m.mu.Unlock() + sess.close() + return + } + sess.disconnected = true + sess.graceTimer = time.AfterFunc(m.cfg.ReconnectGrace, func() { m.onGraceExpired(sess) }) + m.mu.Unlock() + + sess.close() + m.l.Warn("executor session lost; entering reconnect grace", "node", sess.nodeID, "grace", m.cfg.ReconnectGrace) + m.notifyChange() +} + +func (m *Mill) sessionReady(sess *millSession) { + for _, lease := range m.cancelledLeasesForNode(sess.nodeID) { + _, reason := lease.cancelRequested() + m.sendCancel(sess, lease, reason) + } + m.notifyChange() +} + +func (m *Mill) cancelledLeasesForNode(nodeID string) []*RemoteLease { + m.mu.Lock() + var candidates []*RemoteLease + for _, lease := range m.leases { + if lease.nodeID == nodeID { + candidates = append(candidates, lease) + } + } + m.mu.Unlock() + + var leases []*RemoteLease + for _, lease := range candidates { + if cancelled, _ := lease.cancelRequested(); cancelled && lease.getState() != leaseDone { + leases = append(leases, lease) + } + } + return leases +} + +func (m *Mill) onGraceExpired(sess *millSession) { + m.mu.Lock() + if m.sessions[sess.nodeID] != sess || !sess.disconnected { + // reconnected in the meantime. + m.mu.Unlock() + return + } + var dead []*RemoteLease + for _, lease := range m.leases { + if lease.nodeID == sess.nodeID { + dead = append(dead, lease) + } + } + delete(m.sessions, sess.nodeID) + delete(m.nodeOffset, sess.nodeID) + m.mu.Unlock() + + m.l.Warn("executor declared dead; failing its in-flight jobs", "node", sess.nodeID, "jobs", len(dead)) + for _, lease := range dead { + if cancelled, reason := lease.cancelRequested(); cancelled { + lease.deliverCancelled(reason) + } else { + lease.markDead() + } + } + m.notifyChange() +} + +// --- placement (the WorkflowSlotter seam) --------------------------------- + +func (m *Mill) place(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow) (engine.WorkflowSlot, error) { + m.mu.Lock() + if m.cfg.MaxPending > 0 && m.pending >= m.cfg.MaxPending { + max := m.cfg.MaxPending + cur := m.pending + m.mu.Unlock() + return nil, fmt.Errorf("%w: mill has %d pending jobs (max %d)", engine.ErrNoWorkflowSlots, cur, max) + } + m.pending++ + m.mu.Unlock() + defer func() { + m.mu.Lock() + m.pending-- + m.mu.Unlock() + }() + + // stamp wid early so the executor reconstructs logs/status under the same + // pipeline id the mill used for "pending". + if st, ok := wf.Data.(*millWorkflowState); ok && st != nil { + st.Wid = wid + } + + for { + if err := ctx.Err(); err != nil { + return nil, err + } + + lease, err := m.bid(ctx, engineName, wf) + if err != nil { + return nil, err + } + if lease != nil { + lease.wid = wid + m.mu.Lock() + m.leases[lease.id] = lease + if st, ok := wf.Data.(*millWorkflowState); ok && st != nil { + st.Lease = lease + } + m.mu.Unlock() + return &millSlot{fleet: m, lease: lease}, nil + } + + // nobody to place on right now: wait for a change, ctx, or a poll tick. + m.mu.Lock() + ch := m.changeCh + m.mu.Unlock() + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-ch: + case <-time.After(placementPoll): + } + } +} + +func (m *Mill) bid(ctx context.Context, engineName string, wf *models.Workflow) (*RemoteLease, error) { + rawPipeline, rawWorkflow, err := marshalJob(wf) + if err != nil { + return nil, err + } + + candidates := m.rankCandidates(engineName) + if len(candidates) == 0 { + return nil, nil + } + if len(candidates) > m.cfg.TopK { + candidates = candidates[:m.cfg.TopK] + } + + bidCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout) + defer cancel() + + type bidResult struct { + sess *millSession + lease *RemoteLease + score float64 + } + results := make(chan bidResult, len(candidates)) + var wg sync.WaitGroup + for _, sess := range candidates { + wg.Add(1) + go func(sess *millSession) { + defer wg.Done() + leaseID := m.nextLeaseID() + lease := newLease(leaseID, sess.nodeID, engineName) + wid := m.jobWid(wf) + msg := &millproto.Message{ReserveSeat: &millv1.ReserveSeat{ + LeaseId: leaseID, + TargetEngine: engineName, + RawPipelineJson: rawPipeline, + RawWorkflowJson: rawWorkflow, + Knot: wid.Knot, + Rkey: wid.Rkey, + TtlSeconds: uint32(m.cfg.ReconnectGrace / time.Second), + }} + resp, err := sess.request(bidCtx, leaseID, msg) + if err != nil { + return + } + rr := resp.GetReserveResult() + if rr == nil || !rr.GetAccepted() { + return + } + results <- bidResult{sess: sess, lease: lease, score: rr.GetScore()} + }(sess) + } + wg.Wait() + close(results) + + var winner *bidResult + var losers []*RemoteLease + for r := range results { + if winner == nil || r.score > winner.score { + if winner != nil { + losers = append(losers, winner.lease) + } + winner = &r + } else { + losers = append(losers, r.lease) + } + } + + // release the also-rans so they free their held seats immediately. + for _, l := range losers { + m.releaseRemote(l) + } + + if winner == nil { + return nil, nil + } + return winner.lease, nil +} + +func (m *Mill) rankCandidates(engineName string) []*millSession { + m.mu.Lock() + defer m.mu.Unlock() + + type ranked struct { + sess *millSession + seats uint32 + } + var rs []ranked + for _, s := range m.sessions { + if s.disconnected { + continue + } + if s.snapshot == nil { + continue + } + es, ok := s.snapshot.GetEngines()[engineName] + if !ok || es.GetFreeSeats() == 0 { + continue + } + rs = append(rs, ranked{sess: s, seats: es.GetFreeSeats()}) + } + sort.SliceStable(rs, func(i, j int) bool { return rs[i].seats > rs[j].seats }) + + out := make([]*millSession, len(rs)) + for i := range rs { + out[i] = rs[i].sess + } + return out +} + +// --- commit + wait (RunStep) ---------------------------------------------- + +func (m *Mill) commitAndWait(ctx context.Context, wf *models.Workflow, unlocked []secrets.UnlockedSecret) error { + st, ok := wf.Data.(*millWorkflowState) + if !ok || st == nil || st.Lease == nil { + return fmt.Errorf("mill workflow state missing lease") + } + lease := st.Lease + + pbSecrets := make([]*millv1.Secret, len(unlocked)) + for i, s := range unlocked { + pbSecrets[i] = &millv1.Secret{Key: s.Key, Value: s.Value} + } + + commit := &millproto.Message{CommitLease: &millv1.CommitLease{ + LeaseId: lease.id, + Secrets: pbSecrets, + }} + + for { + if res, ok := pollTerminal(lease); ok { + return terminalError(res.GetTerminalStatus()) + } + if !lease.markCommitting() { + return engine.ErrWorkflowFailed + } + + sess := m.sessionForNode(lease.nodeID) + if sess == nil { + if done, err := m.waitCommitRetry(ctx, lease); done || err != nil { + return err + } + continue + } + + reqCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout) + resp, err := sess.request(reqCtx, lease.id, commit) + cancel() + if err != nil { + switch { + case errors.Is(err, errSessionClosed): + if done, err := m.waitCommitRetry(ctx, lease); done || err != nil { + return err + } + continue + case errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil: + continue + case errors.Is(err, context.DeadlineExceeded): + return engine.ErrTimedOut + case errors.Is(err, context.Canceled): + return err + default: + m.l.Warn("commit lease send failed; waiting for reconnect", "lease", lease.id, "node", lease.nodeID, "err", err) + if done, err := m.waitCommitRetry(ctx, lease); done || err != nil { + return err + } + continue + } + } + if resp.GetCommitted() == nil { + return engine.ErrWorkflowFailed + } + lease.markRunning() + if cancelled, reason := lease.cancelRequested(); cancelled { + m.sendCancel(sess, lease, reason) + } + break + } + + select { + case res := <-lease.terminal: + return terminalError(res.GetTerminalStatus()) + case <-lease.dead: + m.l.Warn("executor lost while running job", "lease", lease.id, "node", lease.nodeID) + return engine.ErrWorkflowFailed + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + return engine.ErrTimedOut + } + return ctx.Err() + } +} + +func (m *Mill) waitCommitRetry(ctx context.Context, lease *RemoteLease) (bool, error) { + res, err := m.waitLeaseChange(ctx, lease) + if err != nil { + return true, err + } + if res != nil { + return true, terminalError(res.GetTerminalStatus()) + } + return false, nil +} + +func terminalError(status string) error { + switch models.StatusKind(status) { + case models.StatusKindSuccess: + return nil + case models.StatusKindTimeout: + return engine.ErrTimedOut + case models.StatusKindCancelled: + return engine.ErrCancelled + default: + return engine.ErrWorkflowFailed + } +} + +func pollTerminal(lease *RemoteLease) (*millv1.AttemptResult, bool) { + select { + case res := <-lease.terminal: + return res, true + default: + return nil, false + } +} + +func (m *Mill) waitLeaseChange(ctx context.Context, lease *RemoteLease) (*millv1.AttemptResult, error) { + ch := m.currentChangeCh() + select { + case res := <-lease.terminal: + return res, nil + case <-lease.dead: + return nil, engine.ErrWorkflowFailed + case <-ctx.Done(): + if ctx.Err() == context.DeadlineExceeded { + return nil, engine.ErrTimedOut + } + return nil, ctx.Err() + case <-ch: + return nil, nil + case <-time.After(placementPoll): + return nil, nil + } +} + +// --- cleanup --------------------------------------------------------------- + +func (m *Mill) destroy(wid models.WorkflowId) { + m.mu.Lock() + var lease *RemoteLease + for _, l := range m.leases { + if l.wid == wid { + lease = l + break + } + } + m.mu.Unlock() + if lease == nil { + return + } + reason := "workflow destroyed" + switch lease.requestCancel(reason) { + case cancelLocal: + if sess := m.sessionForNode(lease.nodeID); sess != nil { + _ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}}) + } + lease.deliverCancelled(reason) + case cancelRemote: + if sess := m.sessionForNode(lease.nodeID); sess != nil { + m.sendCancel(sess, lease, reason) + } + } +} + +func (m *Mill) releaseSlot(s *millSlot) { + lease := s.lease + if lease.getState() == leaseReserved { + m.releaseRemote(lease) + } + lease.closeLog() + + m.mu.Lock() + delete(m.leases, lease.id) + m.mu.Unlock() + + m.notifyChange() +} + +func (m *Mill) releaseRemote(lease *RemoteLease) { + lease.setState(leaseDone) + if sess := m.sessionForNode(lease.nodeID); sess != nil { + _ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}}) + } +} + +func (m *Mill) sendCancel(sess *millSession, lease *RemoteLease, reason string) { + _ = sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{ + LeaseId: lease.id, + Reason: reason, + }}) +} + +func (m *Mill) sessionForNode(nodeID string) *millSession { + m.mu.Lock() + defer m.mu.Unlock() + s := m.sessions[nodeID] + if s == nil || s.disconnected { + return nil + } + return s +} + +// --- relay handlers -------------------------------------------------------- + +func (m *Mill) shouldAcceptOffset(sess *millSession, offset uint64) bool { + m.mu.Lock() + defer m.mu.Unlock() + current := m.nodeOffset[sess.nodeID] + return offset == current+1 +} + +func (m *Mill) ackOffset(sess *millSession, offset uint64) { + m.mu.Lock() + if offset > m.nodeOffset[sess.nodeID] { + m.nodeOffset[sess.nodeID] = offset + } + m.mu.Unlock() + + _ = sess.send(&millproto.Message{Ack: &millv1.Ack{UpToOffset: offset}}) +} + +func (m *Mill) processRelay(sess *millSession, offset uint64, kind string, apply func() error) { + if !m.shouldAcceptOffset(sess, offset) { + return + } + // ack only after the side effect lands, or a transient write error makes + // the message unreplayable. + if err := apply(); err != nil { + m.l.Error("process relayed message failed", "kind", kind, "node", sess.nodeID, "offset", offset, "err", err) + return + } + m.ackOffset(sess, offset) +} + +func (m *Mill) onSnapshot(sess *millSession, snap *millv1.NodeSnapshot) { + m.mu.Lock() + sess.snapshot = snap + m.mu.Unlock() + m.notifyChange() +} + +func (m *Mill) onStatusRelay(sess *millSession, ev *millv1.StatusEvent) { + m.processRelay(sess, ev.GetOffset(), "status", func() error { + if m.db == nil { + return fmt.Errorf("mill db not attached") + } + m.mu.Lock() + lease := m.leases[ev.GetLeaseId()] + m.mu.Unlock() + if lease == nil || lease.nodeID != sess.nodeID { + // not a lease this executor holds; don't let it author status rows for + // arbitrary pipelines. + return nil + } + var errMsg *string + if e := ev.GetError(); e != "" { + errMsg = &e + } + var exit *int64 + if c := ev.GetExitCode(); c != 0 { + exit = &c + } + // authoritative pipeline/workflow come from the lease, not whatever the + // executor sent. + return m.db.InsertRelayedStatus(string(lease.wid.PipelineId.AtUri()), lease.wid.Name, ev.GetStatus(), errMsg, exit, m.n) + }) +} + +func (m *Mill) onLogRelay(sess *millSession, ll *millv1.LogLine) { + m.processRelay(sess, ll.GetOffset(), "log", func() error { + m.mu.Lock() + lease := m.leases[ll.GetLeaseId()] + m.mu.Unlock() + if lease == nil || lease.nodeID != sess.nodeID { + return nil + } + path := models.LogFilePath(m.cfg.LogDir, lease.wid) + line := append(ll.GetRawJson(), '\n') + return lease.writeLog(path, line) + }) +} + +func (m *Mill) onAttemptResult(sess *millSession, ar *millv1.AttemptResult) { + m.processRelay(sess, ar.GetOffset(), "attempt-result", func() error { + m.mu.Lock() + lease := m.leases[ar.GetLeaseId()] + m.mu.Unlock() + if lease == nil || lease.nodeID != sess.nodeID { + return nil + } + lease.deliverTerminal(ar) + return nil + }) +} + +// --- helpers --------------------------------------------------------------- + +func marshalJob(wf *models.Workflow) (pipeline string, workflow string, err error) { + st, ok := wf.Data.(*millWorkflowState) + if !ok || st == nil { + return "", "", fmt.Errorf("mill workflow state missing") + } + p, err := json.Marshal(st.RawPipeline) + if err != nil { + return "", "", fmt.Errorf("marshal pipeline: %w", err) + } + w, err := json.Marshal(st.RawWorkflow) + if err != nil { + return "", "", fmt.Errorf("marshal workflow: %w", err) + } + return string(p), string(w), nil +} + +func (m *Mill) jobWid(wf *models.Workflow) models.WorkflowId { + st, ok := wf.Data.(*millWorkflowState) + if !ok || st == nil { + return models.WorkflowId{} + } + return st.Wid +} diff --git a/spindle/mill/mill_test.go b/spindle/mill/mill_test.go new file mode 100644 index 00000000..35416d9a --- /dev/null +++ b/spindle/mill/mill_test.go @@ -0,0 +1,154 @@ +package mill + +import ( + "context" + "errors" + "io" + "log/slog" + "testing" + "time" + + "tangled.org/core/api/tangled" + "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" +) + +type scriptedEncoder func(*millproto.Message) error + +func (e scriptedEncoder) Encode(msg *millproto.Message) error { return e(msg) } + +func testWorkflow(name string) *models.Workflow { + return &models.Workflow{ + Name: name, + Environment: map[string]string{}, + Steps: []models.Step{remoteStep{}}, + Data: &millWorkflowState{ + TargetEngine: "dummy", + RawWorkflow: tangled.Pipeline_Workflow{Name: name}, + RawPipeline: tangled.Pipeline{}, + }, + } +} + +func TestCommitRetriesAfterSessionCloseBeforeCommitted(t *testing.T) { + l := slog.New(slog.NewTextHandler(io.Discard, nil)) + m := New(l, Config{BidTimeout: 25 * time.Millisecond, ReconnectGrace: time.Second}) + wf := testWorkflow("build") + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"} + lease := newLease("lease-1", "node-1", "dummy") + lease.wid = wid + wf.Data.(*millWorkflowState).Lease = lease + + m.mu.Lock() + m.leases[lease.id] = lease + m.mu.Unlock() + + var sess1 *millSession + firstCommit := make(chan struct{}) + sess1 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error { + if msg.GetCommitLease() != nil { + close(firstCommit) + m.detachSession(sess1) + } + return nil + }), l) + m.attachSession(sess1) + + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + done := make(chan error, 1) + go func() { done <- m.commitAndWait(ctx, wf, nil) }() + + select { + case <-firstCommit: + case <-ctx.Done(): + t.Fatal("first commit was not sent") + } + + var sess2 *millSession + sess2 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error { + if msg.GetCommitLease() == nil { + return nil + } + leaseID := msg.GetCommitLease().GetLeaseId() + sess2.deliver(leaseID, &millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}}) + m.onAttemptResult(sess2, &millv1.AttemptResult{ + Offset: 1, + LeaseId: leaseID, + TerminalStatus: string(models.StatusKindSuccess), + }) + return nil + }), l) + m.attachSession(sess2) + m.sessionReady(sess2) + + select { + case err := <-done: + if err != nil { + t.Fatalf("commitAndWait() error = %v, want success after reconnect", err) + } + case <-ctx.Done(): + t.Fatal("commitAndWait() did not finish after reconnect") + } +} + +func TestDestroyRunningLeaseDoesNotDropCancelledTerminal(t *testing.T) { + l := slog.New(slog.NewTextHandler(io.Discard, nil)) + m := New(l, Config{}) + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"} + lease := newLease("lease-1", "node-1", "dummy") + lease.wid = wid + lease.setState(leaseRunning) + + m.mu.Lock() + m.leases[lease.id] = lease + m.mu.Unlock() + + m.destroy(wid) + if lease.getState() == leaseDone { + t.Fatal("destroy sealed the lease before the terminal result") + } + + lease.deliverTerminal(&millv1.AttemptResult{ + LeaseId: lease.id, + TerminalStatus: string(models.StatusKindCancelled), + }) + res := <-lease.terminal + if err := terminalError(res.GetTerminalStatus()); !errors.Is(err, engine.ErrCancelled) { + t.Fatalf("terminalError() = %v, want ErrCancelled", err) + } +} + +func TestPlaceBlocksWhenNoCapacity(t *testing.T) { + l := slog.New(slog.NewTextHandler(io.Discard, nil)) + m := New(l, Config{}) + wf := testWorkflow("build") + wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"} + + // no executors at all: place must block until ctx expires (user sees pending). + ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond) + defer cancel() + + _, err := m.place(ctx, "dummy", wid, wf) + if err != context.DeadlineExceeded { + t.Fatalf("place() error = %v, want DeadlineExceeded", err) + } +} + +func TestMaxPendingRejects(t *testing.T) { + l := slog.New(slog.NewTextHandler(io.Discard, nil)) + m := New(l, Config{MaxPending: 1}) + + m.mu.Lock() + m.pending = 1 + m.mu.Unlock() + + wf2 := testWorkflow("b") + _, err := m.place(context.Background(), "dummy", models.WorkflowId{Name: "b"}, wf2) + if err == nil { + t.Fatal("place() past maxPending should error") + } +} diff --git a/spindle/mill/session.go b/spindle/mill/session.go new file mode 100644 index 00000000..e1959a81 --- /dev/null +++ b/spindle/mill/session.go @@ -0,0 +1,139 @@ +package mill + +import ( + "context" + "errors" + "log/slog" + "sync" + "time" + + millproto "tangled.org/core/spindle/mill/proto" + millv1 "tangled.org/core/spindle/mill/proto/gen" +) + +var errSessionClosed = errors.New("mill: executor session closed") + +// millSession is one live websocket to an executor. Unlike the microVM +// AgentSession it is multiplexed: many leases plus async status/log relay share +// it, so a single reader goroutine demuxes by message type and correlates +// request/response by lease id. We never hold a lock across a decode. +type millSession struct { + nodeID string + enc messageEncoder + l *slog.Logger + + // snapshot, disconnected and graceTimer are guarded by Mill.mu, not the + // session mutex below (the fleet ranks across sessions under its own lock). + snapshot *millv1.NodeSnapshot + disconnected bool + graceTimer *time.Timer + + mu sync.Mutex + pending map[string]chan *millproto.Message // lease_id -> response waiter + + closeOnce sync.Once + closed chan struct{} +} + +type messageEncoder interface { + Encode(*millproto.Message) error +} + +func newSession(nodeID string, enc messageEncoder, l *slog.Logger) *millSession { + return &millSession{ + nodeID: nodeID, + enc: enc, + l: l, + pending: make(map[string]chan *millproto.Message), + closed: make(chan struct{}), + } +} + +func (s *millSession) send(msg *millproto.Message) error { + return s.enc.Encode(msg) +} + +func (s *millSession) close() { + s.closeOnce.Do(func() { close(s.closed) }) +} + +// await registers a one-shot waiter for the next response addressed to leaseID, +// returning the channel and a cancel that unregisters it. +func (s *millSession) await(leaseID string) (<-chan *millproto.Message, func()) { + ch := make(chan *millproto.Message, 1) + s.mu.Lock() + s.pending[leaseID] = ch + s.mu.Unlock() + return ch, func() { + s.mu.Lock() + if s.pending[leaseID] == ch { + delete(s.pending, leaseID) + } + s.mu.Unlock() + } +} + +// deliver routes a correlated response to a waiter, if any is registered. +func (s *millSession) deliver(leaseID string, msg *millproto.Message) { + s.mu.Lock() + ch := s.pending[leaseID] + delete(s.pending, leaseID) + s.mu.Unlock() + if ch != nil { + select { + case ch <- msg: + default: + } + } +} + +// request sends msg and waits for the correlated response addressed to leaseID, +// honouring ctx and session close. +func (s *millSession) request(ctx context.Context, leaseID string, msg *millproto.Message) (*millproto.Message, error) { + ch, cancel := s.await(leaseID) + defer cancel() + + if err := s.send(msg); err != nil { + return nil, err + } + + select { + case resp := <-ch: + return resp, nil + case <-ctx.Done(): + return nil, ctx.Err() + case <-s.closed: + return nil, errSessionClosed + } +} + +// readLoop demuxes incoming frames until the decoder errors (connection gone). +func (s *millSession) readLoop(m *Mill, dec *millproto.Decoder) { + for { + msg, err := dec.Decode() + if err != nil { + s.l.Debug("session read ended", "node", s.nodeID, "err", err) + return + } + s.dispatch(m, msg) + } +} + +func (s *millSession) dispatch(m *Mill, msg *millproto.Message) { + switch { + case msg.GetNodeSnapshot() != nil: + m.onSnapshot(s, msg.GetNodeSnapshot()) + case msg.GetReserveResult() != nil: + s.deliver(msg.GetReserveResult().GetLeaseId(), msg) + case msg.GetCommitted() != nil: + s.deliver(msg.GetCommitted().GetLeaseId(), msg) + case msg.GetStatusEvent() != nil: + m.onStatusRelay(s, msg.GetStatusEvent()) + case msg.GetLogLine() != nil: + m.onLogRelay(s, msg.GetLogLine()) + case msg.GetAttemptResult() != nil: + m.onAttemptResult(s, msg.GetAttemptResult()) + default: + s.l.Warn("session received unexpected message", "node", s.nodeID) + } +}