Monorepo for Tangled
Something went wrong. Try again.
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187package mill
import ( "context" "encoding/json" "errors" "fmt" "log/slog" "os" "path/filepath" "slices" "strings" "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 defaultMaxPending = 100 defaultQuarantineStrikes = 3)
// marks session-ending errors that are the executor's fault, enough of them// in a row gets the node quarantinedvar errProtocolViolation = errors.New("executor protocol violation")
func protoErrf(format string, args ...any) error { return fmt.Errorf("%w: %s", errProtocolViolation, fmt.Sprintf(format, args...))}
type Config struct { // mill appends live-tailed executor lines here so logview can follow running remote jobs LogDir string MaxPending int ReconnectGrace time.Duration JobTimeout time.Duration BidTimeout time.Duration TopK int CancelTimeout time.Duration // how many protocol-violating session deaths in a row quarantine the node QuarantineStrikes int}
type Mill struct { l *slog.Logger cfg Config
db *db.DB n *notifier.Notifier
mu sync.Mutex sessions map[string]*millSession leases map[string]*RemoteLease reservations map[string]*RemoteLease nodeSeqno map[string]uint64 protoStrikes map[string]int pending int changeCh chan struct{} // closed and 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 } if cfg.MaxPending <= 0 { cfg.MaxPending = defaultMaxPending } if cfg.CancelTimeout <= 0 { cfg.CancelTimeout = 10 * time.Second } if cfg.QuarantineStrikes <= 0 { cfg.QuarantineStrikes = defaultQuarantineStrikes } return &Mill{ l: l, cfg: cfg, sessions: make(map[string]*millSession), leases: make(map[string]*RemoteLease), reservations: make(map[string]*RemoteLease), nodeSeqno: make(map[string]uint64), protoStrikes: make(map[string]int), 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) dropReservation(id string) { m.mu.Lock() delete(m.reservations, id) m.mu.Unlock()}
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}
func (m *Mill) attachSession(sess *millSession) (uint64, bool) { // read before taking mu so the db never runs under the mill lock, this // also picks up operator cursor resets without a mill restart var dbCursor uint64 haveCursor := false if m.db != nil { if cur, err := m.db.GetExecutorCursor(sess.nodeID, sess.epoch); err == nil { dbCursor, haveCursor = cur, true } }
m.mu.Lock() defer m.mu.Unlock()
if old := m.sessions[sess.nodeID]; old != nil { if old.live(m.cfg.ReconnectGrace) { // a second live session for the same identity is a hijack // attempt, reject it return 0, false } if old.graceTimer != nil { old.graceTimer.Stop() } old.close() if old.disconnected { m.l.Info("executor reconnected", "node", sess.nodeID) } else { m.l.Warn("replacing silent executor session", "node", sess.nodeID) } } m.sessions[sess.nodeID] = sess // wakes commit retries waiting out reconnect grace m.notifyChangeLocked() key := sess.nodeID + "/" + sess.epoch if haveCursor { m.nodeSeqno[key] = dbCursor } return m.nodeSeqno[key], true}
func (m *Mill) touchSession(sess *millSession) bool { m.mu.Lock() defer m.mu.Unlock() if m.sessions[sess.nodeID] != sess || sess.disconnected { return false } sess.lastSeen = time.Now() return true}
func (m *Mill) detachSession(sess *millSession) { m.mu.Lock() if m.sessions[sess.nodeID] != sess { // already replaced by a reconnect m.mu.Unlock() sess.close() return } sess.disconnected = true sess.graceTimer = time.AfterFunc(m.cfg.ReconnectGrace, func() { m.failLeasesAfterGrace(sess) }) m.mu.Unlock()
sess.close() m.l.Warn("executor session lost; entering reconnect grace", "node", sess.nodeID, "grace", m.cfg.ReconnectGrace) m.notifyChange()}
// an executor whose sessions keep dying on protocol errors is stuck, eg. an// unrecoverable stream gap, and its reconnects keep resetting grace so its// leases never fail. after enough strikes quarantine it until an operator// fixes its stream. clean deaths reset the countfunc (m *Mill) noteSessionError(sess *millSession, err error) { if !errors.Is(err, errProtocolViolation) { m.mu.Lock() delete(m.protoStrikes, sess.nodeID) m.mu.Unlock() return }
m.mu.Lock() m.protoStrikes[sess.nodeID]++ strikes := m.protoStrikes[sess.nodeID] m.mu.Unlock()
m.l.Warn("executor session died on protocol violation", "node", sess.nodeID, "strikes", strikes, "err", err) if strikes < m.cfg.QuarantineStrikes || m.db == nil { return }
reason := fmt.Sprintf("repeated protocol violations (%d): %v", strikes, err) if err := m.db.QuarantineExecutor(sess.nodeID, reason); err != nil { m.l.Error("failed to quarantine executor", "node", sess.nodeID, "err", err) return } m.l.Warn("executor quarantined; reset its stream state and unquarantine to readmit", "node", sess.nodeID)}
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) hasLiveExecutor(nodeID string) bool { m.mu.Lock() defer m.mu.Unlock() sess := m.sessions[nodeID] return sess != nil && sess.live(m.cfg.ReconnectGrace)}
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}
// reconnect grace ran out without the executor coming back. fails every// lease the node held, and reschedules itself if any failure didn't stickfunc (m *Mill) failLeasesAfterGrace(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) } } m.mu.Unlock()
m.l.Warn("executor declared dead; failing its in-flight jobs", "node", sess.nodeID, "jobs", len(dead)) deadReason := "executor lost" success := true for _, lease := range dead { switch { case lease.orphaned: // restored leases have no RunStep waiting, fail them straight // into the event stream if err := m.finishOrphan(lease, string(models.StatusKindFailed), &deadReason, nil); err != nil { m.l.Error("finish orphan after executor loss failed, will retry", "lease", lease.id, "err", err) success = false } default: // live leases have a blocked RunStep, keep a pending cancellation // as the terminal reason status := string(models.StatusKindFailed) reason := deadReason if cancelled, cancelReason := lease.cancelRequested(); cancelled { status = string(models.StatusKindCancelled) reason = cancelReason } if err := m.finishLiveLease(lease, status, reason); err != nil { m.l.Error("finish lease after executor loss failed, will retry", "lease", lease.id, "err", err) success = false } } }
m.mu.Lock() if !success { // something didn't finish cleanly, try again shortly. finished // leases skip themselves on the next pass sess.graceTimer = time.AfterFunc(5*time.Second, func() { m.failLeasesAfterGrace(sess) }) m.mu.Unlock() return }
// everything failed cleanly. forget the node and wake placement if m.sessions[sess.nodeID] == sess { delete(m.sessions, sess.nodeID) } m.mu.Unlock() m.notifyChange() m.sweepUnclaimedOrphans()}
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() }()
for { if err := ctx.Err(); err != nil { return nil, err }
// grab the channel before bidding. a change mid-bid closes it, so // the wait below re-bids right away ch := m.currentChangeCh()
lease, err := m.bid(ctx, engineName, wid, wf) if err != nil { return nil, err } if lease != nil { lease.wid = wid if err := m.persistLease(lease, leaseRowReserved); err != nil { m.releaseRemote(lease) m.mu.Lock() delete(m.reservations, lease.id) m.mu.Unlock() return nil, fmt.Errorf("persist reserved mill lease: %w", err) } m.mu.Lock() delete(m.reservations, lease.id) 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 }
// no executor available. wait for a change or ctx select { case <-ctx.Done(): return nil, ctx.Err() case <-ch: } }}
func (m *Mill) bid(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow) (*RemoteLease, error) { rawPipeline, rawWorkflow, err := marshalJob(wf) if err != nil { return nil, err }
requiredLabels := requiredLabels(wf) candidates := m.rankCandidates(engineName, requiredLabels) if len(candidates) == 0 { return nil, nil }
type bidResult struct { sess *millSession lease *RemoteLease rank int incompatible bool reason string } limit := m.cfg.TopK if limit <= 0 { limit = len(candidates) } if len(candidates) > limit { candidates = candidates[:limit] } results := make(chan bidResult, limit) ask := func(rank int, sess *millSession) { bidCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout) defer cancel() leaseID := m.nextLeaseID() lease := newLease(leaseID, sess.nodeID, sess.epoch, engineName) m.mu.Lock() m.reservations[leaseID] = lease m.mu.Unlock() 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 { m.dropReservation(leaseID) results <- bidResult{} return } rr := resp.GetReserveResult() if rr == nil { m.dropReservation(leaseID) results <- bidResult{} return } if !rr.GetAccepted() { m.dropReservation(leaseID) if rr.GetRejectClass() == millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE { results <- bidResult{sess: sess, rank: rank, incompatible: true, reason: rr.GetRejectReason()} return } results <- bidResult{} return } results <- bidResult{sess: sess, lease: lease, rank: rank} } next := 0 inFlight := 0 for next < len(candidates) && inFlight < limit { inFlight++ go ask(next, candidates[next]) next++ }
var winner *bidResult var losers []*RemoteLease var incompatible []string // any soft reject (transient or timeout) means the fleet was just // busy, so an all-incompatible outcome isn't a hard placement error softReject := false for inFlight > 0 { r := <-results inFlight-- // incompatible rejects get reported to the user, any other failure // just means the fleet is busy if r.incompatible { if r.reason != "" { incompatible = append(incompatible, r.reason) } } else if r.lease == nil { softReject = true } // a failed bid means ask the next candidate, unless someone won if r.lease == nil { for winner == nil && next < len(candidates) && inFlight < limit { inFlight++ go ask(next, candidates[next]) next++ } continue } // if this bid is worse than the winner it goes to the losers pile if winner != nil && r.rank >= winner.rank { losers = append(losers, r.lease) continue } // otherwise it's the new best and the old winner joins the losers if winner != nil { losers = append(losers, winner.lease) } winner = &r }
// let the losers go so they free their seats right away for _, l := range losers { m.dropReservation(l.id) m.releaseRemote(l) }
if winner == nil { if len(incompatible) > 0 && !softReject { return nil, fmt.Errorf("no compatible executor for %s: %s", engineName, strings.Join(incompatible, "; ")) } return nil, nil } return winner.lease, nil}
// ranks nodes that are least busy first. if a resource is used a lot// then that node will lose to one that is more even across the board.func (m *Mill) rankCandidates(engineName string, requiredLabels []string) []*millSession { m.mu.Lock() defer m.mu.Unlock()
type ranked struct { sess *millSession worst float64 sum float64 } var rs []ranked for _, s := range m.sessions { // only live, reporting sessions can take work if s.disconnected { continue } if s.snapshot == nil { continue } // the engine has to exist and have room right now ea, ok := s.snapshot.GetEngines()[engineName] if !ok || !ea.GetAvailable() { continue } // and satisfy the wf's label requirements if !hasLabels(s.labels, requiredLabels) { continue } worst, sum := loadScore(ea.GetLoad()) rs = append(rs, ranked{sess: s, worst: worst, sum: sum}) } slices.SortStableFunc(rs, func(a, b ranked) int { if a.worst < b.worst { return -1 } if a.worst > b.worst { return 1 } if a.sum < b.sum { return -1 } if a.sum > b.sum { return 1 } return 0 })
out := make([]*millSession, len(rs)) for i := range rs { out[i] = rs[i].sess } return out}
func loadScore(load map[string]float64) (worst, sum float64) { for _, v := range load { if v > worst { worst = v } sum += v } return worst, sum}
func requiredLabels(wf *models.Workflow) []string { st, ok := wf.Data.(*millWorkflowState) if !ok || st == nil { return nil } return st.RawWorkflow.RunsOn}
func hasLabels(labels []string, required []string) bool { for _, want := range required { if !slices.Contains(labels, want) { return false } } return true}
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, }}
// commit retries ride reconnects, a reservation outlives one // disconnect. lost session or slow executor just means wait and retry, // only job timeout or a dead lease stops the loop for { if res, ok := pollTerminal(lease); ok { return terminalError(res.Status) } 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): // session died mid-request. wait out the grace, then retry on // the new one if done, err := m.waitCommitRetry(ctx, lease); done || err != nil { return err } continue case errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil: // executor didn't answer in time, but its seat is still held so // retrying is safe continue case errors.Is(err, context.DeadlineExceeded): // the job ctx itself ran out, a real timeout 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 err := m.persistLease(lease, leaseRowRunning); err != nil { m.l.Error("persist running mill lease", "lease", lease.id, "err", err) } if cancelled, reason := lease.cancelRequested(); cancelled { m.sendCancel(sess, lease, reason) } break }
select { case res := <-lease.terminal: return terminalError(res.Status) 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) { // grab the channel before checking for a live session again // a reconnect will still close the channel if it happens in between ch := m.currentChangeCh() if m.sessionForNode(lease.nodeID) != nil { return false, nil } select { case res := <-lease.terminal: return true, terminalError(res.Status) case <-ctx.Done(): if ctx.Err() == context.DeadlineExceeded { return true, engine.ErrTimedOut } return true, ctx.Err() case <-ch: return false, nil }}
func terminalError(status millv1.TerminalStatus) error { switch status { case millv1.TerminalStatus_SUCCESS: return nil case millv1.TerminalStatus_TIMEOUT: return engine.ErrTimedOut case millv1.TerminalStatus_CANCELLED: return engine.ErrWorkflowCanceled 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) 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 state, cancelled := lease.releaseState() if state == leaseReserved { m.releaseRemote(lease) } else if cancelled && state != leaseDone { return } if err := m.cleanupLease(lease); err != nil { m.l.Error("releaseSlot cleanupLease failed", "lease", lease.id, "err", err) }}func (m *Mill) cleanupLease(lease *RemoteLease) error { lease.finishMu.Lock() defer lease.finishMu.Unlock() return m.cleanupLeaseLocked(lease)}
func (m *Mill) cleanupLeaseLocked(lease *RemoteLease) error { if lease.cleanedUp { return nil } if m.db != nil { if err := m.db.DeleteMillLease(lease.id); err != nil { m.scheduleCleanupLocked(lease) return err } } m.mu.Lock() delete(m.leases, lease.id) m.mu.Unlock() lease.cleanedUp = true m.notifyChange() return nil}
func (m *Mill) scheduleCleanupLocked(lease *RemoteLease) { if lease.cleanedUp || lease.cleanupRetry { return } lease.cleanupRetry = true time.AfterFunc(5*time.Second, func() { lease.finishMu.Lock() lease.cleanupRetry = false err := m.cleanupLeaseLocked(lease) lease.finishMu.Unlock() if err != nil { m.l.Error("retry lease cleanup failed", "lease", lease.id, "err", err) } })}
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) { if err := sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{ LeaseId: lease.id, Reason: reason, }}); err != nil { return } time.AfterFunc(m.cfg.CancelTimeout, func() { m.checkCancelDeadline(lease) })}
func (m *Mill) sessionForNode(nodeID string) *millSession { m.mu.Lock() defer m.mu.Unlock() sess := m.sessions[nodeID] if sess == nil || sess.disconnected { return nil } return sess}
func (m *Mill) onSnapshot(sess *millSession, snap *millv1.NodeSnapshot) error { m.mu.Lock() if sess.snapshot != nil && snap.Seqno <= sess.snapshot.Seqno { m.mu.Unlock() return protoErrf("snapshot seqno regression. Got %d, last seen %d", snap.Seqno, sess.snapshot.Seqno) } sess.snapshot = snap m.mu.Unlock()
if err := m.reconcileLeases(sess, snap.GetActiveLeaseIds()); err != nil { return err } m.mu.Lock() for _, id := range snap.GetActiveLeaseIds() { if lease := m.leases[id]; lease != nil && lease.nodeID == sess.nodeID && lease.epoch == sess.epoch { lease.claimed = true } } m.mu.Unlock() m.notifyChange() return nil}
func (m *Mill) onEventBatch(sess *millSession, batch *millv1.EventBatch) error { if batch == nil { return nil }
if batch.Epoch != sess.epoch { return protoErrf("batch epoch %q does not match session %q", batch.Epoch, sess.epoch) }
m.mu.Lock() currentKey := sess.nodeID + "/" + sess.epoch current := m.nodeSeqno[currentKey] m.mu.Unlock()
expected := current + 1 var newEntries []*millv1.Event for _, entry := range batch.Events { // reconnect replays old seqnos, drop those if entry.Seqno <= current { continue } // gap means executor lost rows. so dont apply a partial batch if entry.Seqno != expected { return protoErrf("gap in stream seqnos. Expected %d, got %d", expected, entry.Seqno) } newEntries = append(newEntries, entry) expected++ }
// all replays, still ack so the executor can trim its outbox if len(newEntries) == 0 { return m.sendAck(sess, current) }
type pendingTerminal struct { lease *RemoteLease ar *millv1.AttemptResult } var pendingTerminals []pendingTerminal var artifactLeases []*RemoteLease finishedInBatch := make(map[string]struct{})
var highestSeqno uint64 = current
applyFunc := func(tx *db.EventBatchTx) error { for _, entry := range newEntries { m.mu.Lock() lease := m.leases[entry.LeaseId] m.mu.Unlock()
// events for leases this node doesn't own are skipped but still // count as processed if lease == nil || lease.nodeID != sess.nodeID { highestSeqno = entry.Seqno continue }
// events arriving for a different epoch are invalid (different session) if lease.epoch != "" && lease.epoch != sess.epoch { return protoErrf("lease %q epoch %q does not match session %q", lease.id, lease.epoch, sess.epoch) }
lease.mu.Lock() state := lease.state lease.mu.Unlock()
// done leases can replay terminals on reconnect, skip them if state == leaseDone { highestSeqno = entry.Seqno continue }
if _, finished := finishedInBatch[lease.id]; finished { return protoErrf("stream entry follows terminal for lease %q", lease.id) }
switch { case entry.GetStatusEvent() != nil: ev := entry.GetStatusEvent() statusStr := string(models.StatusKindRunning) var errMsg *string if e := ev.GetError(); e != "" { errMsg = &e } var exitCode *int64 if c := ev.GetExitCode(); c != 0 { exitCode = &c } pipelineAtUri := string(lease.wid.PipelineId.AtUri()) if tx != nil { if err := tx.InsertStatusEvent(pipelineAtUri, lease.wid.Name, statusStr, errMsg, exitCode); err != nil { return err } }
case entry.GetAttemptResult() != nil: ar := entry.GetAttemptResult() statusStr := "success" switch ar.Status { case millv1.TerminalStatus_SUCCESS: statusStr = "success" case millv1.TerminalStatus_FAILED: statusStr = "failed" case millv1.TerminalStatus_TIMEOUT: statusStr = "timeout" case millv1.TerminalStatus_CANCELLED: statusStr = "cancelled" default: return protoErrf("unsupported terminal status %v", ar.Status) } var errMsg *string if e := ar.GetError(); e != "" { errMsg = &e } var exitCode *int64 if c := ar.GetExitCode(); c != 0 { exitCode = &c } finishedInBatch[lease.id] = struct{}{} pipelineAtUri := string(lease.wid.PipelineId.AtUri()) if tx != nil { if err := tx.InsertStatusEvent(pipelineAtUri, lease.wid.Name, statusStr, errMsg, exitCode); err != nil { return err } if err := tx.DeleteLease(lease.id); err != nil { return err } if a := ar.GetLogArtifact(); a != nil { if a.GetRef() == "" { return protoErrf("empty log artifact ref") } if !strings.HasPrefix(a.GetHash(), "sha256:") { return protoErrf("invalid log artifact hash %q", a.GetHash()) } if tx != nil { if err := tx.InsertArtifactRef(lease.id, lease.wid.Name, a.GetRef(), a.GetHash()); err != nil { return err } } artifactLeases = append(artifactLeases, lease) } } pendingTerminals = append(pendingTerminals, pendingTerminal{ lease: lease, ar: ar, }) }
highestSeqno = entry.Seqno }
if tx != nil { return tx.AdvanceCursor(sess.nodeID, sess.epoch, highestSeqno) } return nil }
var err error if m.db != nil { err = m.db.ApplyEventBatch(m.n, applyFunc) } else { err = applyFunc(nil) }
if err != nil { return err }
if m.cfg.LogDir != "" { for _, lease := range artifactLeases { path := models.LogFilePath(m.cfg.LogDir, lease.wid) if err := os.Remove(path); err != nil && !errors.Is(err, os.ErrNotExist) { m.l.Warn("failed to remove live log file after artifact recorded", "path", path, "err", err) } } }
m.mu.Lock() if highestSeqno > m.nodeSeqno[currentKey] { m.nodeSeqno[currentKey] = highestSeqno } m.mu.Unlock()
for _, pt := range pendingTerminals { // orphans have no waiting RunStep, just mark and clean up if pt.lease.orphaned { pt.lease.markDone() _ = m.cleanupLease(pt.lease) continue } pt.lease.deliverTerminal(pt.ar) if pt.lease.cleanupReady() { _ = m.cleanupLease(pt.lease) } }
return m.sendAck(sess, highestSeqno)}
func (m *Mill) sendAck(sess *millSession, seqno uint64) error { msg := &millproto.Message{Ack: &millv1.Ack{ Epoch: sess.epoch, UpToSeqno: seqno, }} if err := sess.send(msg); err != nil { return fmt.Errorf("send ack message: %w", err) } return nil}func (m *Mill) onLiveLog(sess *millSession, ll *millv1.LiveLog) error { if ll == nil || ll.GetLeaseId() == "" { return nil } m.mu.Lock() lease := m.leases[ll.GetLeaseId()] m.mu.Unlock() if lease == nil || lease.nodeID != sess.nodeID { return nil } raw := ll.GetRawJson() if m.cfg.LogDir == "" || len(raw) == 0 { if m.n != nil { m.n.NotifyAll() } return nil } lease.mu.Lock() isDone := (lease.state == leaseDone) lease.mu.Unlock() if isDone { if m.n != nil { m.n.NotifyAll() } return nil }
logPath := models.LogFilePath(m.cfg.LogDir, lease.wid) if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil { m.l.Warn("failed to create log dir", "path", filepath.Dir(logPath), "err", err) if m.n != nil { m.n.NotifyAll() } return nil } f, err := os.OpenFile(logPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600) if err != nil { m.l.Warn("failed to open log file", "path", logPath, "err", err) if m.n != nil { m.n.NotifyAll() } return nil } if _, err := f.Write(raw); err != nil { m.l.Warn("failed to write log file", "path", logPath, "err", err) } _ = f.Close()
if m.n != nil { m.n.NotifyAll() } return nil}
func (m *Mill) onCancelAck(sess *millSession, ca *millv1.CancelAck) { m.mu.Lock() lease := m.leases[ca.GetLeaseId()] m.mu.Unlock() if lease == nil || lease.nodeID != sess.nodeID { return } lease.mu.Lock() lease.cancelAcked = true lease.mu.Unlock()}
func (m *Mill) checkCancelDeadline(lease *RemoteLease) { lease.mu.Lock() isDone := (lease.state == leaseDone) isAcked := lease.cancelAcked lease.mu.Unlock()
if isDone && isAcked { return }
m.l.Warn("node failed to comply with cancel request within deadline, quarantining", "node", lease.nodeID, "lease", lease.id, "done", isDone, "acked", isAcked)
reason := fmt.Sprintf("cancel noncompliance for lease %s (done: %t, acked: %t)", lease.id, isDone, isAcked) if m.db != nil { if err := m.db.QuarantineExecutor(lease.nodeID, reason); err != nil { m.l.Error("failed to quarantine executor", "node", lease.nodeID, "err", err) } }
m.mu.Lock() sess := m.sessions[lease.nodeID] m.mu.Unlock() if sess != nil { sess.close() }}
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}