Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561package executor
import ( "context" "encoding/json" "errors" "fmt" "log/slog" "maps" "net/http" "runtime" "strings" "sync" "time"
"github.com/bluesky-social/indigo/atproto/syntax"
"tangled.org/core/api/tangled" "tangled.org/core/notifier" "tangled.org/core/spindle/config" "tangled.org/core/spindle/db" "tangled.org/core/spindle/engine" "tangled.org/core/spindle/models" "tangled.org/core/util/netutil"
millproto "tangled.org/core/spindle/mill/proto" millv1 "tangled.org/core/spindle/mill/proto/gen")
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
relay *relayLog relayMu sync.Mutex // offset assign + send must be serialized so wire order == offset order
connMu sync.Mutex enc messageEncoder
mu sync.Mutex active map[string]*reservation draining bool}
type reservation struct { leaseID string wid models.WorkflowId realEngine models.Engine slot engine.WorkflowSlot wf *models.Workflow repoDid syntax.DID
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) *Executor { seats := cfg.Mill.Seats if seats <= 0 { seats = defaultSeats } labels := normalizeLabels(cfg.Mill.Labels) return &Executor{ millURL: cfg.Mill.URL, token: cfg.Mill.SharedSecret, nodeID: cfg.Server.Hostname, seats: seats, labels: labels, engines: engines, db: d, n: n, cfg: cfg, l: l.With("component", "mill.executor"), relay: newRelayLog(), active: make(map[string]*reservation), }}
func (e *Executor) Connect(ctx context.Context) { 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) go e.observeLoop(ctx, sub, cursor) defer e.n.Unsubscribe(sub)
backoff := dialBackoffMin for { if ctx.Err() != nil { return } err := e.runSession(ctx) if ctx.Err() != nil { return } e.l.Warn("mill session ended; reconnecting", "err", err, "backoff", backoff) select { case <-ctx.Done(): return case <-time.After(backoff): } backoff = min(backoff*2, dialBackoffMax) }}
func (e *Executor) runSession(ctx context.Context) error { dev := 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, }} 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") }
// install encoder and replay anything the mill missed. e.connMu.Lock() e.enc = enc e.connMu.Unlock() defer func() { e.connMu.Lock() e.enc = 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) } }()
e.replay(resume.GetAckOffset()) e.pushSnapshot() e.l.Info("connected to mill", "node", e.nodeID, "resumeFrom", resume.GetAckOffset())
go e.snapshotLoop(sessionCtx, enc) return <-readErr}
func (e *Executor) replay(ackOffset uint64) { e.relayMu.Lock() defer e.relayMu.Unlock() e.relay.ack(ackOffset) for _, msg := range e.relay.since(ackOffset) { e.send(msg) }}
func (e *Executor) send(msg *millproto.Message) { e.connMu.Lock() enc := e.enc e.connMu.Unlock() if enc != nil { _ = enc.Encode(msg) }}
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.relay.ack(msg.GetAck().GetUpToOffset()) default: e.l.Warn("executor received unexpected message") }}
// --- reserve / commit / release / cancel ----------------------------------
func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) { reject := func(reason string, class millv1.RejectClass) { e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ LeaseId: rs.GetLeaseId(), Accepted: false, RejectReason: reason, RejectClass: 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 }
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 } 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 } } // the job skipped processPipeline, so inject TANGLED_* env here. if wf.Environment == nil { wf.Environment = make(map[string]string) } maps.Copy(wf.Environment, models.PipelineEnvVars(tpl.TriggerMetadata, pipelineId))
// NoWait: the executor doesn't queue locally; the mill owns the backlog. 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 }
var repoDid syntax.DID if tpl.TriggerMetadata != nil && tpl.TriggerMetadata.Repo != nil && tpl.TriggerMetadata.Repo.RepoDid != nil { repoDid, _ = syntax.ParseDID(*tpl.TriggerMetadata.Repo.RepoDid) }
res := &reservation{ leaseID: rs.GetLeaseId(), wid: wid, realEngine: realEngine, slot: slot, wf: wf, repoDid: repoDid, } res.ttlTimer = time.AfterFunc(ttlDuration(rs.GetTtlSeconds()), func() { e.expireReservation(res.leaseID) })
e.mu.Lock() e.active[res.leaseID] = res e.mu.Unlock()
e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ LeaseId: rs.GetLeaseId(), Accepted: true, }}) e.pushSnapshot()}
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.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{ LeaseId: cl.GetLeaseId(), Accepted: false, RejectReason: "reservation missing or expired", RejectClass: millv1.RejectClass_REJECT_CLASS_TRANSIENT, }}) return } if res.committed { e.mu.Unlock() e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: cl.GetLeaseId()}}) return } res.committed = true if res.ttlTimer != nil { res.ttlTimer.Stop() } jobCtx, cancel := context.WithCancel(ctx) res.cancel = cancel e.mu.Unlock()
vault := newMemVault(cl.GetSecrets()) re := newReservedEngine(res.realEngine, res.slot) pipeline := &models.Pipeline{ RepoDid: res.repoDid, Workflows: map[models.Engine][]models.Workflow{re: {*res.wf}}, TrustedSource: true, }
e.startTail(res)
go engine.StartWorkflows(e.l, vault, e.cfg, e.db, e.n, jobCtx, pipeline, res.wid.PipelineId)
e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: cl.GetLeaseId()}})}
func (e *Executor) handleRelease(leaseID string) { cleanup, ok := e.takeUncommittedReservation(leaseID, true) if !ok { return } cleanup() e.pushSnapshot()}
func (e *Executor) handleCancel(leaseID string) { e.mu.Lock() res := e.active[leaseID] if res == nil { e.mu.Unlock() return } res.cancelled = true cancel := res.cancel committed := res.committed var cleanup func() if !committed { cleanup = e.removeReservationLocked(res, true) } e.mu.Unlock()
if cancel != nil { cancel() } // committed jobs clean up when the observe loop sees the terminal row; an // uncommitted cancel still needs the slot released. if !committed { cleanup() e.pushSnapshot() }}
func (e *Executor) expireReservation(leaseID string) { cleanup, ok := e.takeUncommittedReservation(leaseID, true) if !ok { return } e.l.Warn("reservation expired before commit", "lease", leaseID) cleanup() e.pushSnapshot()}
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 func() {}, false } return e.removeReservationLocked(res, releaseSlot), true}
func (e *Executor) finishReservation(res *reservation) (func(), bool, bool) { e.mu.Lock() defer e.mu.Unlock() if e.active[res.leaseID] != res { return func() {}, false, false } cancelled := res.cancelled return e.removeReservationLocked(res, false), cancelled, true}
func (e *Executor) removeReservationLocked(res *reservation, releaseSlot bool) func() { delete(e.active, res.leaseID) if res.ttlTimer != nil { res.ttlTimer.Stop() res.ttlTimer = nil } stopTail := res.stopTail res.stopTail = nil slot := res.slot if releaseSlot { res.slot = nil } return func() { if stopTail != nil { stopTail() } if releaseSlot && slot != nil { slot.Release() } }}
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}
// --- snapshots -------------------------------------------------------------
func (e *Executor) snapshotLoop(ctx context.Context, enc *millproto.Encoder) { t := time.NewTicker(snapshotEvery) defer t.Stop() for { select { case <-ctx.Done(): return case <-t.C: // stop pushing once this connection is replaced. e.connMu.Lock() cur := e.enc e.connMu.Unlock() if cur != enc { return } e.pushSnapshot() } }}
func (e *Executor) pushSnapshot() { e.mu.Lock() draining := e.draining active := len(e.active) leaseIDs := make([]string, 0, len(e.active)) for id := range e.active { leaseIDs = append(leaseIDs, id) } e.mu.Unlock()
load := 0.0 if e.seats > 0 { load = float64(active) / float64(e.seats) } available := !draining && active < e.seats engines := make(map[string]*millv1.EngineAvailability, len(e.engines)) for name, realEngine := range e.engines { var capabilities []string if provider, ok := realEngine.(engine.CapabilityProvider); ok { var err error capabilities, err = engine.NormalizeCapabilities(provider.Capabilities()) if err != nil { e.l.Error("engine advertised invalid placement capabilities", "engine", name, "err", err) capabilities = nil } } engines[name] = &millv1.EngineAvailability{ Available: available, Load: map[string]float64{"slots": load}, Capabilities: capabilities, } } e.send(&millproto.Message{NodeSnapshot: &millv1.NodeSnapshot{ NodeId: e.nodeID, Engines: engines, ActiveLeaseIds: leaseIDs, }})}
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