package 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
CacheQuota *db.QuotaStore
CacheMaxBytes 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
NixCaches bool
}
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
debugKeys DebugKeyChecker
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{}
changeGen uint64
waiters map[*placeWaiter]struct{}
waiterSeq uint64
cacheDeleteOnce sync.Once
cacheDeleteMu sync.Mutex
cacheDeleteQueue []cacheDeletion
cacheDeleteWake chan struct{}
nixUseSlots 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),
waiters: make(map[*placeWaiter]struct{}),
nixUseSlots: make(chan struct{}, nixUseConcurrency),
}
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()
defer m.mu.Unlock()
return m.nextLeaseIDLocked()
}
func (m *Mill) nextLeaseIDLocked() string {
m.leaseSeq++
return fmt.Sprintf("%s-%d", tid.TID(), m.leaseSeq)
}
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() {
m.changeGen++
m.wakeWaitersLocked()
}
func (m *Mill) wakeWaitersLocked() {
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) extendGraceLocked(sess *millSession) {
if !sess.recovering || sess.recoveryExpired {
return
}
deadline := time.Now().Add(m.cfg.ReconnectGrace)
if deadline.After(sess.graceDeadline) {
sess.graceDeadline = deadline
if sess.graceTimer != nil {
sess.graceTimer.Stop()
}
delay := time.Until(deadline)
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 stick
func (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))
}()
if err := m.checkExecutorExists(wf); err != nil {
placementResult = "rejected"
return nil, err
}
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()
}()
owner, repo := wf.OwnerDID, wf.RepoDID
if id, ok := m.quotaIdentity(wf); ok {
owner, repo = id.OwnerDID, id.RepoDID
}
w := m.joinPlacement(owner, repo, engineName)
defer m.leavePlacement(w)
for {
if err := ctx.Err(); err != nil {
placementResult = "cancelled"
if errors.Is(err, context.DeadlineExceeded) {
placementResult = "timeout"
}
return nil, err
}
// the channel comes from before the bid, so a change during the bid
// makes the next wait return immediately
turn, gen, ch := m.placementTurn(w)
if turn {
lease, err := m.bid(ctx, engineName, wid, wf)
if err != nil {
return nil, err
}
if lease == nil {
m.passTurn(w, gen)
} else {
wf.CacheNamespace = lease.cacheNamespace
m.setAdmitting(w, true)
slot, retry, err := m.admit(ctx, engineName, wid, wf, lease)
m.setAdmitting(w, false)
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, res.Err())
}
// don't hold fleet capacity while waiting for quota
m.releaseSeat(lease)
qlease, err = m.qm.Acquire(ctx, req)
if errors.Is(err, quota.ErrDenied) {
return nil, false, engine.ClassifiedFailure(engine.FailureClassPolicy, engine.FailureReasonQuotaDenied, err)
}
if err != nil {
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 once
func (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, string(wid.PipelineId), 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) {
if wf == nil {
return quota.Identity{}, false
}
state, ok := wf.Data.(*millWorkflowState)
if !ok || state == nil || state.RawPipeline.TriggerMetadata == nil {
return quota.Identity{}, false
}
repo := state.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 !id.Valid() {
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.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,
PipelineId: string(wid.PipelineId),
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()
lease.runBudget = time.Duration(rr.GetRunBudgetSeconds()) * time.Second
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
}
// a registered executor counts even while it's offline
func (m *Mill) checkExecutorExists(wf *models.Workflow) error {
required := requiredLabels(wf)
m.mu.Lock()
d := m.db
for _, s := range m.sessions {
if hasLabels(s.labels, required) {
m.mu.Unlock()
return nil
}
}
m.mu.Unlock()
if d == nil {
return nil
}
registrations, err := d.ListExecutorRegistrations()
if err != nil {
m.l.Warn("failed to list registered executors, placing anyway", "err", err)
return nil
}
now := time.Now()
for _, r := range registrations {
if r.ExpiresAt != nil && now.After(*r.ExpiresAt) {
continue
}
if hasLabels(r.Labels, required) {
return nil
}
}
err = errors.New("no executor is registered")
if len(required) > 0 {
err = fmt.Errorf("no executor has the labels %s", strings.Join(required, ", "))
}
return engine.ClassifiedFailure(engine.FailureClassPolicy, engine.FailureReasonNoExecutor, err)
}
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.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),
AcceptsAttemptOutcome: true,
NixCache: nixCacheToProto(wf.NixCache),
}}
// 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 terminalStatusString(status millv1.TerminalStatus) (string, error) {
switch status {
case millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS:
return string(models.StatusKindSuccess), nil
case millv1.TerminalStatus_TERMINAL_STATUS_FAILED:
return string(models.StatusKindFailed), nil
case millv1.TerminalStatus_TERMINAL_STATUS_TIMEOUT:
return string(models.StatusKindTimeout), nil
case millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED:
return string(models.StatusKindCancelled), nil
default:
return "", protoErrf("unsupported terminal status %v", status)
}
}
func statusDetails(errMsg string, exitCode int64) (*string, *int64) {
var e *string
if errMsg != "" {
e = &errMsg
}
var c *int64
if exitCode != 0 {
c = &exitCode
}
return e, c
}
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
}
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()
}
}
// precompute storage sizes so object-storage stat does not run inside the write tx
precomputedSizes := make(map[string]int64)
for _, entry := range newEntries {
if u := entry.GetCacheUpdate(); u != nil && u.GetAction() == millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED {
var storageKey string
if m.db != nil {
_ = m.db.QueryRow(`select storage_key from mill_cache_capabilities where lease_id = ? and action = 'save' and cache_id = ?`, entry.LeaseId, u.GetId()).Scan(&storageKey)
}
if storageKey != "" {
if sz, err := m.cacheObjectSize(sess.ctx, storageKey); err == nil {
precomputedSizes[storageKey] = sz
}
}
}
}
type batchApply struct {
terminalCancelIDs []string
pendingTerminals []pendingTerminal
artifactLeases []*RemoteLease
startupObservations []startupObservation
highestSeqno uint64
cacheDeletes map[string]db.CacheEntry
cacheObjectDeletes map[string]struct{}
}
applyFunc := func(tx *db.EventBatchTx) (batchApply, error) {
var r batchApply
r.highestSeqno = current
r.cacheDeletes = make(map[string]db.CacheEntry)
r.cacheObjectDeletes = make(map[string]struct{})
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
}
if tx != nil {
alreadyRunning, err := tx.HasWorkflowStatus(context.Background(), lease.wid, statusStr)
if err != nil {
return batchApply{}, err
}
if err := tx.InsertStatusEvent(lease.wid, statusStr, errMsg, exitCode); err != nil {
return batchApply{}, err
}
if !alreadyRunning {
delay, ok, err := tx.WorkflowStartupDelay(context.Background(), lease.wid)
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, ok := precomputedSizes[ref]
if !ok {
var err error
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.CacheQuota,
m.cfg.CacheMaxBytes,
at,
)
if err != nil {
return batchApply{}, err
}
for _, old := range superseded {
r.cacheDeletes[old.ID] = old
}
if !ready {
if err := tx.QueueCacheObjectDeletion(context.Background(), ref, at); err != nil {
return batchApply{}, err
}
r.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 {
r.cacheDeletes[discarded.ID] = *discarded
}
}
}
case entry.GetAttemptOutcome() != nil:
out := entry.GetAttemptOutcome()
statusStr, err := terminalStatusString(out.GetStatus())
if err != nil {
return batchApply{}, err
}
if tx != nil {
errMsg, exitCode := statusDetails(out.GetError(), out.GetExitCode())
if err := tx.InsertStatusChange(lease.wid, statusStr, errMsg, exitCode); err != nil {
return batchApply{}, err
}
}
case entry.GetAttemptResult() != nil:
ar := entry.GetAttemptResult()
statusStr, err := terminalStatusString(ar.GetStatus())
if err != nil {
return batchApply{}, err
}
finishedInBatch[lease.id] = struct{}{}
if tx != nil {
// an AttemptOutcome may have published this status already
errMsg, exitCode := statusDetails(ar.GetError(), ar.GetExitCode())
if err := tx.InsertStatusChange(lease.wid, 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.PipelineId == "" || 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.ApplyEventBatchCtx(sess.ctx, 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(applied.cacheDeletes, applied.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
}
if sess.recovering && !sess.recoveryExpired {
m.extendGraceLocked(sess)
}
m.maybeFinishSessionRecoveryLocked(sess)
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 stuck
func (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)
},
)
metrics.RegisterMillWorkflows(m.workflowCounts)
}
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")
}
nixObject, err := nixObjectFromProto(req.GetNixObject(), kind, m.leaseNixCache(id.RepoDID))
if err != nil {
return quota.Reservation{}, err
}
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,
NixObject: nixObject,
})
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 {
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()
}