Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
Go
12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336package mill
import ( "context" "encoding/json" "errors" "fmt" "io" "log/slog" "maps" "os" "path/filepath" "slices" "sort" "strings" "sync" "time"
"go.opentelemetry.io/otel/attribute" "go.opentelemetry.io/otel/codes" "tangled.org/core/notifier" "tangled.org/core/spindle/db" "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" "tangled.org/core/spindle/observability" "tangled.org/core/spindle/quota" "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 defaultCancelAckTimeout = 10 * time.Second defaultCancelTeardownTimeout = 10 * time.Minute defaultCacheDeleteTimeout = 30 * time.Second)
var 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 CancelAckTimeout time.Duration // teardown can take minutes after the cancel is acknowledged CancelTeardownTimeout time.Duration // bounds time spent deleting an object from the shared cache CacheDeleteTimeout time.Duration // how many protocol-violating session deaths in a row quarantine the node QuarantineStrikes int}
type cacheDeletion struct { entries []db.CacheEntry objects []string}
type Mill struct { l *slog.Logger cfg Config
db *db.DB n *notifier.Notifier metrics *observability.Metrics qm *quota.Manager cache storage.Storage mu sync.Mutex quotaLifecycleMu sync.Mutex sessions map[string]*millSession leases map[string]*RemoteLease reservations map[string]*RemoteLease nodeSeqno map[string]uint64 quotaLeases map[string]*RemoteLease protoStrikes map[string]int pending int changeCh chan struct{}
cacheDeleteOnce sync.Once cacheDeleteMu sync.Mutex cacheDeleteQueue []cacheDeletion cacheDeleteWake chan struct{}
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.CancelAckTimeout <= 0 { cfg.CancelAckTimeout = defaultCancelAckTimeout } if cfg.CancelTeardownTimeout <= 0 { cfg.CancelTeardownTimeout = defaultCancelTeardownTimeout } if cfg.CacheDeleteTimeout <= 0 { cfg.CacheDeleteTimeout = defaultCacheDeleteTimeout } m := &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), changeCh: make(chan struct{}), quotaLeases: make(map[string]*RemoteLease), } m.startCacheDeletionWorker() return m}
func (m *Mill) Attach(d *db.DB, n *notifier.Notifier, qm *quota.Manager) { m.mu.Lock() m.db = d m.n = n m.qm = qm m.mu.Unlock()}
func (m *Mill) AttachCache(store storage.Storage) { m.mu.Lock() m.cache = store m.mu.Unlock() m.ensureCacheSentinel(store)}
func (m *Mill) ensureCacheSentinel(store storage.Storage) { if store == nil { return } rc, err := store.Get(context.Background(), millproto.CacheSentinelKey) if err == nil { defer rc.Close() value, readErr := io.ReadAll(io.LimitReader(rc, int64(len(millproto.CacheSentinelValue)+1))) if readErr != nil || string(value) != millproto.CacheSentinelValue { m.l.Warn("cache store sentinel has unexpected contents", "key", millproto.CacheSentinelKey) } return } if !errors.Is(err, storage.ErrNotExist) { m.l.Warn("failed to read cache store sentinel", "key", millproto.CacheSentinelKey, "err", err) return } if err := store.Put(context.Background(), millproto.CacheSentinelKey, strings.NewReader(millproto.CacheSentinelValue)); err != nil { m.l.Warn("failed to initialize cache store sentinel", "key", millproto.CacheSentinelKey, "err", err) }}
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) { // load db state before taking the mill lock 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) { return 0, false } if old.recoveryExpired { m.l.Warn("rejecting executor reconnect while lease failure retries", "node", sess.nodeID) return 0, false } if old.graceTimer != nil { old.graceTimer.Stop() } old.close() if old.recovering { m.armSessionRecoveryLocked(sess, old.graceDeadline) } else { m.armSessionRecoveryLocked(sess, time.Now().Add(m.cfg.ReconnectGrace)) } if old.disconnected { m.metrics.RecordMillReconnect() 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) armSessionRecoveryLocked(sess *millSession, deadline time.Time) { sess.recovering = true sess.recoveryReady = false sess.graceDeadline = deadline delay := time.Until(deadline) if delay < 0 { delay = 0 } sess.graceTimer = time.AfterFunc(delay, func() { m.failLeasesAfterGrace(sess) })}
func (m *Mill) maybeFinishSessionRecoveryLocked(sess *millSession) { if m.sessions[sess.nodeID] != sess || sess.disconnected || !sess.recovering || sess.recoveryExpired || !sess.recoveryReady || len(sess.pendingCancels) != 0 { return } if sess.graceTimer != nil { sess.graceTimer.Stop() } sess.recovering = false sess.graceDeadline = time.Time{} sess.graceTimer = nil}
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 if !sess.recovering { m.armSessionRecoveryLocked(sess, time.Now().Add(m.cfg.ReconnectGrace)) } 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) noteSessionError(sess *millSession, err error) { if errors.Is(err, errProtocolViolation) { m.l.Warn("executor session rejected on protocol violation", "node", sess.nodeID, "err", err) }}
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 lease.cancelPending() { leases = append(leases, lease) } } return leases}
// reconnect recovery did not complete before grace expired. fails every// lease the node held, and reschedules itself if any failure didn't stickfunc (m *Mill) failLeasesAfterGrace(sess *millSession) { m.failLeasesAfterGraceAttempt(sess, false)}
func (m *Mill) failLeasesAfterGraceAttempt(sess *millSession, retry bool) { m.mu.Lock() if m.sessions[sess.nodeID] != sess || !sess.recovering || (sess.recoveryExpired && !retry) { m.mu.Unlock() return } sess.recoveryExpired = true 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.failLeasesAfterGraceAttempt(sess, true) }) m.mu.Unlock() return }
// everything failed cleanly. forget the node and wake placement current := m.sessions[sess.nodeID] == sess if current { delete(m.sessions, sess.nodeID) } m.mu.Unlock() if current { sess.close() } m.notifyChange() m.sweepUnclaimedOrphans()}
func (m *Mill) place(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow) (engine.WorkflowSlot, error) { placementStart := time.Now() placementResult := "error" defer func() { m.metrics.RecordMillPlacementResult(ctx, engineName, placementResult, time.Since(placementStart)) }()
m.mu.Lock() if m.cfg.MaxPending > 0 && m.pending >= m.cfg.MaxPending { max := m.cfg.MaxPending cur := m.pending m.mu.Unlock() placementResult = "rejected" m.metrics.RecordMillPlacementAdmission(false, "max_pending_reached") return nil, fmt.Errorf("%w: mill has %d pending jobs (max %d)", engine.ErrNoWorkflowSlots, cur, max) } m.metrics.RecordMillPlacementAdmission(true, "allowed") m.pending++ m.mu.Unlock() defer func() { m.mu.Lock() m.pending-- m.mu.Unlock() }()
for { if err := ctx.Err(); err != nil { placementResult = "cancelled" if errors.Is(err, context.DeadlineExceeded) { placementResult = "timeout" } return nil, err }
// capture the channel before bidding // a change during the bid makes the next wait return immediately ch := m.currentChangeCh()
lease, err := m.bid(ctx, engineName, wid, wf) if err != nil { return nil, err } if lease != nil { wf.CacheNamespace = lease.cacheNamespace slot, retry, err := m.admit(ctx, engineName, wid, wf, lease) if err != nil { var failure *engine.WorkflowFailure if errors.As(err, &failure) && failure.Class == engine.FailureClassPolicy { placementResult = "rejected" } return nil, err } if slot != nil { placementResult = "success" return slot, nil } if retry { continue } }
// no executor available. wait for a change or ctx select { case <-ctx.Done(): placementResult = "cancelled" if errors.Is(ctx.Err(), context.DeadlineExceeded) { placementResult = "timeout" } return nil, ctx.Err() case <-ch: } }}
func (m *Mill) admit(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow, lease *RemoteLease) (engine.WorkflowSlot, bool, error) { req, ok := m.quotaRequest(wid, wf, lease.resources) if !ok { return m.publish(wid, wf, lease, nil) }
qlease, res, err := m.qm.TryAcquire(ctx, req) switch { case err != nil: m.releaseSeat(lease) return nil, false, fmt.Errorf("acquire workflow quota: %w", err)
case res.Allowed: return m.publish(wid, wf, lease, qlease)
case !res.Temporary: m.releaseSeat(lease) return nil, false, engine.ClassifiedFailure( engine.FailureClassPolicy, engine.FailureReasonQuotaDenied, fmt.Errorf("%w: workflow quota denied: %s", engine.ErrWorkflowFailed, res.Reason), ) }
// don't hold fleet capacity while waiting for quota m.releaseSeat(lease)
qlease, err = m.qm.Acquire(ctx, req) if err != nil { if errors.Is(err, quota.ErrDenied) { err = engine.ClassifiedFailure(engine.FailureClassPolicy, engine.FailureReasonQuotaDenied, err) } return nil, false, fmt.Errorf("wait for workflow quota: %w", err) }
// a re-bid must match the reserved resources next, err := m.bid(ctx, engineName, wid, wf) if err != nil { _ = m.releaseAbandonedQuota(qlease) return nil, false, err } if next == nil { _ = m.releaseAbandonedQuota(qlease) return nil, false, nil } if !maps.Equal(next.resources, req.Resources) { m.l.Warn("re-bid reported different resources than the quota reserved, retrying placement", "lease", next.id, "node", next.nodeID, "engine", engineName) m.releaseSeat(next) _ = m.releaseAbandonedQuota(qlease) return nil, true, nil } return m.publish(wid, wf, next, qlease)}
func (m *Mill) releaseAbandonedQuota(qlease quota.Lease) error { if qlease == nil { return nil } err := qlease.ReleaseWithError() if err != nil { dummy := &RemoteLease{ id: "abandoned-" + qlease.ID(), quotaID: qlease.ID(), state: leaseDone, cleanedUp: false, } dummy.finishMu.Lock() m.scheduleCleanupLocked(dummy) dummy.finishMu.Unlock() m.l.Error("abandoned quota release failed, retry scheduled", "quota_id", qlease.ID(), "err", err) } return err}
// teardown releases the attached quota exactly oncefunc (m *Mill) publish(wid models.WorkflowId, wf *models.Workflow, lease *RemoteLease, qlease quota.Lease) (engine.WorkflowSlot, bool, error) { lease.wid = wid if id, ok := m.quotaIdentity(wf); ok { lease.ownerDID = id.OwnerDID lease.repoDID = id.RepoDID } if qlease != nil { lease.quotaLease = qlease lease.quotaID = qlease.ID() } if err := m.persistLease(lease, leaseRowReserved); err != nil { m.releaseSeat(lease) // release an unrecorded reservation releaseErr := m.releaseQuotaLocked(lease) if releaseErr != nil { lease.finishMu.Lock() m.scheduleCleanupLocked(lease) lease.finishMu.Unlock() } return nil, false, errors.Join(fmt.Errorf("persist reserved mill lease: %w", err), releaseErr) } 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}, false, nil}
func (m *Mill) releaseSeat(lease *RemoteLease) { m.dropReservation(lease.id) m.releaseRemote(lease)}
func (m *Mill) quotaRequest(wid models.WorkflowId, wf *models.Workflow, resources quota.Resources) (quota.ReserveRequest, bool) { if m.qm == nil || len(resources) == 0 { return quota.ReserveRequest{}, false } id, ok := m.quotaIdentity(wf) if !ok { m.l.Warn("placing workflow without quota because the pipeline has no repo identity", "workflow", wid.String()) return quota.ReserveRequest{}, false } resID := quota.WorkflowReservationID(wf.RunID, id.OwnerDID, id.RepoDID, wid.Knot, wid.Rkey, wid.Name) return quota.ReserveRequest{ ID: resID, Kind: quota.KindWorkflow, Key: resID, Identity: id, Resources: resources, }, true}
func (m *Mill) quotaIdentity(wf *models.Workflow) (quota.Identity, bool) { st, ok := wf.Data.(*millWorkflowState) if !ok || st == nil || st.RawPipeline.TriggerMetadata == nil { return quota.Identity{}, false } repo := st.RawPipeline.TriggerMetadata.Repo if repo == nil { return quota.Identity{}, false } id := quota.Identity{OwnerDID: repo.Did, RepoDID: repo.Did} if repo.RepoDid != nil && *repo.RepoDid != "" { id.RepoDID = *repo.RepoDid } if !strings.HasPrefix(id.OwnerDID, "did:") || !strings.HasPrefix(id.RepoDID, "did:") { return quota.Identity{}, false } return id, true}
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 failureClass string failureReason 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() bidCtx, span := observability.Tracer().Start(bidCtx, "mill.placement.bid") if span.IsRecording() { var ownerDID, repoDID string if wf != nil { ownerDID = wf.OwnerDID repoDID = wf.RepoDID } if (ownerDID == "" || repoDID == "") && wf != nil { if id, ok := m.quotaIdentity(wf); ok { ownerDID = id.OwnerDID repoDID = id.RepoDID } } attrs := []attribute.KeyValue{ attribute.String(observability.WorkflowIDKey, wid.String()), attribute.String(observability.LeaseIDKey, leaseID), attribute.String(observability.ExecutorNodeIDKey, sess.nodeID), attribute.String(observability.PipelineIDKey, wid.PipelineId.AtUri().String()), } if ownerDID != "" { attrs = append(attrs, attribute.String(observability.OwnerDIDKey, ownerDID)) } if repoDID != "" { attrs = append(attrs, attribute.String(observability.RepoDIDKey, repoDID)) } span.SetAttributes(attrs...) } accepted := false defer func() { if accepted { span.SetStatus(codes.Ok, "accepted") } else { span.SetStatus(codes.Error, "placement bid failed") } span.End() }()
traceparent, tracestate := observability.InjectToTraceparentAndTracestate(bidCtx)
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), Traceparent: traceparent, Tracestate: tracestate, 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(), failureClass: rr.GetFailureClass(), failureReason: rr.GetFailureReason(), } return } results <- bidResult{} return } // invalid resource claims are protocol errors if err := quota.ValidateResources(rr.GetQuotaResources()); err != nil { m.l.WarnContext(bidCtx, "executor reported invalid workflow resources", "node", sess.nodeID, "lease", leaseID, "err", err) m.dropReservation(leaseID) m.releaseRemote(lease) results <- bidResult{} return } lease.resources = maps.Clone(rr.GetQuotaResources()) lease.millRecordsTerminalMetrics = rr.GetSupportsMillTerminalMetrics() accepted = true 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 var incompatibleClass, incompatibleReason string attributionConsistent := true attributionSet := false // 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) } if r.failureClass == "" || r.failureReason == "" { attributionConsistent = false } else if !attributionSet { incompatibleClass = r.failureClass incompatibleReason = r.failureReason attributionSet = true } else if incompatibleClass != r.failureClass || incompatibleReason != r.failureReason { attributionConsistent = false } } 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 { err := fmt.Errorf("no compatible executor for %s: %s", engineName, strings.Join(incompatible, "; ")) if attributionSet && attributionConsistent { return nil, engine.ClassifiedFailure( engine.FailureClass(incompatibleClass), engine.FailureReason(incompatibleReason), err, ) } return nil, err } 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 || s.recovering || len(s.pendingCancels) != 0 { continue } if s.snapshot == nil { continue } if cacheRequired && (m.cfg.CacheStoreID == "" || s.cacheStoreID != m.cfg.CacheStoreID || s.legacyProtocol) { 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) (err error) { st, ok := wf.Data.(*millWorkflowState) if !ok || st == nil || st.Lease == nil { return fmt.Errorf("mill workflow state missing lease") } lease := st.Lease
ctx, span := observability.Tracer().Start(ctx, "mill.commit") if span.IsRecording() { attrs := []attribute.KeyValue{ attribute.String(observability.WorkflowIDKey, lease.wid.String()), attribute.String(observability.LeaseIDKey, lease.id), attribute.String(observability.ExecutorNodeIDKey, lease.nodeID), attribute.String(observability.PipelineIDKey, lease.wid.PipelineId.AtUri().String()), } if lease.ownerDID != "" { attrs = append(attrs, attribute.String(observability.OwnerDIDKey, lease.ownerDID)) } if lease.repoDID != "" { attrs = append(attrs, attribute.String(observability.RepoDIDKey, lease.repoDID)) } span.SetAttributes(attrs...) } defer func() { if err != nil { span.SetStatus(codes.Error, "commit failed") } else { span.SetStatus(codes.Ok, "success") } span.End() }()
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} }
traceparent, tracestate := observability.InjectToTraceparentAndTracestate(ctx)
commit := &millproto.Message{CommitLease: &millv1.CommitLease{ LeaseId: lease.id, Secrets: pbSecrets, Traceparent: traceparent, Tracestate: tracestate, MillRecordsTerminalMetrics: lease.millRecordsTerminalMetrics, 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, 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_TERMINAL_STATUS_SUCCESS: return nil case millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT: return engine.ErrTimedOut case millv1.TerminalStatus_TERMINAL_STATUS_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 { // fence unreported save uploads before dropping the lease capabilities; // a late object is then covered by a durable deletion tombstone pendingKeys, err := m.db.DiscardPendingCacheEntriesForLease(lease.id, time.Now()) if err != nil { m.scheduleCleanupLocked(lease) return err } if len(pendingKeys) > 0 { objects := make(map[string]struct{}, len(pendingKeys)) for _, key := range pendingKeys { objects[key] = struct{}{} } m.enqueueCacheDeletions(nil, objects) } // delete the row before releasing the reservation 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()
if err := m.releaseQuotaLocked(lease); err != nil { m.scheduleCleanupLocked(lease) return err } lease.cleanedUp = true m.notifyChange() return nil}
func (m *Mill) releaseQuotaLocked(lease *RemoteLease) error { if m.qm == nil { return nil } lease.mu.Lock() workflowID := lease.quotaID if workflowID == "" && lease.quotaLease != nil { workflowID = lease.quotaLease.ID() } remoteQuotas := maps.Clone(lease.remoteQuotas) lease.mu.Unlock()
ctx := context.Background() var releaseErr error if workflowID != "" { if err := m.qm.Release(ctx, workflowID); err != nil { releaseErr = errors.Join(releaseErr, fmt.Errorf("release workflow quota: %w", err)) } else { lease.mu.Lock() if lease.quotaID == workflowID { lease.quotaLease = nil lease.quotaID = "" } lease.mu.Unlock() } } for id, rq := range remoteQuotas { if rq.committing { if err := m.qm.Commit(ctx, id); err != nil { releaseErr = errors.Join(releaseErr, fmt.Errorf("commit remote quota %s: %w", id, err)) continue } } else { if err := m.qm.Release(ctx, id); err != nil { releaseErr = errors.Join(releaseErr, fmt.Errorf("release remote quota %s: %w", id, err)) continue } } lease.mu.Lock() delete(lease.remoteQuotas, id) lease.mu.Unlock()
m.quotaLifecycleMu.Lock() if m.quotaLeases[id] == lease { delete(m.quotaLeases, id) } m.quotaLifecycleMu.Unlock() } return releaseErr}
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) { lease.finishMu.Lock() m.mu.Lock() if m.sessions[lease.nodeID] != sess || sess.disconnected || m.leases[lease.id] != lease || lease.getState() == leaseDone { m.mu.Unlock() lease.finishMu.Unlock() return } sess.pendingCancels[lease.id] = struct{}{} m.mu.Unlock() lease.finishMu.Unlock()
if err := sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{ LeaseId: lease.id, Reason: reason, }}); err != nil { m.settleSessionCancel(sess, lease.id) sess.close() return } time.AfterFunc(m.cfg.CancelAckTimeout, func() { m.checkCancelAck(sess, lease) })}
func (m *Mill) sendUntrackedCancel(sess *millSession, leaseID, reason string) error { m.mu.Lock() if m.sessions[sess.nodeID] != sess || sess.disconnected { m.mu.Unlock() return errSessionClosed } sess.pendingCancels[leaseID] = struct{}{} m.mu.Unlock()
if err := sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{ LeaseId: leaseID, Reason: reason, }}); err != nil { m.settleSessionCancel(sess, leaseID) sess.close() return err } time.AfterFunc(m.cfg.CancelAckTimeout, func() { m.checkUntrackedCancelAck(sess, leaseID) }) return nil}
func (m *Mill) checkUntrackedCancelAck(sess *millSession, leaseID string) { m.mu.Lock() _, pending := sess.pendingCancels[leaseID] isCurrent := m.sessions[sess.nodeID] == sess && !sess.disconnected m.mu.Unlock() if !pending || !isCurrent { return } m.l.Warn("node did not acknowledge unknown lease cancel within deadline; closing session", "node", sess.nodeID, "lease", leaseID, "deadline", m.cfg.CancelAckTimeout) sess.close()}
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 } } sess.recoveryReady = true m.maybeFinishSessionRecoveryLocked(sess) 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 } cacheDeletes := make(map[string]db.CacheEntry) cacheObjectDeletes := make(map[string]struct{}) type startupObservation struct { engine string delay time.Duration } leaseSet := make(map[*RemoteLease]struct{}) for _, entry := range newEntries { m.mu.Lock() lease := m.leases[entry.LeaseId] m.mu.Unlock() if lease == nil || lease.nodeID != sess.nodeID { continue } if lease.epoch != "" && lease.epoch != sess.epoch { return protoErrf("lease %q epoch %q does not match session %q", lease.id, lease.epoch, sess.epoch) } leaseSet[lease] = struct{}{} } lockedLeases := make([]*RemoteLease, 0, len(leaseSet)) for lease := range leaseSet { lockedLeases = append(lockedLeases, lease) } sort.Slice(lockedLeases, func(i, j int) bool { return lockedLeases[i].id < lockedLeases[j].id }) for _, lease := range lockedLeases { lease.finishMu.Lock() } unlockLeases := func() { for i := len(lockedLeases) - 1; i >= 0; i-- { lockedLeases[i].finishMu.Unlock() } }
type batchApply struct { terminalCancelIDs []string pendingTerminals []pendingTerminal artifactLeases []*RemoteLease startupObservations []startupObservation highestSeqno uint64 }
applyFunc := func(tx *db.EventBatchTx) (batchApply, error) { var r batchApply r.highestSeqno = current finishedInBatch := make(map[string]struct{})
for _, entry := range newEntries { if entry.GetAttemptResult() != nil { r.terminalCancelIDs = append(r.terminalCancelIDs, entry.LeaseId) } 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 { r.highestSeqno = entry.Seqno continue }
// events arriving for a different epoch are invalid (different session) if lease.epoch != "" && lease.epoch != sess.epoch { return batchApply{}, 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 { r.highestSeqno = entry.Seqno continue }
if _, finished := finishedInBatch[lease.id]; finished { return batchApply{}, 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 { alreadyRunning, err := tx.HasWorkflowStatus(context.Background(), pipelineAtUri, lease.wid.Name, statusStr) if err != nil { return batchApply{}, err } if err := tx.InsertStatusEvent(pipelineAtUri, lease.wid.Name, statusStr, errMsg, exitCode); err != nil { return batchApply{}, err } if !alreadyRunning { delay, ok, err := tx.WorkflowStartupDelay(context.Background(), pipelineAtUri, lease.wid.Name) if err != nil { return batchApply{}, err } if ok { r.startupObservations = append(r.startupObservations, startupObservation{ engine: lease.engine, delay: delay, }) } } }
case entry.GetCacheUpdate() != nil: update := entry.GetCacheUpdate() if tx != nil { at := time.Now() capabilityAction := "" switch update.GetAction() { case millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_USED, millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_MISSING: capabilityAction = "restore" case millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED, millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_DISCARDED: capabilityAction = "save" default: return batchApply{}, protoErrf("unsupported cache update action %v", update.GetAction()) } ref, authorized, err := tx.ConsumeMillCacheCapability( context.Background(), lease.id, capabilityAction, update.GetId(), ) if err != nil { return batchApply{}, err } if !authorized { return batchApply{}, protoErrf( "cache update %q was not planned for lease %q", update.GetId(), lease.id, ) }
switch update.GetAction() { case millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_USED: if err := tx.TouchCacheEntry(context.Background(), update.GetId(), at); err != nil { return batchApply{}, err } case millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED: actualSize, err := m.cacheObjectSize(context.Background(), ref) if err != nil { return batchApply{}, fmt.Errorf("stat stored cache object %q: %w", ref, err) } superseded, ready, err := tx.MarkCacheEntryReadyWithChecksum( context.Background(), update.GetId(), actualSize, update.GetChecksum(), m.cfg.CacheMaxBytesPerOwner, at, ) if err != nil { return batchApply{}, err } for _, old := range superseded { cacheDeletes[old.ID] = old } if !ready { if err := tx.QueueCacheObjectDeletion(context.Background(), ref, at); err != nil { return batchApply{}, err } cacheObjectDeletes[ref] = struct{}{} } case millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_DISCARDED, millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_MISSING: // A missing restore object must not let one executor delete a // ready generation that another executor published. Both // actions can only discard entries still owned by this save. pendingOnly := true discarded, err := tx.DiscardCacheEntry( context.Background(), update.GetId(), pendingOnly, ) if err != nil { return batchApply{}, err } if discarded != nil { cacheDeletes[discarded.ID] = *discarded } } }
case entry.GetAttemptResult() != nil: ar := entry.GetAttemptResult() statusStr := "success" switch ar.Status { case millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS: statusStr = "success" case millv1.TerminalStatus_TERMINAL_STATUS_FAILED: statusStr = "failed" case millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT: statusStr = "timeout" case millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED: statusStr = "cancelled" default: return batchApply{}, 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 batchApply{}, err } if err := tx.DeleteLease(lease.id); err != nil { return batchApply{}, err } if a := ar.GetLogArtifact(); a != nil { if a.GetRef() == "" { return batchApply{}, protoErrf("empty log artifact ref") } // the ref names an object the mill deletes on wipe; // pin it to the derived log key so an executor cannot // point cleanup at someone else's objects if a.GetRef() != "logs/"+lease.id+".log" { return batchApply{}, protoErrf("log artifact ref %q does not match lease %s", a.GetRef(), lease.id) } if !strings.HasPrefix(a.GetHash(), "sha256:") { return batchApply{}, protoErrf("invalid log artifact hash %q", a.GetHash()) } if lease.wid.Knot == "" || lease.wid.Rkey == "" || lease.wid.Name == "" { m.l.Error("skipping log artifact with incomplete workflow identity", "lease", lease.id, "wid", lease.wid) } else { if tx != nil { if err := tx.InsertArtifactRef(lease.id, lease.repoDID, lease.wid, a.GetRef(), a.GetHash()); err != nil { return batchApply{}, err } } r.artifactLeases = append(r.artifactLeases, lease) } } } r.pendingTerminals = append(r.pendingTerminals, pendingTerminal{ lease: lease, ar: ar, }) }
r.highestSeqno = entry.Seqno }
if tx != nil { if err := tx.AdvanceCursor(sess.nodeID, sess.epoch, r.highestSeqno); err != nil { return batchApply{}, err } } return r, nil }
var applied batchApply var err error if m.db != nil { applied, err = db.ApplyEventBatch(m.db, m.n, applyFunc) } else { applied, err = applyFunc(nil) }
if err != nil { unlockLeases() return err }
for _, observation := range applied.startupObservations { m.metrics.RecordWorkflowStartupDelay(context.Background(), observation.engine, "mill", observation.delay) }
m.enqueueCacheDeletions(cacheDeletes, cacheObjectDeletes)
if m.cfg.LogDir != "" { for _, lease := range applied.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 applied.highestSeqno > m.nodeSeqno[currentKey] { m.nodeSeqno[currentKey] = applied.highestSeqno } m.mu.Unlock()
for _, leaseID := range applied.terminalCancelIDs { m.settleSessionCancel(sess, leaseID) } for _, pt := range applied.pendingTerminals { if pt.lease.millRecordsTerminalMetrics && pt.ar.GetMillRecordsTerminalMetrics() && pt.ar.GetFailureClass() != "" && pt.ar.GetFailureReason() != "" { m.metrics.RecordWorkflowTerminal( pt.lease.engine, terminalMetricResult(pt.ar.GetStatus()), pt.ar.GetFailureClass(), pt.ar.GetFailureReason(), ) } // orphans have no waiting RunStep, just mark and clean up if pt.lease.orphaned { pt.lease.markDone() _ = m.cleanupLeaseLocked(pt.lease) continue } pt.lease.deliverTerminal(pt.ar) if pt.lease.cleanupReady() { _ = m.cleanupLeaseLocked(pt.lease) } } unlockLeases()
return m.sendAck(sess, applied.highestSeqno)}func terminalMetricResult(status millv1.TerminalStatus) string { switch status { case millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS: return "success" case millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT: return "timeout" case millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED: return "cancelled" default: return "failure" }}
func (m *Mill) cacheObjectSize(ctx context.Context, key string) (int64, error) { m.mu.Lock() store := m.cache m.mu.Unlock() if store == nil { return 0, errors.New("cache store is not configured") } if statter, ok := store.(storage.StatStorage); ok { return statter.Stat(ctx, key) } r, err := store.Get(ctx, key) if err != nil { return 0, err } defer r.Close() n, err := io.Copy(io.Discard, r) if err != nil { return 0, err } return n, nil}
func (m *Mill) startCacheDeletionWorker() { m.cacheDeleteOnce.Do(func() { m.cacheDeleteWake = make(chan struct{}, 1) go m.cacheDeletionLoop() })}
func (m *Mill) enqueueCacheDeletions(entries map[string]db.CacheEntry, objects map[string]struct{}) { if len(entries) == 0 && len(objects) == 0 { return } m.startCacheDeletionWorker() deletion := cacheDeletion{ entries: make([]db.CacheEntry, 0, len(entries)), objects: make([]string, 0, len(objects)), } for _, entry := range entries { deletion.entries = append(deletion.entries, entry) } for key := range objects { deletion.objects = append(deletion.objects, key) } m.cacheDeleteMu.Lock() m.cacheDeleteQueue = append(m.cacheDeleteQueue, deletion) m.cacheDeleteMu.Unlock() select { case m.cacheDeleteWake <- struct{}{}: default: }}
func (m *Mill) cacheDeletionLoop() { for range m.cacheDeleteWake { for { m.cacheDeleteMu.Lock() if len(m.cacheDeleteQueue) == 0 { m.cacheDeleteMu.Unlock() break } deletion := m.cacheDeleteQueue[0] m.cacheDeleteQueue[0] = cacheDeletion{} m.cacheDeleteQueue = m.cacheDeleteQueue[1:] m.cacheDeleteMu.Unlock() m.deleteCacheEntries(deletion.entries) m.deleteQueuedCacheObjects(deletion.objects) } }}
func (m *Mill) cacheDeleteContext() (context.Context, context.CancelFunc) { timeout := m.cfg.CacheDeleteTimeout if timeout <= 0 { timeout = defaultCacheDeleteTimeout } return context.WithTimeout(context.Background(), timeout)}
func (m *Mill) deleteCacheEntries(entries []db.CacheEntry) { if len(entries) == 0 { return } m.mu.Lock() cache := m.cache database := m.db m.mu.Unlock() if cache == nil { m.l.Warn("cache objects need deletion but the mill has no cache store", "count", len(entries)) return } for _, entry := range entries { ctx, cancel := m.cacheDeleteContext() err := cache.Delete(ctx, entry.StorageKey) cancel() if err != nil { m.l.Warn("delete cache object", "id", entry.ID, "ref", entry.StorageKey, "err", err) continue } if database != nil { ctx, cancel = m.cacheDeleteContext() err = database.DeleteCacheEntry(ctx, entry.ID) cancel() if err != nil { m.l.Warn("delete cache metadata", "id", entry.ID, "err", err) } } }}
func (m *Mill) deleteQueuedCacheObjects(keys []string) { if len(keys) == 0 { return } m.mu.Lock() cache := m.cache database := m.db m.mu.Unlock() if cache == nil { m.l.Warn("cache objects need deletion but the mill has no cache store", "count", len(keys)) return } for _, key := range keys { ctx, cancel := m.cacheDeleteContext() err := cache.Delete(ctx, key) cancel() if err != nil { m.l.Warn("delete unindexed cache object", "ref", key, "err", err) continue } if database != nil { ctx, cancel = m.cacheDeleteContext() err = database.CompleteCacheObjectDeletion(ctx, key) cancel() if 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) settleSessionCancel(sess *millSession, leaseID string) { m.mu.Lock() _, pending := sess.pendingCancels[leaseID] delete(sess.pendingCancels, leaseID) m.maybeFinishSessionRecoveryLocked(sess) m.mu.Unlock() if pending { m.notifyChange() }}
func (m *Mill) onCancelAck(sess *millSession, ca *millv1.CancelAck) { leaseID := ca.GetLeaseId() m.settleSessionCancel(sess, leaseID) m.mu.Lock() lease := m.leases[leaseID] m.mu.Unlock() if lease == nil || lease.nodeID != sess.nodeID { return } if !lease.markCancelAcked() { return } time.AfterFunc(m.cfg.CancelTeardownTimeout, func() { m.checkCancelTeardown(lease) })}
func (m *Mill) checkCancelAck(sess *millSession, lease *RemoteLease) { lease.mu.Lock() isDone := (lease.state == leaseDone) isAcked := lease.cancelAcked lease.mu.Unlock()
if isDone || isAcked { m.settleSessionCancel(sess, lease.id) return }
m.mu.Lock() isCurrent := m.sessions[lease.nodeID] == sess && !sess.disconnected m.mu.Unlock() if !isCurrent { return }
m.l.Warn("node did not acknowledge cancel request within deadline; closing session", "node", lease.nodeID, "lease", lease.id, "deadline", m.cfg.CancelAckTimeout) sess.close()}
// don't fail the node's other jobs because one teardown got stuckfunc (m *Mill) checkCancelTeardown(lease *RemoteLease) { if lease.getState() == leaseDone { return }
_, cancelReason := lease.cancelRequested() reason := fmt.Sprintf("%s (node teardown exceeded %s)", cancelReason, m.cfg.CancelTeardownTimeout) m.l.Warn("acked cancel exceeded teardown deadline; finishing lease without the node's terminal", "node", lease.nodeID, "lease", lease.id, "deadline", m.cfg.CancelTeardownTimeout)
status := string(models.StatusKindCancelled) if lease.orphaned { if err := m.finishOrphan(lease, status, &reason, nil); err != nil { m.l.Error("finish orphan after cancel teardown deadline failed", "lease", lease.id, "err", err) } return } if err := m.finishLiveLease(lease, status, reason); err != nil { m.l.Error("finish lease after cancel teardown deadline failed", "lease", lease.id, "err", err) }}
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) RegisterMetrics(metrics *observability.Metrics) { if metrics == nil { return } m.metrics = metrics metrics.RegisterMillGauges( func() float64 { m.mu.Lock() defer m.mu.Unlock() return float64(m.pending) }, func() float64 { return float64(m.cfg.MaxPending) }, func() float64 { m.mu.Lock() defer m.mu.Unlock() return float64(len(m.leases)) }, func() float64 { m.mu.Lock() defer m.mu.Unlock() return float64(len(m.reservations)) }, func() float64 { m.mu.Lock() defer m.mu.Unlock() active := 0 for _, sess := range m.sessions { if !sess.disconnected { active++ } } return float64(active) }, func() float64 { m.mu.Lock() defer m.mu.Unlock() disconnected := 0 for _, sess := range m.sessions { if sess.disconnected { disconnected++ } } return float64(disconnected) }, )}
func (m *Mill) authorizeAndLockQuota(s *millSession, reservationID, requestedLeaseID string) (*RemoteLease, error) { // snapshot the current owner before locking its teardown path m.quotaLifecycleMu.Lock() lease, ok := m.quotaLeases[reservationID] m.quotaLifecycleMu.Unlock()
if !ok || lease == nil { return nil, fmt.Errorf("quota reservation %q is not attached to a live lease", reservationID) }
// keep teardown from changing the lease during authorization lease.finishMu.Lock()
// serialize the final ownership check with quota phase transitions m.quotaLifecycleMu.Lock()
if m.quotaLeases[reservationID] != lease { m.quotaLifecycleMu.Unlock() lease.finishMu.Unlock() return nil, fmt.Errorf("quota reservation %q is not attached to a live lease", reservationID) } if lease.id != requestedLeaseID { m.quotaLifecycleMu.Unlock() lease.finishMu.Unlock() return nil, fmt.Errorf("quota reservation %q belongs to another workflow", reservationID) } if lease.nodeID != s.nodeID { m.quotaLifecycleMu.Unlock() lease.finishMu.Unlock() return nil, fmt.Errorf("quota reservation %q belongs to another executor", reservationID) } if lease.epoch != s.epoch { m.quotaLifecycleMu.Unlock() lease.finishMu.Unlock() return nil, fmt.Errorf("quota reservation %q has stale epoch %q (current: %q)", reservationID, lease.epoch, s.epoch) }
lease.mu.Lock() rq, hasQuota := lease.remoteQuotas[reservationID] active := hasQuota && (lease.state != leaseDone || rq.committing) lease.mu.Unlock() if !active { m.quotaLifecycleMu.Unlock() lease.finishMu.Unlock() return nil, fmt.Errorf("quota reservation %q is not active", reservationID) }
// the caller releases both locks after its phase transition return lease, nil}
func (m *Mill) handleQuotaRequest(s *millSession, req *millv1.QuotaRequest) { var ( res quota.Reservation err error ) switch req.GetOperation() { case millv1.QuotaOperation_QUOTA_OPERATION_RESERVE: res, err = m.reserveRemoteQuota(s, req) case millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT, millv1.QuotaOperation_QUOTA_OPERATION_COMMIT, millv1.QuotaOperation_QUOTA_OPERATION_RELEASE: err = m.transitionRemoteQuota(s, req) default: err = fmt.Errorf("unsupported quota operation %q", req.GetOperation()) }
resp := &millv1.QuotaResponse{ RequestId: req.GetRequestId(), ReservationId: res.ID, Allowed: res.Allowed, Temporary: res.Temporary, Reason: res.Reason, Resource: string(res.Resource), } if err != nil { resp.Error = err.Error() } _ = s.send(&millproto.Message{QuotaResp: resp})}
func (m *Mill) reserveRemoteQuota(s *millSession, req *millv1.QuotaRequest) (quota.Reservation, error) { if m.qm == nil { return quota.Reservation{}, errors.New("quota manager not attached") } if req.GetReservationId() != "" { return quota.Reservation{}, errors.New("reserve operation cannot name a reservation") } kind, resources, err := remoteQuotaPolicy(req) if err != nil { return quota.Reservation{}, err }
m.mu.Lock() lease := m.leases[req.GetLeaseId()] m.mu.Unlock() if lease == nil || lease.nodeID != s.nodeID || lease.epoch != s.epoch { return quota.Reservation{}, fmt.Errorf("no live lease %q on this node with matching epoch", req.GetLeaseId()) } lease.mu.Lock() id := quota.Identity{OwnerDID: lease.ownerDID, RepoDID: lease.repoDID} lease.mu.Unlock() if id.OwnerDID == "" || id.RepoDID == "" { return quota.Reservation{}, errors.New("lease has no charged subject") }
lease.finishMu.Lock() defer lease.finishMu.Unlock() if lease.cleanedUp || lease.getState() == leaseDone { return quota.Reservation{}, fmt.Errorf("lease %q ended before quota reservation", lease.id) }
m.quotaLifecycleMu.Lock() defer m.quotaLifecycleMu.Unlock() reservationKey := string(kind) + "\x00" + req.GetKey() lease.mu.Lock() var existingID string mismatch := false for resID, rq := range lease.remoteQuotas { if rq.key == reservationKey { existingID = resID mismatch = !maps.Equal(rq.resources, resources) break } } lease.mu.Unlock() if existingID != "" { // a repeat is only idempotent for the exact same charge; anything // else is an executor trying to inflate an existing reservation if mismatch { return quota.Reservation{}, fmt.Errorf("cache key %q already reserved with different resources", req.GetKey()) } return quota.Reservation{ ID: existingID, Allowed: true, Reason: quota.ReasonWithinLimit, }, nil }
qlease, res, err := m.qm.TryAcquire(context.Background(), quota.ReserveRequest{ Kind: kind, Key: req.GetKey(), Identity: id, Resources: resources, }) if err != nil { return quota.Reservation{}, err } if qlease == nil || res.ID == "" { return res, nil }
existingLease := m.quotaLeases[res.ID] if existingLease != nil && existingLease != lease { return quota.Reservation{}, fmt.Errorf("quota reservation %q is already owned by another workflow", res.ID) } m.quotaLeases[res.ID] = lease lease.mu.Lock() rq := lease.remoteQuotas[res.ID] rq.key = reservationKey rq.resources = resources lease.remoteQuotas[res.ID] = rq lease.mu.Unlock() return res, nil}
func remoteQuotaPolicy(req *millv1.QuotaRequest) (quota.Kind, quota.Resources, error) { kind := quota.Kind(req.GetKind()) if kind != quota.KindNixCache && kind != quota.KindGenericCache { return "", nil, fmt.Errorf("unsupported remote quota kind %q", req.GetKind()) } resources := req.GetResources() if len(resources) == 0 { return "", nil, errors.New("empty remote quota resources") } if err := quota.ValidateResources(resources); err != nil { return "", nil, fmt.Errorf("invalid remote quota resources: %w", err) } cacheBytes, ok := resources[quota.ResourceCacheStorageBytes] if !ok || cacheBytes <= 0 || len(resources) != 1 { return "", nil, errors.New("remote cache quota requires one positive cache_storage_bytes resource") } return kind, maps.Clone(resources), nil}
func (m *Mill) transitionRemoteQuota(s *millSession, req *millv1.QuotaRequest) error { reservationID := req.GetReservationId() if reservationID == "" { return errors.New("quota transition requires a reservation") } if req.GetKind() != "" || req.GetKey() != "" || req.GetResources() != nil { return errors.New("quota transition cannot include reservation parameters") }
// successful terminal transitions are idempotent if req.GetOperation() != millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT { m.quotaLifecycleMu.Lock() _, ok := m.quotaLeases[reservationID] m.quotaLifecycleMu.Unlock() if !ok { return nil } }
lease, err := m.authorizeAndLockQuota(s, reservationID, req.GetLeaseId()) if err != nil { return err } defer lease.finishMu.Unlock() defer m.quotaLifecycleMu.Unlock()
if m.qm == nil { return errors.New("quota manager not attached") } switch req.GetOperation() { case millv1.QuotaOperation_QUOTA_OPERATION_BEGIN_COMMIT: if err := m.qm.BeginCommit(context.Background(), reservationID); err != nil { return err } lease.mu.Lock() rq := lease.remoteQuotas[reservationID] rq.committing = true lease.remoteQuotas[reservationID] = rq lease.mu.Unlock() case millv1.QuotaOperation_QUOTA_OPERATION_COMMIT: if err := m.qm.Commit(context.Background(), reservationID); err != nil { return err } m.forgetQuotaReservationLocked(reservationID) case millv1.QuotaOperation_QUOTA_OPERATION_RELEASE: if err := m.qm.Release(context.Background(), reservationID); err != nil { return err } m.forgetQuotaReservationLocked(reservationID) } return nil}
func (m *Mill) forgetQuotaReservationLocked(reservationID string) { lease := m.quotaLeases[reservationID] delete(m.quotaLeases, reservationID) if lease == nil { return } lease.mu.Lock() delete(lease.remoteQuotas, reservationID) lease.mu.Unlock()}