Monorepo for Tangled
Something went wrong. Try again.
Go
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324package 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/spindle/storage" "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 CacheStoreID string CacheMaxBytesPerOwner int64 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 cache storage.Storage
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) AttachCache(store storage.Storage) { m.mu.Lock() m.cache = store 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) { m.sendCancel(sess, lease) } 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}
// 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 wf.CacheNamespace = lease.cacheNamespace 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, len(wf.Caches) > 0) 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) lease.cacheNamespace = sess.cacheNamespace 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), RepoDid: wf.RepoDid, }} 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, cacheRequired bool) []*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 } if cacheRequired && (m.cfg.CacheStoreID == "" || s.cacheStoreID != m.cfg.CacheStoreID) { 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
capabilities, err := cacheCapabilities(wf.RepoDid, wf.CacheBindings) if err != nil { return err } if m.db != nil { if err := m.db.SaveMillCacheCapabilities(lease.id, capabilities); err != nil { return fmt.Errorf("persist mill cache capabilities: %w", err) } }
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, CacheBindings: cacheBindingsToProto(wf.CacheBindings), }}
// 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, _ := lease.cancelRequested(); cancelled { m.sendCancel(sess, lease) } 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) } }}
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) { if err := sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{ LeaseId: lease.id, }}); 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 cacheDeletes := make(map[string]db.CacheEntry) cacheObjectDeletes := make(map[string]struct{}) 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.GetCacheUpdate() != nil: update := entry.GetCacheUpdate() if tx != nil { at := time.Now() capabilityAction := "" switch update.GetAction() { case millv1.CacheUpdateAction_CACHE_USED, millv1.CacheUpdateAction_CACHE_MISSING: capabilityAction = "restore" case millv1.CacheUpdateAction_CACHE_STORED, millv1.CacheUpdateAction_CACHE_DISCARDED: capabilityAction = "save" default: return protoErrf("unsupported cache update action %v", update.GetAction()) } ref, authorized, err := tx.ConsumeMillCacheCapability( context.Background(), lease.id, capabilityAction, update.GetId(), ) if err != nil { return err } if !authorized { return protoErrf( "cache update %q was not planned for lease %q", update.GetId(), lease.id, ) }
switch update.GetAction() { case millv1.CacheUpdateAction_CACHE_USED: if err := tx.TouchCacheEntry(context.Background(), update.GetId(), at); err != nil { return err } case millv1.CacheUpdateAction_CACHE_STORED: superseded, ready, err := tx.MarkCacheEntryReady( context.Background(), update.GetId(), update.GetSizeBytes(), m.cfg.CacheMaxBytesPerOwner, at, ) if err != nil { return err } for _, old := range superseded { cacheDeletes[old.ID] = old } if !ready { if err := tx.QueueCacheObjectDeletion(context.Background(), ref, at); err != nil { return err } cacheObjectDeletes[ref] = struct{}{} } case millv1.CacheUpdateAction_CACHE_DISCARDED, millv1.CacheUpdateAction_CACHE_MISSING: pendingOnly := update.GetAction() == millv1.CacheUpdateAction_CACHE_DISCARDED discarded, err := tx.DiscardCacheEntry( context.Background(), update.GetId(), pendingOnly, ) if err != nil { return err } if discarded != nil { cacheDeletes[discarded.ID] = *discarded } } }
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 }
m.deleteCacheEntries(cacheDeletes)
m.deleteQueuedCacheObjects(cacheObjectDeletes)
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) deleteCacheEntries(entries map[string]db.CacheEntry) { if len(entries) == 0 { return } if m.cache == nil { m.l.Warn("cache objects need deletion but the mill has no cache store", "count", len(entries)) return } for _, entry := range entries { if err := m.cache.Delete(context.Background(), entry.StorageKey); err != nil { m.l.Warn("delete cache object", "id", entry.ID, "ref", entry.StorageKey, "err", err) continue } if m.db != nil { if err := m.db.DeleteCacheEntry(context.Background(), entry.ID); err != nil { m.l.Warn("delete cache metadata", "id", entry.ID, "err", err) } } }}
func (m *Mill) deleteQueuedCacheObjects(keys map[string]struct{}) { if len(keys) == 0 { return } if m.cache == nil { m.l.Warn("cache objects need deletion but the mill has no cache store", "count", len(keys)) return } for key := range keys { if err := m.cache.Delete(context.Background(), key); err != nil { m.l.Warn("delete unindexed cache object", "ref", key, "err", err) continue } if m.db != nil { if err := m.db.CompleteCacheObjectDeletion(context.Background(), key); err != nil { m.l.Warn("complete cache object deletion", "ref", key, "err", err) } } }}
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}