Monorepo for Tangled
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634package executor
import ( "context" "encoding/json" "errors" "fmt" "maps" "net/http" "runtime" "sync" "time"
"github.com/bluesky-social/indigo/atproto/syntax" "log/slog" "strings"
"tangled.org/core/api/tangled" "tangled.org/core/netutil" "tangled.org/core/notifier" "tangled.org/core/spindle/artifactstore" "tangled.org/core/spindle/config" "tangled.org/core/spindle/db" "tangled.org/core/spindle/engine" millproto "tangled.org/core/spindle/mill/proto" millv1 "tangled.org/core/spindle/mill/proto/gen" "tangled.org/core/spindle/models" "tangled.org/core/spindle/storage")
const ( dialBackoffMin = 1 * time.Second dialBackoffMax = 30 * time.Second snapshotEvery = 15 * time.Second defaultSeats = 4)
type Executor struct { millURL string token string nodeID string seats int labels []string
engines map[string]models.Engine db *db.DB n *notifier.Notifier cfg *config.Config l *slog.Logger writer artifactstore.Writer cache storage.Storage
epoch string outboxBytes int64 maxOutboxBytes int64 // 10 MiB default outbox cap
eventMu sync.Mutex sendMu sync.Mutex flushMu sync.Mutex
sentSeqno uint64
connMu sync.Mutex enc messageEncoder sessionCancel context.CancelFunc
mu sync.Mutex active map[string]*reservation draining bool
snapshotMu sync.Mutex nextSeqno uint64
lifecycleCtx context.Context jobsWG sync.WaitGroup}
type reservation struct { leaseID string wid models.WorkflowId realEngine models.Engine slot engine.WorkflowSlot wf *models.Workflow pipeline *models.Pipeline vault *memVault
committed bool cancelled bool cancel context.CancelFunc ttlTimer *time.Timer stopTail func()}
type messageEncoder interface { Encode(*millproto.Message) error}
func New(cfg *config.Config, engines map[string]models.Engine, d *db.DB, n *notifier.Notifier, l *slog.Logger, writer artifactstore.Writer, cache storage.Storage) (*Executor, error) { seats := defaultSeats millURL := "" token := "" nodeID := "" var labels []string if cfg != nil { if cfg.Mill.Seats > 0 { seats = cfg.Mill.Seats } labels = normalizeLabels(cfg.Mill.Labels) millURL = cfg.Mill.URL token = cfg.Mill.SharedSecret nodeID = cfg.Server.Hostname } if d == nil || n == nil { return nil, fmt.Errorf("executor requires a database and notifier") } e := &Executor{ millURL: millURL, token: token, nodeID: nodeID, seats: seats, labels: labels, engines: engines, db: d, n: n, cfg: cfg, l: l.With("component", "mill.executor"), writer: writer, cache: cache, active: make(map[string]*reservation), maxOutboxBytes: 10 * 1024 * 1024, } if err := e.initOutbox(); err != nil { return nil, fmt.Errorf("initialize executor outbox: %w", err) } return e, nil}
func (e *Executor) Connect(ctx context.Context) { e.lifecycleCtx = ctx sub := e.n.Subscribe() cursor, err := e.db.EventHighWater() if err != nil { e.n.Unsubscribe(sub) e.l.Error("establish event cursor failed", "err", err) return } e.drainEvents(&cursor) observerCtx, stopObserver := context.WithCancel(ctx) observerDone := make(chan struct{}) go func() { defer close(observerDone) e.observeLoop(observerCtx, sub, cursor) }() defer e.n.Unsubscribe(sub)
backoff := dialBackoffMin for { if ctx.Err() != nil { break } err := e.runSession(ctx) if ctx.Err() != nil { break } e.l.Warn("mill session ended; reconnecting", "err", err, "backoff", backoff) select { case <-ctx.Done(): break case <-time.After(backoff): } backoff = min(backoff*2, dialBackoffMax) }
e.jobsWG.Wait() stopObserver() <-observerDone e.drainEvents(&cursor)}
func (e *Executor) runSession(ctx context.Context) error { dev := e.cfg == nil || e.cfg.Server.Dev if _, err := netutil.EnforceWSSURL(e.millURL, dev); err != nil { return fmt.Errorf("mill url: %w", err) } header := http.Header{} if e.token != "" { header.Set("Authorization", "Bearer "+e.token) } conn, _, err := netutil.SSRFWebsocketDialer(dev).DialContext(ctx, e.millURL, header) if err != nil { return fmt.Errorf("dial mill: %w", err) } defer conn.Close()
sessionCtx, cancelSession := context.WithCancel(ctx) defer cancelSession() stopClose := context.AfterFunc(sessionCtx, func() { _ = conn.Close() }) defer stopClose()
stream := millproto.NewWSStream(conn) enc := millproto.NewEncoder(stream) dec := millproto.NewDecoder(stream)
hello := &millproto.Message{Hello: &millv1.Hello{ ProtocolVersion: millproto.ProtocolVersion, Arch: runtime.GOARCH, Labels: e.labels, Epoch: e.epoch, CacheStoreId: e.cfg.Cache.StoreID, CacheNamespace: models.CacheNamespace(), }} if err := enc.Encode(hello); err != nil { return fmt.Errorf("send hello: %w", err) }
resumeMsg, err := dec.Decode() if err != nil { return fmt.Errorf("read resume: %w", err) } resume := resumeMsg.GetResume() if resume == nil { return fmt.Errorf("expected resume, got something else") } if resume.GetEpoch() != e.epoch { return fmt.Errorf("resume epoch mismatch: got %q, want %q", resume.GetEpoch(), e.epoch) }
e.connMu.Lock() e.sessionCancel = cancelSession e.enc = enc e.connMu.Unlock() defer func() { e.connMu.Lock() e.enc = nil e.sessionCancel = nil e.connMu.Unlock() }()
readErr := make(chan error, 1) go func() { for { msg, err := dec.Decode() if err != nil { readErr <- fmt.Errorf("read: %w", err) return } e.dispatch(sessionCtx, msg) } }()
if err := e.replay(resume.GetAckSeqno()); err != nil { cancelSession() <-readErr return fmt.Errorf("replay failed: %w", err) } e.pushSnapshot() e.l.Info("connected to mill", "node", e.nodeID, "resumeFrom", resume.GetAckSeqno())
go e.snapshotLoop(sessionCtx, enc) return <-readErr}func (e *Executor) send(msg *millproto.Message) { e.connMu.Lock() enc := e.enc cancel := e.sessionCancel e.connMu.Unlock() if enc != nil { e.sendMu.Lock() err := enc.Encode(msg) e.sendMu.Unlock() if err != nil { e.l.Error("send failed, ending session", "err", err) if cancel != nil { cancel() } } }}
func (e *Executor) dispatch(ctx context.Context, msg *millproto.Message) { switch { case msg.GetReserveSeat() != nil: e.handleReserve(ctx, msg.GetReserveSeat()) case msg.GetCommitLease() != nil: e.handleCommit(ctx, msg.GetCommitLease()) case msg.GetReleaseLease() != nil: e.handleRelease(msg.GetReleaseLease().GetLeaseId()) case msg.GetCancelAttempt() != nil: e.handleCancel(msg.GetCancelAttempt().GetLeaseId()) case msg.GetAck() != nil: e.handleAck(msg.GetAck()) default: e.l.Warn("unhandled incoming message", "type", fmt.Sprintf("%T", msg)) }}
func (e *Executor) sendReject(leaseID string, reason string, class millv1.RejectClass) { e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ LeaseId: leaseID, Accepted: false, RejectReason: reason, RejectClass: class, }})}
func (e *Executor) sendCommitted(leaseID string) { e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}})}
func (e *Executor) sendCancelAck(leaseID string) { e.send(&millproto.Message{CancelAck: &millv1.CancelAck{LeaseId: leaseID}})}
func (e *Executor) releaseReservation(cleanup func()) { if cleanup != nil { cleanup() } e.pushSnapshot()}
func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { reject := func(reason string, class millv1.RejectClass) { e.sendReject(rs.GetLeaseId(), reason, class) }
e.mu.Lock() draining := e.draining e.mu.Unlock() if draining { reject("draining", millv1.RejectClass_REJECT_CLASS_TRANSIENT) return }
realEngine, ok := e.engines[rs.GetTargetEngine()] if !ok { reject("unknown engine "+rs.GetTargetEngine(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } slotter, ok := realEngine.(engine.WorkflowSlotter) if !ok { reject("engine does not support workflow slots", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return }
var twf tangled.Pipeline_Workflow if err := json.Unmarshal([]byte(rs.GetRawWorkflowJson()), &twf); err != nil { reject("bad workflow json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } var tpl tangled.Pipeline if err := json.Unmarshal([]byte(rs.GetRawPipelineJson()), &tpl); err != nil { reject("bad pipeline json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } if tpl.TriggerMetadata == nil { reject("pipeline missing trigger metadata", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return }
repoDid, err := syntax.ParseDID(rs.GetRepoDid()) if err != nil { reject("bad repository did", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } trustedSource := models.TrustedPipelineSource(tpl.TriggerMetadata, repoDid.String())
pipelineId := models.PipelineId{Knot: rs.GetKnot(), Rkey: rs.GetRkey()} wid := models.WorkflowId{PipelineId: pipelineId, Name: twf.Name}
wf, err := realEngine.InitWorkflow(twf, tpl) if err != nil { reject("init workflow: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } wf.Engine = rs.GetTargetEngine() wf.RepoDid = rs.GetRepoDid() if validator, ok := realEngine.(engine.WorkflowPlacementValidator); ok { if err := validator.ValidateWorkflowPlacement(wf); err != nil { reject("validate workflow placement: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } } if wf.Environment == nil { wf.Environment = make(map[string]string) } maps.Copy(wf.Environment, models.PipelineEnvVars(tpl.TriggerMetadata, pipelineId))
slot, err := slotter.AcquireWorkflowSlot(ctx, wid, wf, engine.NoWait) if err != nil { class := millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE if errors.Is(err, engine.ErrNoWorkflowSlots) { class = millv1.RejectClass_REJECT_CLASS_TRANSIENT } reject(err.Error(), class) return }
res := &reservation{ leaseID: rs.GetLeaseId(), wid: wid, realEngine: realEngine, slot: slot, wf: wf, pipeline: &models.Pipeline{ RepoDid: repoDid, TrustedSource: trustedSource, TriggerMetadata: tpl.TriggerMetadata, }, }
e.snapshotMu.Lock() e.mu.Lock() e.active[res.leaseID] = res res.ttlTimer = time.AfterFunc(ttlDuration(rs.GetTtlSeconds()), func() { e.expireReservation(res.leaseID) }) e.mu.Unlock()
e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ LeaseId: rs.GetLeaseId(), Accepted: true, }}) e.pushSnapshotLocked() e.snapshotMu.Unlock()}
func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) { e.mu.Lock() res := e.active[cl.GetLeaseId()] if res == nil { e.mu.Unlock() e.sendReject(cl.GetLeaseId(), "reservation missing or expired", millv1.RejectClass_REJECT_CLASS_TRANSIENT) return } if res.committed { e.mu.Unlock() e.sendCommitted(cl.GetLeaseId()) return } bindings, err := cacheBindingsFromProto(res.pipeline.RepoDid.String(), res.wf.Caches, cl.GetCacheBindings()) if err != nil { e.mu.Unlock() e.sendReject(cl.GetLeaseId(), "invalid cache plan: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE) return } res.wf.CacheBindings = bindings res.committed = true if res.ttlTimer != nil { res.ttlTimer.Stop() }
committedSecrets := cl.GetSecrets() if !res.pipeline.TrustedSource { committedSecrets = nil } res.vault = newMemVault(committedSecrets) jobCtx, cancel := context.WithCancel(e.lifecycleCtx) res.cancel = cancel e.mu.Unlock()
vault := res.vault re := newReservedEngine(res.realEngine, res.slot) res.pipeline.Workflows = map[models.Engine][]models.Workflow{re: {*res.wf}} cacheController := engine.NewPreplannedCacheController(func(ctx context.Context, update engine.CacheUpdate) error { return e.appendCacheUpdate(res.leaseID, update) })
e.startTail(res)
e.jobsWG.Add(1) go func() { defer e.jobsWG.Done() engine.StartWorkflows(e.l, vault, e.cfg, nil, e.db, e.n, e.cache, cacheController, jobCtx, res.pipeline, res.wid.PipelineId) }()
e.sendCommitted(cl.GetLeaseId())}
func (e *Executor) handleRelease(leaseID string) { cleanup, ok := e.takeUncommittedReservation(leaseID, true) if !ok { return } e.releaseReservation(cleanup)}
func (e *Executor) handleCancel(leaseID string) { e.mu.Lock() res := e.active[leaseID] if res == nil { e.mu.Unlock() if err := e.appendTerminal(leaseID, string(models.StatusKindCancelled), nil); err != nil { e.l.Error("persist cancelled reservation terminal", "lease", leaseID, "err", err) return } e.sendCancelAck(leaseID) return } res.cancelled = true cancel := res.cancel committed := res.committed var cleanup func() if !committed { cleanup = e.removeReservationLocked(res, true) } e.mu.Unlock()
if !committed { if err := e.appendTerminal(leaseID, string(models.StatusKindCancelled), nil); err != nil { e.l.Error("persist cancelled reservation terminal", "lease", leaseID, "err", err) e.releaseReservation(cleanup) return } e.sendCancelAck(leaseID) e.releaseReservation(cleanup) return }
e.sendCancelAck(leaseID) if cancel != nil { cancel() }}
func (e *Executor) expireReservation(leaseID string) { cleanup, ok := e.takeUncommittedReservation(leaseID, true) if !ok { return } e.releaseReservation(cleanup)}
func (e *Executor) takeUncommittedReservation(leaseID string, releaseSlot bool) (func(), bool) { e.mu.Lock() defer e.mu.Unlock() res := e.active[leaseID] if res == nil || res.committed { return nil, false } if res.ttlTimer != nil { res.ttlTimer.Stop() } return e.removeReservationLocked(res, releaseSlot), true}
func (e *Executor) removeReservationLocked(res *reservation, releaseSlot bool) func() { delete(e.active, res.leaseID) slot := res.slot return func() { if releaseSlot && slot != nil { slot.Release() } }}
func (e *Executor) snapshotLoop(ctx context.Context, enc *millproto.Encoder) { ticker := time.NewTicker(snapshotEvery) defer ticker.Stop()
for { select { case <-ctx.Done(): return case <-ticker.C: e.pushSnapshot() } }}
func (e *Executor) pushSnapshot() { e.snapshotMu.Lock() defer e.snapshotMu.Unlock() e.pushSnapshotLocked()}
func (e *Executor) pushSnapshotLocked() { e.mu.Lock() activeLeases := make([]string, 0, len(e.active)) for leaseID := range e.active { activeLeases = append(activeLeases, leaseID) } e.mu.Unlock()
avail := make(map[string]*millv1.EngineAvailability) for name, eng := range e.engines { a := &millv1.EngineAvailability{Available: true} if getter, ok := eng.(interface{ Load() map[string]float64 }); ok { a.Load = getter.Load() } avail[name] = a }
e.nextSeqno++ snap := &millproto.Message{ NodeSnapshot: &millv1.NodeSnapshot{ Seqno: e.nextSeqno, Engines: avail, ActiveLeaseIds: activeLeases, }, } e.send(snap)}
func (e *Executor) Drain() { e.mu.Lock() e.draining = true e.mu.Unlock() e.pushSnapshot()}
func ttlDuration(secs uint32) time.Duration { if secs == 0 { return defaultReservationTTL } return time.Duration(secs) * time.Second}
const defaultReservationTTL = 60 * time.Second
func normalizeLabels(labels []string) []string { seen := make(map[string]struct{}, len(labels)) out := make([]string, 0, len(labels)) for _, label := range labels { label = strings.TrimSpace(label) if label == "" { continue } if _, ok := seen[label]; ok { continue } seen[label] = struct{}{} out = append(out, label) } return out}