alice.pds.demo.boltless.dev
This repository has no description
1package executor
2
3import (
4 "context"
5 "encoding/json"
6 "errors"
7 "fmt"
8 "maps"
9 "net/http"
10 "runtime"
11 "sync"
12 "time"
13
14 "github.com/bluesky-social/indigo/atproto/syntax"
15 "log/slog"
16 "strings"
17
18 "tangled.org/core/api/tangled"
19 "tangled.org/core/netutil"
20 "tangled.org/core/notifier"
21 "tangled.org/core/spindle/artifactstore"
22 "tangled.org/core/spindle/config"
23 "tangled.org/core/spindle/db"
24 "tangled.org/core/spindle/engine"
25 millproto "tangled.org/core/spindle/mill/proto"
26 millv1 "tangled.org/core/spindle/mill/proto/gen"
27 "tangled.org/core/spindle/models"
28)
29
30const (
31 dialBackoffMin = 1 * time.Second
32 dialBackoffMax = 30 * time.Second
33 snapshotEvery = 15 * time.Second
34 defaultSeats = 4
35)
36
37type Executor struct {
38 millURL string
39 token string
40 nodeID string
41 seats int
42 labels []string
43
44 engines map[string]models.Engine
45 db *db.DB
46 n *notifier.Notifier
47 cfg *config.Config
48 l *slog.Logger
49 writer artifactstore.Writer
50
51 epoch string
52 outboxBytes int64
53 maxOutboxBytes int64 // 10 MiB default outbox cap
54
55 eventMu sync.Mutex
56 sendMu sync.Mutex
57 flushMu sync.Mutex
58
59 sentSeqno uint64
60
61 connMu sync.Mutex
62 enc messageEncoder
63 sessionCancel context.CancelFunc
64
65 mu sync.Mutex
66 active map[string]*reservation
67 draining bool
68
69 snapshotMu sync.Mutex
70 nextSeqno uint64
71
72 lifecycleCtx context.Context
73 jobsWG sync.WaitGroup
74}
75
76type reservation struct {
77 leaseID string
78 wid models.WorkflowId
79 realEngine models.Engine
80 slot engine.WorkflowSlot
81 wf *models.Workflow
82 repoDid syntax.DID
83 vault *memVault
84
85 committed bool
86 cancelled bool
87 cancel context.CancelFunc
88 ttlTimer *time.Timer
89 stopTail func()
90}
91
92type messageEncoder interface {
93 Encode(*millproto.Message) error
94}
95
96func New(cfg *config.Config, engines map[string]models.Engine, d *db.DB, n *notifier.Notifier, l *slog.Logger, writers ...artifactstore.Writer) (*Executor, error) {
97 seats := defaultSeats
98 millURL := ""
99 token := ""
100 nodeID := ""
101 var labels []string
102 if cfg != nil {
103 if cfg.Mill.Seats > 0 {
104 seats = cfg.Mill.Seats
105 }
106 labels = normalizeLabels(cfg.Mill.Labels)
107 millURL = cfg.Mill.URL
108 token = cfg.Mill.SharedSecret
109 nodeID = cfg.Server.Hostname
110 }
111 if d == nil || n == nil {
112 return nil, fmt.Errorf("executor requires a database and notifier")
113 }
114 var writer artifactstore.Writer
115 if len(writers) > 0 {
116 writer = writers[0]
117 }
118 e := &Executor{
119 millURL: millURL,
120 token: token,
121 nodeID: nodeID,
122 seats: seats,
123 labels: labels,
124 engines: engines,
125 db: d,
126 n: n,
127 cfg: cfg,
128 l: l.With("component", "mill.executor"),
129 writer: writer,
130 active: make(map[string]*reservation),
131 maxOutboxBytes: 10 * 1024 * 1024,
132 }
133 if err := e.initOutbox(); err != nil {
134 return nil, fmt.Errorf("initialize executor outbox: %w", err)
135 }
136 return e, nil
137}
138
139func (e *Executor) Connect(ctx context.Context) {
140 e.lifecycleCtx = ctx
141 sub := e.n.Subscribe()
142 cursor, err := e.db.EventHighWater()
143 if err != nil {
144 e.n.Unsubscribe(sub)
145 e.l.Error("establish event cursor failed", "err", err)
146 return
147 }
148 e.drainEvents(&cursor)
149 observerCtx, stopObserver := context.WithCancel(ctx)
150 observerDone := make(chan struct{})
151 go func() {
152 defer close(observerDone)
153 e.observeLoop(observerCtx, sub, cursor)
154 }()
155 defer e.n.Unsubscribe(sub)
156
157 backoff := dialBackoffMin
158 for {
159 if ctx.Err() != nil {
160 break
161 }
162 err := e.runSession(ctx)
163 if ctx.Err() != nil {
164 break
165 }
166 e.l.Warn("mill session ended; reconnecting", "err", err, "backoff", backoff)
167 select {
168 case <-ctx.Done():
169 break
170 case <-time.After(backoff):
171 }
172 backoff = min(backoff*2, dialBackoffMax)
173 }
174
175 e.jobsWG.Wait()
176 stopObserver()
177 <-observerDone
178 e.drainEvents(&cursor)
179}
180
181func (e *Executor) runSession(ctx context.Context) error {
182 dev := e.cfg == nil || e.cfg.Server.Dev
183 if _, err := netutil.EnforceWSSURL(e.millURL, dev); err != nil {
184 return fmt.Errorf("mill url: %w", err)
185 }
186 header := http.Header{}
187 if e.token != "" {
188 header.Set("Authorization", "Bearer "+e.token)
189 }
190 conn, _, err := netutil.SSRFWebsocketDialer(dev).DialContext(ctx, e.millURL, header)
191 if err != nil {
192 return fmt.Errorf("dial mill: %w", err)
193 }
194 defer conn.Close()
195
196 sessionCtx, cancelSession := context.WithCancel(ctx)
197 defer cancelSession()
198 stopClose := context.AfterFunc(sessionCtx, func() { _ = conn.Close() })
199 defer stopClose()
200
201 stream := millproto.NewWSStream(conn)
202 enc := millproto.NewEncoder(stream)
203 dec := millproto.NewDecoder(stream)
204
205 hello := &millproto.Message{Hello: &millv1.Hello{
206 ProtocolVersion: millproto.ProtocolVersion,
207 Arch: runtime.GOARCH,
208 Labels: e.labels,
209 Epoch: e.epoch,
210 }}
211 if err := enc.Encode(hello); err != nil {
212 return fmt.Errorf("send hello: %w", err)
213 }
214
215 resumeMsg, err := dec.Decode()
216 if err != nil {
217 return fmt.Errorf("read resume: %w", err)
218 }
219 resume := resumeMsg.GetResume()
220 if resume == nil {
221 return fmt.Errorf("expected resume, got something else")
222 }
223 if resume.GetEpoch() != e.epoch {
224 return fmt.Errorf("resume epoch mismatch: got %q, want %q", resume.GetEpoch(), e.epoch)
225 }
226
227 e.connMu.Lock()
228 e.sessionCancel = cancelSession
229 e.enc = enc
230 e.connMu.Unlock()
231 defer func() {
232 e.connMu.Lock()
233 e.enc = nil
234 e.sessionCancel = nil
235 e.connMu.Unlock()
236 }()
237
238 readErr := make(chan error, 1)
239 go func() {
240 for {
241 msg, err := dec.Decode()
242 if err != nil {
243 readErr <- fmt.Errorf("read: %w", err)
244 return
245 }
246 e.dispatch(sessionCtx, msg)
247 }
248 }()
249
250 if err := e.replay(resume.GetAckSeqno()); err != nil {
251 cancelSession()
252 <-readErr
253 return fmt.Errorf("replay failed: %w", err)
254 }
255 e.pushSnapshot()
256 e.l.Info("connected to mill", "node", e.nodeID, "resumeFrom", resume.GetAckSeqno())
257
258 go e.snapshotLoop(sessionCtx, enc)
259 return <-readErr
260}
261func (e *Executor) send(msg *millproto.Message) {
262 e.connMu.Lock()
263 enc := e.enc
264 cancel := e.sessionCancel
265 e.connMu.Unlock()
266 if enc != nil {
267 e.sendMu.Lock()
268 err := enc.Encode(msg)
269 e.sendMu.Unlock()
270 if err != nil {
271 e.l.Error("send failed, ending session", "err", err)
272 if cancel != nil {
273 cancel()
274 }
275 }
276 }
277}
278
279func (e *Executor) dispatch(ctx context.Context, msg *millproto.Message) {
280 switch {
281 case msg.GetReserveSeat() != nil:
282 e.handleReserve(ctx, msg.GetReserveSeat())
283 case msg.GetCommitLease() != nil:
284 e.handleCommit(ctx, msg.GetCommitLease())
285 case msg.GetReleaseLease() != nil:
286 e.handleRelease(msg.GetReleaseLease().GetLeaseId())
287 case msg.GetCancelAttempt() != nil:
288 e.handleCancel(msg.GetCancelAttempt().GetLeaseId())
289 case msg.GetAck() != nil:
290 e.handleAck(msg.GetAck())
291 default:
292 e.l.Warn("unhandled incoming message", "type", fmt.Sprintf("%T", msg))
293 }
294}
295
296func (e *Executor) sendReject(leaseID string, reason string, class millv1.RejectClass) {
297 e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{
298 LeaseId: leaseID,
299 Accepted: false,
300 RejectReason: reason,
301 RejectClass: class,
302 }})
303}
304
305func (e *Executor) sendCommitted(leaseID string) {
306 e.send(&millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}})
307}
308
309func (e *Executor) sendCancelAck(leaseID string) {
310 e.send(&millproto.Message{CancelAck: &millv1.CancelAck{LeaseId: leaseID}})
311}
312
313func (e *Executor) releaseReservation(cleanup func()) {
314 if cleanup != nil {
315 cleanup()
316 }
317 e.pushSnapshot()
318}
319
320func (e *Executor) handleReserve(ctx context.Context, rs *millv1.ReserveSeat) {
321 reject := func(reason string, class millv1.RejectClass) {
322 e.sendReject(rs.GetLeaseId(), reason, class)
323 }
324
325 e.mu.Lock()
326 draining := e.draining
327 e.mu.Unlock()
328 if draining {
329 reject("draining", millv1.RejectClass_REJECT_CLASS_TRANSIENT)
330 return
331 }
332
333 realEngine, ok := e.engines[rs.GetTargetEngine()]
334 if !ok {
335 reject("unknown engine "+rs.GetTargetEngine(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
336 return
337 }
338 slotter, ok := realEngine.(engine.WorkflowSlotter)
339 if !ok {
340 reject("engine does not support workflow slots", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
341 return
342 }
343
344 var twf tangled.Pipeline_Workflow
345 if err := json.Unmarshal([]byte(rs.GetRawWorkflowJson()), &twf); err != nil {
346 reject("bad workflow json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
347 return
348 }
349 var tpl tangled.Pipeline
350 if err := json.Unmarshal([]byte(rs.GetRawPipelineJson()), &tpl); err != nil {
351 reject("bad pipeline json", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
352 return
353 }
354 if tpl.TriggerMetadata == nil {
355 reject("pipeline missing trigger metadata", millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
356 return
357 }
358
359 pipelineId := models.PipelineId{Knot: rs.GetKnot(), Rkey: rs.GetRkey()}
360 wid := models.WorkflowId{PipelineId: pipelineId, Name: twf.Name}
361
362 wf, err := realEngine.InitWorkflow(twf, tpl)
363 if err != nil {
364 reject("init workflow: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
365 return
366 }
367 if validator, ok := realEngine.(engine.WorkflowPlacementValidator); ok {
368 if err := validator.ValidateWorkflowPlacement(wf); err != nil {
369 reject("validate workflow placement: "+err.Error(), millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE)
370 return
371 }
372 }
373 if wf.Environment == nil {
374 wf.Environment = make(map[string]string)
375 }
376 maps.Copy(wf.Environment, models.PipelineEnvVars(tpl.TriggerMetadata, pipelineId))
377
378 slot, err := slotter.AcquireWorkflowSlot(ctx, wid, wf, engine.NoWait)
379 if err != nil {
380 class := millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE
381 if errors.Is(err, engine.ErrNoWorkflowSlots) {
382 class = millv1.RejectClass_REJECT_CLASS_TRANSIENT
383 }
384 reject(err.Error(), class)
385 return
386 }
387
388 var repoDid syntax.DID
389 if tpl.TriggerMetadata != nil && tpl.TriggerMetadata.Repo != nil && tpl.TriggerMetadata.Repo.RepoDid != nil {
390 repoDid, _ = syntax.ParseDID(*tpl.TriggerMetadata.Repo.RepoDid)
391 }
392
393 res := &reservation{
394 leaseID: rs.GetLeaseId(),
395 wid: wid,
396 realEngine: realEngine,
397 slot: slot,
398 wf: wf,
399 repoDid: repoDid,
400 }
401
402 e.snapshotMu.Lock()
403 e.mu.Lock()
404 e.active[res.leaseID] = res
405 res.ttlTimer = time.AfterFunc(ttlDuration(rs.GetTtlSeconds()), func() { e.expireReservation(res.leaseID) })
406 e.mu.Unlock()
407
408 e.send(&millproto.Message{ReserveResult: &millv1.ReserveResult{
409 LeaseId: rs.GetLeaseId(),
410 Accepted: true,
411 }})
412 e.pushSnapshotLocked()
413 e.snapshotMu.Unlock()
414}
415
416func (e *Executor) handleCommit(ctx context.Context, cl *millv1.CommitLease) {
417 e.mu.Lock()
418 res := e.active[cl.GetLeaseId()]
419 if res == nil {
420 e.mu.Unlock()
421 e.sendReject(cl.GetLeaseId(), "reservation missing or expired", millv1.RejectClass_REJECT_CLASS_TRANSIENT)
422 return
423 }
424 if res.committed {
425 e.mu.Unlock()
426 e.sendCommitted(cl.GetLeaseId())
427 return
428 }
429 res.committed = true
430 if res.ttlTimer != nil {
431 res.ttlTimer.Stop()
432 }
433
434 jobCtx, cancel := context.WithCancel(e.lifecycleCtx)
435 res.cancel = cancel
436 e.mu.Unlock()
437
438 vault := newMemVault(cl.GetSecrets())
439 re := newReservedEngine(res.realEngine, res.slot)
440 pipeline := &models.Pipeline{
441 RepoDid: res.repoDid,
442 Workflows: map[models.Engine][]models.Workflow{re: {*res.wf}},
443 TrustedSource: true,
444 }
445
446 e.startTail(res)
447
448 e.jobsWG.Add(1)
449 go func() {
450 defer e.jobsWG.Done()
451 engine.StartWorkflows(e.l, vault, e.cfg, nil, e.db, e.n, jobCtx, pipeline, res.wid.PipelineId)
452 }()
453
454 e.sendCommitted(cl.GetLeaseId())
455}
456
457func (e *Executor) handleRelease(leaseID string) {
458 cleanup, ok := e.takeUncommittedReservation(leaseID, true)
459 if !ok {
460 return
461 }
462 e.releaseReservation(cleanup)
463}
464
465func (e *Executor) handleCancel(leaseID string) {
466 e.mu.Lock()
467 res := e.active[leaseID]
468 if res == nil {
469 e.mu.Unlock()
470 if err := e.appendTerminal(leaseID, string(models.StatusKindCancelled), nil); err != nil {
471 e.l.Error("persist cancelled reservation terminal", "lease", leaseID, "err", err)
472 return
473 }
474 e.sendCancelAck(leaseID)
475 return
476 }
477 res.cancelled = true
478 cancel := res.cancel
479 committed := res.committed
480 var cleanup func()
481 if !committed {
482 cleanup = e.removeReservationLocked(res, true)
483 }
484 e.mu.Unlock()
485
486 if !committed {
487 if err := e.appendTerminal(leaseID, string(models.StatusKindCancelled), nil); err != nil {
488 e.l.Error("persist cancelled reservation terminal", "lease", leaseID, "err", err)
489 e.releaseReservation(cleanup)
490 return
491 }
492 e.sendCancelAck(leaseID)
493 e.releaseReservation(cleanup)
494 return
495 }
496
497 e.sendCancelAck(leaseID)
498 if cancel != nil {
499 cancel()
500 }
501}
502
503func (e *Executor) expireReservation(leaseID string) {
504 cleanup, ok := e.takeUncommittedReservation(leaseID, true)
505 if !ok {
506 return
507 }
508 e.releaseReservation(cleanup)
509}
510
511func (e *Executor) takeUncommittedReservation(leaseID string, releaseSlot bool) (func(), bool) {
512 e.mu.Lock()
513 defer e.mu.Unlock()
514 res := e.active[leaseID]
515 if res == nil || res.committed {
516 return nil, false
517 }
518 if res.ttlTimer != nil {
519 res.ttlTimer.Stop()
520 }
521 return e.removeReservationLocked(res, releaseSlot), true
522}
523
524func (e *Executor) removeReservationLocked(res *reservation, releaseSlot bool) func() {
525 delete(e.active, res.leaseID)
526 slot := res.slot
527 return func() {
528 if releaseSlot && slot != nil {
529 slot.Release()
530 }
531 }
532}
533
534func (e *Executor) snapshotLoop(ctx context.Context, enc *millproto.Encoder) {
535 ticker := time.NewTicker(snapshotEvery)
536 defer ticker.Stop()
537
538 for {
539 select {
540 case <-ctx.Done():
541 return
542 case <-ticker.C:
543 e.pushSnapshot()
544 }
545 }
546}
547
548func (e *Executor) pushSnapshot() {
549 e.snapshotMu.Lock()
550 defer e.snapshotMu.Unlock()
551 e.pushSnapshotLocked()
552}
553
554func (e *Executor) pushSnapshotLocked() {
555 e.mu.Lock()
556 activeLeases := make([]string, 0, len(e.active))
557 for leaseID := range e.active {
558 activeLeases = append(activeLeases, leaseID)
559 }
560 e.mu.Unlock()
561
562 avail := make(map[string]*millv1.EngineAvailability)
563 for name, eng := range e.engines {
564 a := &millv1.EngineAvailability{Available: true}
565 if getter, ok := eng.(interface{ Load() map[string]float64 }); ok {
566 a.Load = getter.Load()
567 }
568 avail[name] = a
569 }
570
571 e.nextSeqno++
572 snap := &millproto.Message{
573 NodeSnapshot: &millv1.NodeSnapshot{
574 Seqno: e.nextSeqno,
575 Engines: avail,
576 ActiveLeaseIds: activeLeases,
577 },
578 }
579 e.send(snap)
580}
581
582func (e *Executor) Drain() {
583 e.mu.Lock()
584 e.draining = true
585 e.mu.Unlock()
586 e.pushSnapshot()
587}
588
589func ttlDuration(secs uint32) time.Duration {
590 if secs == 0 {
591 return defaultReservationTTL
592 }
593 return time.Duration(secs) * time.Second
594}
595
596const defaultReservationTTL = 60 * time.Second
597
598func normalizeLabels(labels []string) []string {
599 seen := make(map[string]struct{}, len(labels))
600 out := make([]string, 0, len(labels))
601 for _, label := range labels {
602 label = strings.TrimSpace(label)
603 if label == "" {
604 continue
605 }
606 if _, ok := seen[label]; ok {
607 continue
608 }
609 seen[label] = struct{}{}
610 out = append(out, label)
611 }
612 return out
613}