Monorepo for Tangled
spindle/mill: implement mill host
Authenticates executors with the pre-shared secret and takes node
identity from the Hello. Rejects a second live session for a held
identity and drops relayed status/logs/results for leases the session
does not own, re-authoring status under the lease's workflow id.
+1664
-15
+41
-15
spindle/db/events.go
+41
-15
spindle/db/events.go
···
32
32
return d.insertEvent(event, n)
33
33
}
34
34
35
-
func (d *DB) createStatusEvent(
36
-
workflowId models.WorkflowId,
37
-
statusKind models.StatusKind,
38
-
workflowError *string,
39
-
exitCode *int64,
40
-
n *notifier.Notifier,
41
-
) error {
42
-
now := time.Now()
43
-
pipelineAtUri := workflowId.PipelineId.AtUri()
35
+
// statusEvent builds the eventstream row for a pipeline status. Created is left
36
+
// zero so insertEvent stamps the local clock.
37
+
func statusEvent(pipelineAtUri, workflow, status string, workflowError *string, exitCode *int64) (eventstream.Event, error) {
44
38
s := tangled.PipelineStatus{
45
-
CreatedAt: now.Format(time.RFC3339),
39
+
CreatedAt: time.Now().Format(time.RFC3339),
46
40
Error: workflowError,
47
41
ExitCode: exitCode,
48
-
Pipeline: string(pipelineAtUri),
49
-
Workflow: workflowId.Name,
50
-
Status: string(statusKind),
42
+
Pipeline: pipelineAtUri,
43
+
Workflow: workflow,
44
+
Status: status,
51
45
}
52
46
53
47
eventJson, err := json.Marshal(s)
54
48
if err != nil {
55
-
return err
49
+
return eventstream.Event{}, err
56
50
}
57
51
58
-
event := eventstream.Event{
52
+
return eventstream.Event{
59
53
Rkey: tid.TID(),
60
54
Nsid: tangled.PipelineStatusNSID,
61
55
EventJson: eventJson,
56
+
}, nil
57
+
}
58
+
59
+
func (d *DB) createStatusEvent(
60
+
workflowId models.WorkflowId,
61
+
statusKind models.StatusKind,
62
+
workflowError *string,
63
+
exitCode *int64,
64
+
n *notifier.Notifier,
65
+
) error {
66
+
event, err := statusEvent(string(workflowId.PipelineId.AtUri()), workflowId.Name, string(statusKind), workflowError, exitCode)
67
+
if err != nil {
68
+
return err
62
69
}
70
+
return d.insertEvent(event, n)
71
+
}
63
72
73
+
// InsertRelayedStatus re-authors a status row relayed from an executor into the
74
+
// mill's own eventstream. The mill assigns its own Created timestamp (via
75
+
// insertEvent passing Created=0) so relayed rows order correctly against the
76
+
// mill's cursor, and NotifyAll wakes the mill's /events subscribers exactly
77
+
// as a local status write would.
78
+
func (d *DB) InsertRelayedStatus(
79
+
pipelineAtUri string,
80
+
workflow string,
81
+
status string,
82
+
workflowError *string,
83
+
exitCode *int64,
84
+
n *notifier.Notifier,
85
+
) error {
86
+
event, err := statusEvent(pipelineAtUri, workflow, status, workflowError, exitCode)
87
+
if err != nil {
88
+
return err
89
+
}
64
90
return d.insertEvent(event, n)
65
91
}
66
92
+263
spindle/mill/auth_test.go
+263
spindle/mill/auth_test.go
···
1
+
package mill
2
+
3
+
import (
4
+
"context"
5
+
"io"
6
+
"log/slog"
7
+
"os"
8
+
"path/filepath"
9
+
"testing"
10
+
"time"
11
+
12
+
"tangled.org/core/notifier"
13
+
"tangled.org/core/spindle/db"
14
+
"tangled.org/core/spindle/models"
15
+
16
+
millproto "tangled.org/core/spindle/mill/proto"
17
+
millv1 "tangled.org/core/spindle/mill/proto/gen"
18
+
)
19
+
20
+
func discardLogger() *slog.Logger {
21
+
return slog.New(slog.NewTextHandler(io.Discard, nil))
22
+
}
23
+
24
+
func nopEncoder() scriptedEncoder {
25
+
return scriptedEncoder(func(*millproto.Message) error { return nil })
26
+
}
27
+
28
+
// attachSession is the anti-hijack gate: a valid token must not be able to
29
+
// displace a live executor that already holds that node identity. Only once the
30
+
// incumbent is disconnected (in its reconnect grace window) may a reconnect
31
+
// adopt the identity.
32
+
func TestAttachSessionRejectsSecondLiveSession(t *testing.T) {
33
+
l := discardLogger()
34
+
m := New(l, Config{ReconnectGrace: time.Minute})
35
+
36
+
sessionOf := func(node string) *millSession {
37
+
m.mu.Lock()
38
+
defer m.mu.Unlock()
39
+
return m.sessions[node]
40
+
}
41
+
42
+
sess1 := newSession("node-1", nopEncoder(), l)
43
+
if _, ok := m.attachSession(sess1); !ok {
44
+
t.Fatal("first attach of a node was rejected; want accept")
45
+
}
46
+
if sessionOf("node-1") != sess1 {
47
+
t.Fatal("first session was not registered as the live session")
48
+
}
49
+
50
+
// a second attach for the same node while sess1 is still live must be
51
+
// rejected and must not evict the incumbent.
52
+
sess2 := newSession("node-1", nopEncoder(), l)
53
+
if _, ok := m.attachSession(sess2); ok {
54
+
t.Fatal("second live attach for an already-live node was accepted; a valid token hijacked the executor")
55
+
}
56
+
if sessionOf("node-1") != sess1 {
57
+
t.Fatal("rejected newcomer evicted the incumbent session")
58
+
}
59
+
60
+
// once the incumbent is disconnected (reconnect grace), a reconnect adopts
61
+
// the identity.
62
+
m.detachSession(sess1)
63
+
sess3 := newSession("node-1", nopEncoder(), l)
64
+
if _, ok := m.attachSession(sess3); !ok {
65
+
t.Fatal("attach during the incumbent's reconnect grace was rejected; want adopt")
66
+
}
67
+
if sessionOf("node-1") != sess3 {
68
+
t.Fatal("adopted session was not installed as the live session")
69
+
}
70
+
}
71
+
72
+
// onAttemptResult must only wake the RunStep of a lease the reporting session
73
+
// actually owns. A terminal relayed for another node's lease, or for a lease
74
+
// that does not exist, must not be delivered.
75
+
func TestOnAttemptResultIgnoresForeignLease(t *testing.T) {
76
+
l := discardLogger()
77
+
m := New(l, Config{ReconnectGrace: time.Minute})
78
+
79
+
foreign := newLease("lease-foreign", "node-a", "dummy")
80
+
m.mu.Lock()
81
+
m.leases[foreign.id] = foreign
82
+
m.mu.Unlock()
83
+
84
+
// node-b reports a terminal for a lease owned by node-a.
85
+
sessB := newSession("node-b", nopEncoder(), l)
86
+
m.attachSession(sessB)
87
+
m.onAttemptResult(sessB, &millv1.AttemptResult{
88
+
Offset: 1,
89
+
LeaseId: foreign.id,
90
+
TerminalStatus: string(models.StatusKindSuccess),
91
+
})
92
+
93
+
if _, ok := pollTerminal(foreign); ok {
94
+
t.Fatal("attempt-result on a foreign lease delivered a terminal; an executor forged another node's job result")
95
+
}
96
+
if foreign.getState() == leaseDone {
97
+
t.Fatal("attempt-result on a foreign lease sealed the lease")
98
+
}
99
+
}
100
+
101
+
func TestOnAttemptResultIgnoresAbsentLease(t *testing.T) {
102
+
l := discardLogger()
103
+
m := New(l, Config{ReconnectGrace: time.Minute})
104
+
105
+
// a bystander lease the reporting node does own, to prove an absent-lease
106
+
// relay does not spill onto some other lease.
107
+
bystander := newLease("lease-bystander", "node-b", "dummy")
108
+
m.mu.Lock()
109
+
m.leases[bystander.id] = bystander
110
+
m.mu.Unlock()
111
+
112
+
sessB := newSession("node-b", nopEncoder(), l)
113
+
m.attachSession(sessB)
114
+
115
+
// no lease with this id exists: must be a no-op (and must not panic).
116
+
m.onAttemptResult(sessB, &millv1.AttemptResult{
117
+
Offset: 1,
118
+
LeaseId: "lease-nonexistent",
119
+
TerminalStatus: string(models.StatusKindSuccess),
120
+
})
121
+
122
+
if _, ok := pollTerminal(bystander); ok {
123
+
t.Fatal("attempt-result for an absent lease delivered a terminal to a bystander lease")
124
+
}
125
+
if bystander.getState() == leaseDone {
126
+
t.Fatal("attempt-result for an absent lease sealed a bystander lease")
127
+
}
128
+
}
129
+
130
+
// The owned-lease path proves the ignore tests above are not passing merely
131
+
// because delivery is broken: a correctly owned terminal IS delivered.
132
+
func TestOnAttemptResultDeliversOwnedLease(t *testing.T) {
133
+
l := discardLogger()
134
+
m := New(l, Config{ReconnectGrace: time.Minute})
135
+
136
+
owned := newLease("lease-owned", "node-b", "dummy")
137
+
m.mu.Lock()
138
+
m.leases[owned.id] = owned
139
+
m.mu.Unlock()
140
+
141
+
sessB := newSession("node-b", nopEncoder(), l)
142
+
m.attachSession(sessB)
143
+
m.onAttemptResult(sessB, &millv1.AttemptResult{
144
+
Offset: 1,
145
+
LeaseId: owned.id,
146
+
TerminalStatus: string(models.StatusKindSuccess),
147
+
})
148
+
149
+
res, ok := pollTerminal(owned)
150
+
if !ok {
151
+
t.Fatal("attempt-result on an owned lease was not delivered")
152
+
}
153
+
if got := res.GetTerminalStatus(); got != string(models.StatusKindSuccess) {
154
+
t.Fatalf("delivered terminal status = %q, want %q", got, string(models.StatusKindSuccess))
155
+
}
156
+
if owned.getState() != leaseDone {
157
+
t.Fatal("owned lease was not sealed after its terminal was delivered")
158
+
}
159
+
}
160
+
161
+
// onLogRelay must only write a log line for a lease the reporting session owns.
162
+
// A line relayed for another node's lease is dropped; a line for an owned lease
163
+
// is written verbatim to that lease's log file.
164
+
func TestOnLogRelayRelayOwnership(t *testing.T) {
165
+
l := discardLogger()
166
+
dir := t.TempDir()
167
+
m := New(l, Config{LogDir: dir, ReconnectGrace: time.Minute})
168
+
169
+
foreign := newLease("lease-foreign", "node-x", "dummy")
170
+
foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"}
171
+
owned := newLease("lease-owned", "node-z", "dummy")
172
+
owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"}
173
+
m.mu.Lock()
174
+
m.leases[foreign.id] = foreign
175
+
m.leases[owned.id] = owned
176
+
m.mu.Unlock()
177
+
178
+
sessY := newSession("node-y", nopEncoder(), l)
179
+
m.attachSession(sessY)
180
+
sessZ := newSession("node-z", nopEncoder(), l)
181
+
m.attachSession(sessZ)
182
+
183
+
// node-y relays a log line for a lease owned by node-x: dropped, no file.
184
+
m.onLogRelay(sessY, &millv1.LogLine{
185
+
Offset: 1,
186
+
LeaseId: foreign.id,
187
+
RawJson: []byte(`{"line":"forged"}`),
188
+
})
189
+
if _, err := os.Stat(models.LogFilePath(dir, foreign.wid)); !os.IsNotExist(err) {
190
+
t.Fatalf("log relay for a foreign lease wrote a file (stat err=%v); an executor forged another node's logs", err)
191
+
}
192
+
193
+
// node-z relays a log line for its own lease: written verbatim.
194
+
line := []byte(`{"line":"legit"}`)
195
+
m.onLogRelay(sessZ, &millv1.LogLine{
196
+
Offset: 1,
197
+
LeaseId: owned.id,
198
+
RawJson: line,
199
+
})
200
+
got, err := os.ReadFile(models.LogFilePath(dir, owned.wid))
201
+
if err != nil {
202
+
t.Fatalf("owned log relay did not write the log file: %v", err)
203
+
}
204
+
if want := string(line) + "\n"; string(got) != want {
205
+
t.Fatalf("owned log line = %q, want %q", string(got), want)
206
+
}
207
+
}
208
+
209
+
// onStatusRelay must only author a status row for a lease the reporting session
210
+
// owns. A status relayed for another node's lease is ignored; a status for an
211
+
// owned lease is re-authored under the lease's (authoritative) workflow id.
212
+
func TestOnStatusRelayRelayOwnership(t *testing.T) {
213
+
ctx := context.Background()
214
+
l := discardLogger()
215
+
216
+
bdb, err := db.Make(ctx, filepath.Join(t.TempDir(), "mill.db"))
217
+
if err != nil {
218
+
t.Fatalf("db.Make: %v", err)
219
+
}
220
+
t.Cleanup(func() { bdb.Close() })
221
+
n := notifier.New()
222
+
223
+
m := New(l, Config{ReconnectGrace: time.Minute})
224
+
m.Attach(bdb, &n)
225
+
226
+
foreign := newLease("lease-foreign", "node-x", "dummy")
227
+
foreign.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "foreign"}, Name: "build"}
228
+
owned := newLease("lease-owned", "node-z", "dummy")
229
+
owned.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "owned"}, Name: "build"}
230
+
m.mu.Lock()
231
+
m.leases[foreign.id] = foreign
232
+
m.leases[owned.id] = owned
233
+
m.mu.Unlock()
234
+
235
+
sessY := newSession("node-y", nopEncoder(), l)
236
+
m.attachSession(sessY)
237
+
sessZ := newSession("node-z", nopEncoder(), l)
238
+
m.attachSession(sessZ)
239
+
240
+
// node-y relays a status for a lease owned by node-x: no row authored.
241
+
m.onStatusRelay(sessY, &millv1.StatusEvent{
242
+
Offset: 1,
243
+
LeaseId: foreign.id,
244
+
Status: "running",
245
+
})
246
+
if _, err := bdb.GetStatus(foreign.wid); err == nil {
247
+
t.Fatal("status relay for a foreign lease authored a status row; an executor forged another pipeline's status")
248
+
}
249
+
250
+
// node-z relays a status for its own lease: authored under the lease's wid.
251
+
m.onStatusRelay(sessZ, &millv1.StatusEvent{
252
+
Offset: 1,
253
+
LeaseId: owned.id,
254
+
Status: "running",
255
+
})
256
+
st, err := bdb.GetStatus(owned.wid)
257
+
if err != nil {
258
+
t.Fatalf("owned status relay did not author a status row: %v", err)
259
+
}
260
+
if st.Status != "running" {
261
+
t.Fatalf("owned status = %q, want %q", st.Status, "running")
262
+
}
263
+
}
+93
spindle/mill/engine.go
+93
spindle/mill/engine.go
···
1
+
package mill
2
+
3
+
import (
4
+
"context"
5
+
"log/slog"
6
+
"time"
7
+
8
+
"tangled.org/core/api/tangled"
9
+
"tangled.org/core/spindle/engine"
10
+
"tangled.org/core/spindle/models"
11
+
"tangled.org/core/spindle/secrets"
12
+
)
13
+
14
+
// millWorkflowState is stashed in Workflow.Data (which is `any`, like both real
15
+
// engines use it). InitWorkflow parses nothing here: it just carries the raw
16
+
// pipeline/workflow forward so the executor can run the real InitWorkflow later.
17
+
type millWorkflowState struct {
18
+
TargetEngine string
19
+
RawWorkflow tangled.Pipeline_Workflow
20
+
RawPipeline tangled.Pipeline
21
+
Wid models.WorkflowId // stamped at placement (InitWorkflow can't see it)
22
+
Lease *RemoteLease
23
+
}
24
+
25
+
// Engine is the mill's stand-in for a real engine, registered under the real
26
+
// engine names ("microvm", "nixery"). All registered names share one Mill.
27
+
type Engine struct {
28
+
name string
29
+
mill *Mill
30
+
l *slog.Logger
31
+
}
32
+
33
+
// NewEngine returns a mill engine view for one engine name.
34
+
func NewEngine(name string, mill *Mill) *Engine {
35
+
return &Engine{name: name, mill: mill, l: mill.l.With("engine", "mill:"+name)}
36
+
}
37
+
38
+
// InitWorkflow returns a synthetic one-step workflow so processPipeline injects
39
+
// TANGLED_* env and marks pending normally. The real InitWorkflow runs later on
40
+
// the executor inside ReserveSeat (intended: it runs twice).
41
+
func (e *Engine) InitWorkflow(twf tangled.Pipeline_Workflow, tpl tangled.Pipeline) (*models.Workflow, error) {
42
+
return &models.Workflow{
43
+
Name: twf.Name,
44
+
Environment: map[string]string{},
45
+
Steps: []models.Step{remoteStep{}},
46
+
Data: &millWorkflowState{
47
+
TargetEngine: e.name,
48
+
RawWorkflow: twf,
49
+
RawPipeline: tpl,
50
+
},
51
+
}, nil
52
+
}
53
+
54
+
// WorkflowLogger gives StartWorkflows a no-op logger for the mill's synthetic
55
+
// workflow: the executor's real log lines are relayed into this wid's log file
56
+
// directly (see Loom.onLogRelay), so a local logger here would only write
57
+
// competing, step-id-colliding lines. This is engine's optional logger-provider
58
+
// seam.
59
+
func (e *Engine) WorkflowLogger(wid models.WorkflowId) models.WorkflowLogger {
60
+
return models.NullLogger{}
61
+
}
62
+
63
+
// AcquireWorkflowSlot is the placement seam: it blocks doing remote placement,
64
+
// which the user sees as "pending". mode is irrelevant here (only StartWorkflows
65
+
// calls this, with Wait); the mill always places.
66
+
func (e *Engine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, _ engine.AcquireMode) (engine.WorkflowSlot, error) {
67
+
return e.mill.place(ctx, e.name, wid, wf)
68
+
}
69
+
70
+
// SetupWorkflow is a no-op info log; the real setup happens on the executor.
71
+
func (e *Engine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, wfLogger models.WorkflowLogger) error {
72
+
e.l.Info("remote job placed, awaiting commit", "wid", wid)
73
+
return nil
74
+
}
75
+
76
+
// RunStep is commit-and-wait: hand over secrets and block on the terminal
77
+
// result relayed over the session.
78
+
func (e *Engine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, unlocked []secrets.UnlockedSecret, wfLogger models.WorkflowLogger) error {
79
+
return e.mill.commitAndWait(ctx, w, unlocked)
80
+
}
81
+
82
+
// WorkflowTimeout is deliberately generous: the executor enforces the real
83
+
// per-engine timeout and relays it. The mill only caps a hung/silent
84
+
// executor, with a true death caught separately by reconnect grace.
85
+
func (e *Engine) WorkflowTimeout() time.Duration {
86
+
return e.mill.cfg.JobTimeout
87
+
}
88
+
89
+
// DestroyWorkflow cancels a still-running attempt; no-op if already terminal.
90
+
func (e *Engine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error {
91
+
e.mill.destroy(wid)
92
+
return nil
93
+
}
+72
spindle/mill/handler.go
+72
spindle/mill/handler.go
···
1
+
package mill
2
+
3
+
import (
4
+
"net/http"
5
+
6
+
"github.com/gorilla/websocket"
7
+
8
+
millproto "tangled.org/core/spindle/mill/proto"
9
+
millv1 "tangled.org/core/spindle/mill/proto/gen"
10
+
)
11
+
12
+
var upgrader = websocket.Upgrader{
13
+
ReadBufferSize: 1024,
14
+
WriteBufferSize: 1024,
15
+
}
16
+
17
+
// HandleExecutorConn terminates one executor websocket. Auth is a pre-shared
18
+
// secret in the Authorization header, checked before the upgrade so a bad token
19
+
// never opens a socket.
20
+
func (m *Mill) HandleExecutorConn(w http.ResponseWriter, r *http.Request) {
21
+
if m.cfg.SharedSecret != "" {
22
+
if r.Header.Get("Authorization") != "Bearer "+m.cfg.SharedSecret {
23
+
http.Error(w, "unauthorized", http.StatusUnauthorized)
24
+
return
25
+
}
26
+
}
27
+
28
+
conn, err := upgrader.Upgrade(w, r, nil)
29
+
if err != nil {
30
+
m.l.Error("fleet ws upgrade failed", "err", err)
31
+
return
32
+
}
33
+
defer conn.Close()
34
+
35
+
stream := millproto.NewWSStream(conn)
36
+
enc := millproto.NewEncoder(stream)
37
+
dec := millproto.NewDecoder(stream)
38
+
39
+
// first frame must be Hello.
40
+
hello, err := dec.Decode()
41
+
if err != nil {
42
+
m.l.Error("fleet read hello failed", "err", err)
43
+
return
44
+
}
45
+
h := hello.GetHello()
46
+
if h == nil {
47
+
m.l.Error("fleet first frame was not hello")
48
+
return
49
+
}
50
+
if h.GetProtocolVersion() != millproto.ProtocolVersion {
51
+
m.l.Error("fleet protocol version mismatch", "got", h.GetProtocolVersion(), "want", millproto.ProtocolVersion)
52
+
return
53
+
}
54
+
55
+
sess := newSession(h.GetNodeId(), enc, m.l)
56
+
resume, ok := m.attachSession(sess)
57
+
if !ok {
58
+
m.l.Warn("rejecting duplicate live executor session", "node", sess.nodeID)
59
+
return
60
+
}
61
+
m.l.Info("executor connected", "node", sess.nodeID, "engines", h.GetEngines(), "arch", h.GetArch(), "resume", resume)
62
+
63
+
if err := sess.send(&millproto.Message{Resume: &millv1.Resume{AckOffset: resume}}); err != nil {
64
+
m.l.Error("fleet send resume failed", "err", err)
65
+
m.detachSession(sess)
66
+
return
67
+
}
68
+
m.sessionReady(sess)
69
+
70
+
sess.readLoop(m, dec)
71
+
m.detachSession(sess)
72
+
}
+205
spindle/mill/lease.go
+205
spindle/mill/lease.go
···
1
+
package mill
2
+
3
+
import (
4
+
"os"
5
+
"sync"
6
+
7
+
"tangled.org/core/spindle/models"
8
+
9
+
millv1 "tangled.org/core/spindle/mill/proto/gen"
10
+
)
11
+
12
+
// leaseState is the mill's view of a remote attempt. It collapses the doc's
13
+
// full state diagram to what the mill actually needs to drive cleanup and the
14
+
// commit-and-wait in RunStep.
15
+
type leaseState int32
16
+
17
+
const (
18
+
// won a bid, executor is holding a seat, not yet committed.
19
+
leaseReserved leaseState = iota
20
+
// CommitLease has been sent; the executor may already be running the job,
21
+
// but the mill has not necessarily received Committed yet.
22
+
leaseCommitting
23
+
// CommitLease sent and acked; the job is running on the executor.
24
+
leaseRunning
25
+
// a terminal result arrived (or we gave up). no further action.
26
+
leaseDone
27
+
)
28
+
29
+
type cancelAction int
30
+
31
+
const (
32
+
cancelNoop cancelAction = iota
33
+
cancelLocal
34
+
cancelRemote
35
+
)
36
+
37
+
// RemoteLease is the mill-side fencing token for one placed job.
38
+
type RemoteLease struct {
39
+
id string
40
+
nodeID string
41
+
engine string
42
+
wid models.WorkflowId // the job this lease carries; set once placed
43
+
44
+
mu sync.Mutex
45
+
state leaseState
46
+
cancel bool
47
+
reason string
48
+
terminal chan *millv1.AttemptResult // buffered(1); RunStep waits here
49
+
dead chan struct{} // closed when the executor is lost past grace
50
+
deadOnce sync.Once
51
+
52
+
// log file for relayed lines, opened lazily and held open for the lease's
53
+
// life so we don't reopen per line. Guarded by logMu, closed on cleanup.
54
+
logMu sync.Mutex
55
+
logFile *os.File
56
+
logClosed bool
57
+
}
58
+
59
+
func newLease(id, nodeID, engine string) *RemoteLease {
60
+
return &RemoteLease{
61
+
id: id,
62
+
nodeID: nodeID,
63
+
engine: engine,
64
+
state: leaseReserved,
65
+
terminal: make(chan *millv1.AttemptResult, 1),
66
+
dead: make(chan struct{}),
67
+
}
68
+
}
69
+
70
+
func (l *RemoteLease) setState(s leaseState) {
71
+
l.mu.Lock()
72
+
l.state = s
73
+
l.mu.Unlock()
74
+
}
75
+
76
+
func (l *RemoteLease) markCommitting() bool {
77
+
l.mu.Lock()
78
+
defer l.mu.Unlock()
79
+
if l.state == leaseDone {
80
+
return false
81
+
}
82
+
if l.state == leaseReserved {
83
+
l.state = leaseCommitting
84
+
}
85
+
return true
86
+
}
87
+
88
+
func (l *RemoteLease) markRunning() {
89
+
l.mu.Lock()
90
+
if l.state != leaseDone {
91
+
l.state = leaseRunning
92
+
}
93
+
l.mu.Unlock()
94
+
}
95
+
96
+
func (l *RemoteLease) getState() leaseState {
97
+
l.mu.Lock()
98
+
defer l.mu.Unlock()
99
+
return l.state
100
+
}
101
+
102
+
func (l *RemoteLease) requestCancel(reason string) cancelAction {
103
+
l.mu.Lock()
104
+
defer l.mu.Unlock()
105
+
if l.state == leaseDone {
106
+
return cancelNoop
107
+
}
108
+
l.cancel = true
109
+
l.reason = reason
110
+
if l.state == leaseReserved {
111
+
return cancelLocal
112
+
}
113
+
return cancelRemote
114
+
}
115
+
116
+
func (l *RemoteLease) cancelRequested() (bool, string) {
117
+
l.mu.Lock()
118
+
defer l.mu.Unlock()
119
+
return l.cancel, l.reason
120
+
}
121
+
122
+
func (l *RemoteLease) deliverCancelled(reason string) {
123
+
l.deliverTerminal(&millv1.AttemptResult{
124
+
LeaseId: l.id,
125
+
TerminalStatus: string(models.StatusKindCancelled),
126
+
Error: reason,
127
+
})
128
+
}
129
+
130
+
// deliverTerminal hands the terminal result to a waiting RunStep without
131
+
// blocking. Idempotent: a duplicate (e.g. replayed on reconnect) is dropped
132
+
// because the channel is buffered to one and the lease is already done.
133
+
func (l *RemoteLease) deliverTerminal(res *millv1.AttemptResult) {
134
+
l.mu.Lock()
135
+
if l.state == leaseDone {
136
+
l.mu.Unlock()
137
+
return
138
+
}
139
+
l.state = leaseDone
140
+
l.mu.Unlock()
141
+
142
+
select {
143
+
case l.terminal <- res:
144
+
default:
145
+
}
146
+
}
147
+
148
+
// markDead signals that the executor running this lease is gone for good.
149
+
func (l *RemoteLease) markDead() {
150
+
l.deadOnce.Do(func() { close(l.dead) })
151
+
}
152
+
153
+
// writeLog appends a relayed log line to the lease's log file, opening it once
154
+
// and holding the handle open for the lease's life (relayed lines for a lease
155
+
// all arrive on one session reader goroutine, so this is serialized). A write
156
+
// after closeLog is dropped so a replayed line can't resurrect the file.
157
+
func (l *RemoteLease) writeLog(path string, line []byte) error {
158
+
l.logMu.Lock()
159
+
defer l.logMu.Unlock()
160
+
if l.logClosed {
161
+
return nil
162
+
}
163
+
if l.logFile == nil {
164
+
f, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
165
+
if err != nil {
166
+
return err
167
+
}
168
+
l.logFile = f
169
+
}
170
+
_, err := l.logFile.Write(line)
171
+
return err
172
+
}
173
+
174
+
// closeLog closes the lease's log file. Idempotent.
175
+
func (l *RemoteLease) closeLog() {
176
+
l.logMu.Lock()
177
+
defer l.logMu.Unlock()
178
+
l.logClosed = true
179
+
if l.logFile != nil {
180
+
_ = l.logFile.Close()
181
+
l.logFile = nil
182
+
}
183
+
}
184
+
185
+
// remoteStep is the mill's synthetic single step. The real steps run on the
186
+
// executor; the mill never mirrors them.
187
+
type remoteStep struct{}
188
+
189
+
func (remoteStep) Name() string { return "remote execution" }
190
+
func (remoteStep) Command() string { return "" }
191
+
func (remoteStep) Kind() models.StepKind { return models.StepKindSystem }
192
+
193
+
// millSlot is what AcquireWorkflowSlot returns; its Release unwinds placement.
194
+
type millSlot struct {
195
+
fleet *Mill
196
+
lease *RemoteLease
197
+
once sync.Once
198
+
}
199
+
200
+
func (s *millSlot) Release() {
201
+
if s == nil {
202
+
return
203
+
}
204
+
s.once.Do(func() { s.fleet.releaseSlot(s) })
205
+
}
+697
spindle/mill/mill.go
+697
spindle/mill/mill.go
···
1
+
package mill
2
+
3
+
import (
4
+
"context"
5
+
"encoding/json"
6
+
"errors"
7
+
"fmt"
8
+
"log/slog"
9
+
"sort"
10
+
"sync"
11
+
"time"
12
+
13
+
"tangled.org/core/notifier"
14
+
"tangled.org/core/spindle/db"
15
+
"tangled.org/core/spindle/engine"
16
+
"tangled.org/core/spindle/models"
17
+
"tangled.org/core/spindle/secrets"
18
+
"tangled.org/core/tid"
19
+
20
+
millproto "tangled.org/core/spindle/mill/proto"
21
+
millv1 "tangled.org/core/spindle/mill/proto/gen"
22
+
)
23
+
24
+
const (
25
+
defaultReconnectGrace = 45 * time.Second
26
+
defaultJobTimeout = 24 * time.Hour
27
+
defaultBidTimeout = 5 * time.Second
28
+
defaultTopK = 3
29
+
// re-poll in case a change notification is missed.
30
+
placementPoll = 2 * time.Second
31
+
)
32
+
33
+
type Config struct {
34
+
SharedSecret string
35
+
MaxPending int
36
+
ReconnectGrace time.Duration
37
+
JobTimeout time.Duration
38
+
BidTimeout time.Duration
39
+
TopK int
40
+
LogDir string
41
+
}
42
+
43
+
type Mill struct {
44
+
l *slog.Logger
45
+
cfg Config
46
+
47
+
// re-author relayed status rows into the mill's stream so appview sees
48
+
// them unchanged.
49
+
db *db.DB
50
+
n *notifier.Notifier
51
+
52
+
mu sync.Mutex
53
+
sessions map[string]*millSession
54
+
leases map[string]*RemoteLease
55
+
nodeOffset map[string]uint64
56
+
pending int
57
+
changeCh chan struct{} // closed + replaced to wake placement waiters
58
+
59
+
leaseSeq uint64
60
+
}
61
+
62
+
func New(l *slog.Logger, cfg Config) *Mill {
63
+
if cfg.ReconnectGrace <= 0 {
64
+
cfg.ReconnectGrace = defaultReconnectGrace
65
+
}
66
+
if cfg.JobTimeout <= 0 {
67
+
cfg.JobTimeout = defaultJobTimeout
68
+
}
69
+
if cfg.BidTimeout <= 0 {
70
+
cfg.BidTimeout = defaultBidTimeout
71
+
}
72
+
if cfg.TopK <= 0 {
73
+
cfg.TopK = defaultTopK
74
+
}
75
+
return &Mill{
76
+
l: l,
77
+
cfg: cfg,
78
+
sessions: make(map[string]*millSession),
79
+
leases: make(map[string]*RemoteLease),
80
+
nodeOffset: make(map[string]uint64),
81
+
changeCh: make(chan struct{}),
82
+
}
83
+
}
84
+
85
+
func (m *Mill) Attach(d *db.DB, n *notifier.Notifier) {
86
+
m.mu.Lock()
87
+
m.db = d
88
+
m.n = n
89
+
m.mu.Unlock()
90
+
}
91
+
92
+
func (m *Mill) nextLeaseID() string {
93
+
m.mu.Lock()
94
+
m.leaseSeq++
95
+
seq := m.leaseSeq
96
+
m.mu.Unlock()
97
+
return fmt.Sprintf("%s-%d", tid.TID(), seq)
98
+
}
99
+
100
+
func (m *Mill) notifyChange() {
101
+
m.mu.Lock()
102
+
m.notifyChangeLocked()
103
+
m.mu.Unlock()
104
+
}
105
+
106
+
func (m *Mill) notifyChangeLocked() {
107
+
close(m.changeCh)
108
+
m.changeCh = make(chan struct{})
109
+
}
110
+
111
+
func (m *Mill) currentChangeCh() <-chan struct{} {
112
+
m.mu.Lock()
113
+
defer m.mu.Unlock()
114
+
return m.changeCh
115
+
}
116
+
117
+
// --- session lifecycle -----------------------------------------------------
118
+
119
+
func (m *Mill) attachSession(sess *millSession) (uint64, bool) {
120
+
m.mu.Lock()
121
+
defer m.mu.Unlock()
122
+
123
+
if old := m.sessions[sess.nodeID]; old != nil {
124
+
if !old.disconnected {
125
+
// reject a second live session for the same identity so a valid token
126
+
// can't hijack an in-flight executor.
127
+
return 0, false
128
+
}
129
+
// the old session is in its reconnect grace window: adopt its leases.
130
+
if old.graceTimer != nil {
131
+
old.graceTimer.Stop()
132
+
}
133
+
old.close()
134
+
m.l.Info("executor reconnected", "node", sess.nodeID)
135
+
}
136
+
m.sessions[sess.nodeID] = sess
137
+
return m.nodeOffset[sess.nodeID], true
138
+
}
139
+
140
+
func (m *Mill) detachSession(sess *millSession) {
141
+
m.mu.Lock()
142
+
if m.sessions[sess.nodeID] != sess {
143
+
// already replaced by a reconnect; nothing to do.
144
+
m.mu.Unlock()
145
+
sess.close()
146
+
return
147
+
}
148
+
sess.disconnected = true
149
+
sess.graceTimer = time.AfterFunc(m.cfg.ReconnectGrace, func() { m.onGraceExpired(sess) })
150
+
m.mu.Unlock()
151
+
152
+
sess.close()
153
+
m.l.Warn("executor session lost; entering reconnect grace", "node", sess.nodeID, "grace", m.cfg.ReconnectGrace)
154
+
m.notifyChange()
155
+
}
156
+
157
+
func (m *Mill) sessionReady(sess *millSession) {
158
+
for _, lease := range m.cancelledLeasesForNode(sess.nodeID) {
159
+
_, reason := lease.cancelRequested()
160
+
m.sendCancel(sess, lease, reason)
161
+
}
162
+
m.notifyChange()
163
+
}
164
+
165
+
func (m *Mill) cancelledLeasesForNode(nodeID string) []*RemoteLease {
166
+
m.mu.Lock()
167
+
var candidates []*RemoteLease
168
+
for _, lease := range m.leases {
169
+
if lease.nodeID == nodeID {
170
+
candidates = append(candidates, lease)
171
+
}
172
+
}
173
+
m.mu.Unlock()
174
+
175
+
var leases []*RemoteLease
176
+
for _, lease := range candidates {
177
+
if cancelled, _ := lease.cancelRequested(); cancelled && lease.getState() != leaseDone {
178
+
leases = append(leases, lease)
179
+
}
180
+
}
181
+
return leases
182
+
}
183
+
184
+
func (m *Mill) onGraceExpired(sess *millSession) {
185
+
m.mu.Lock()
186
+
if m.sessions[sess.nodeID] != sess || !sess.disconnected {
187
+
// reconnected in the meantime.
188
+
m.mu.Unlock()
189
+
return
190
+
}
191
+
var dead []*RemoteLease
192
+
for _, lease := range m.leases {
193
+
if lease.nodeID == sess.nodeID {
194
+
dead = append(dead, lease)
195
+
}
196
+
}
197
+
delete(m.sessions, sess.nodeID)
198
+
delete(m.nodeOffset, sess.nodeID)
199
+
m.mu.Unlock()
200
+
201
+
m.l.Warn("executor declared dead; failing its in-flight jobs", "node", sess.nodeID, "jobs", len(dead))
202
+
for _, lease := range dead {
203
+
if cancelled, reason := lease.cancelRequested(); cancelled {
204
+
lease.deliverCancelled(reason)
205
+
} else {
206
+
lease.markDead()
207
+
}
208
+
}
209
+
m.notifyChange()
210
+
}
211
+
212
+
// --- placement (the WorkflowSlotter seam) ---------------------------------
213
+
214
+
func (m *Mill) place(ctx context.Context, engineName string, wid models.WorkflowId, wf *models.Workflow) (engine.WorkflowSlot, error) {
215
+
m.mu.Lock()
216
+
if m.cfg.MaxPending > 0 && m.pending >= m.cfg.MaxPending {
217
+
max := m.cfg.MaxPending
218
+
cur := m.pending
219
+
m.mu.Unlock()
220
+
return nil, fmt.Errorf("%w: mill has %d pending jobs (max %d)", engine.ErrNoWorkflowSlots, cur, max)
221
+
}
222
+
m.pending++
223
+
m.mu.Unlock()
224
+
defer func() {
225
+
m.mu.Lock()
226
+
m.pending--
227
+
m.mu.Unlock()
228
+
}()
229
+
230
+
// stamp wid early so the executor reconstructs logs/status under the same
231
+
// pipeline id the mill used for "pending".
232
+
if st, ok := wf.Data.(*millWorkflowState); ok && st != nil {
233
+
st.Wid = wid
234
+
}
235
+
236
+
for {
237
+
if err := ctx.Err(); err != nil {
238
+
return nil, err
239
+
}
240
+
241
+
lease, err := m.bid(ctx, engineName, wf)
242
+
if err != nil {
243
+
return nil, err
244
+
}
245
+
if lease != nil {
246
+
lease.wid = wid
247
+
m.mu.Lock()
248
+
m.leases[lease.id] = lease
249
+
if st, ok := wf.Data.(*millWorkflowState); ok && st != nil {
250
+
st.Lease = lease
251
+
}
252
+
m.mu.Unlock()
253
+
return &millSlot{fleet: m, lease: lease}, nil
254
+
}
255
+
256
+
// nobody to place on right now: wait for a change, ctx, or a poll tick.
257
+
m.mu.Lock()
258
+
ch := m.changeCh
259
+
m.mu.Unlock()
260
+
select {
261
+
case <-ctx.Done():
262
+
return nil, ctx.Err()
263
+
case <-ch:
264
+
case <-time.After(placementPoll):
265
+
}
266
+
}
267
+
}
268
+
269
+
func (m *Mill) bid(ctx context.Context, engineName string, wf *models.Workflow) (*RemoteLease, error) {
270
+
rawPipeline, rawWorkflow, err := marshalJob(wf)
271
+
if err != nil {
272
+
return nil, err
273
+
}
274
+
275
+
candidates := m.rankCandidates(engineName)
276
+
if len(candidates) == 0 {
277
+
return nil, nil
278
+
}
279
+
if len(candidates) > m.cfg.TopK {
280
+
candidates = candidates[:m.cfg.TopK]
281
+
}
282
+
283
+
bidCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout)
284
+
defer cancel()
285
+
286
+
type bidResult struct {
287
+
sess *millSession
288
+
lease *RemoteLease
289
+
score float64
290
+
}
291
+
results := make(chan bidResult, len(candidates))
292
+
var wg sync.WaitGroup
293
+
for _, sess := range candidates {
294
+
wg.Add(1)
295
+
go func(sess *millSession) {
296
+
defer wg.Done()
297
+
leaseID := m.nextLeaseID()
298
+
lease := newLease(leaseID, sess.nodeID, engineName)
299
+
wid := m.jobWid(wf)
300
+
msg := &millproto.Message{ReserveSeat: &millv1.ReserveSeat{
301
+
LeaseId: leaseID,
302
+
TargetEngine: engineName,
303
+
RawPipelineJson: rawPipeline,
304
+
RawWorkflowJson: rawWorkflow,
305
+
Knot: wid.Knot,
306
+
Rkey: wid.Rkey,
307
+
TtlSeconds: uint32(m.cfg.ReconnectGrace / time.Second),
308
+
}}
309
+
resp, err := sess.request(bidCtx, leaseID, msg)
310
+
if err != nil {
311
+
return
312
+
}
313
+
rr := resp.GetReserveResult()
314
+
if rr == nil || !rr.GetAccepted() {
315
+
return
316
+
}
317
+
results <- bidResult{sess: sess, lease: lease, score: rr.GetScore()}
318
+
}(sess)
319
+
}
320
+
wg.Wait()
321
+
close(results)
322
+
323
+
var winner *bidResult
324
+
var losers []*RemoteLease
325
+
for r := range results {
326
+
if winner == nil || r.score > winner.score {
327
+
if winner != nil {
328
+
losers = append(losers, winner.lease)
329
+
}
330
+
winner = &r
331
+
} else {
332
+
losers = append(losers, r.lease)
333
+
}
334
+
}
335
+
336
+
// release the also-rans so they free their held seats immediately.
337
+
for _, l := range losers {
338
+
m.releaseRemote(l)
339
+
}
340
+
341
+
if winner == nil {
342
+
return nil, nil
343
+
}
344
+
return winner.lease, nil
345
+
}
346
+
347
+
func (m *Mill) rankCandidates(engineName string) []*millSession {
348
+
m.mu.Lock()
349
+
defer m.mu.Unlock()
350
+
351
+
type ranked struct {
352
+
sess *millSession
353
+
seats uint32
354
+
}
355
+
var rs []ranked
356
+
for _, s := range m.sessions {
357
+
if s.disconnected {
358
+
continue
359
+
}
360
+
if s.snapshot == nil {
361
+
continue
362
+
}
363
+
es, ok := s.snapshot.GetEngines()[engineName]
364
+
if !ok || es.GetFreeSeats() == 0 {
365
+
continue
366
+
}
367
+
rs = append(rs, ranked{sess: s, seats: es.GetFreeSeats()})
368
+
}
369
+
sort.SliceStable(rs, func(i, j int) bool { return rs[i].seats > rs[j].seats })
370
+
371
+
out := make([]*millSession, len(rs))
372
+
for i := range rs {
373
+
out[i] = rs[i].sess
374
+
}
375
+
return out
376
+
}
377
+
378
+
// --- commit + wait (RunStep) ----------------------------------------------
379
+
380
+
func (m *Mill) commitAndWait(ctx context.Context, wf *models.Workflow, unlocked []secrets.UnlockedSecret) error {
381
+
st, ok := wf.Data.(*millWorkflowState)
382
+
if !ok || st == nil || st.Lease == nil {
383
+
return fmt.Errorf("mill workflow state missing lease")
384
+
}
385
+
lease := st.Lease
386
+
387
+
pbSecrets := make([]*millv1.Secret, len(unlocked))
388
+
for i, s := range unlocked {
389
+
pbSecrets[i] = &millv1.Secret{Key: s.Key, Value: s.Value}
390
+
}
391
+
392
+
commit := &millproto.Message{CommitLease: &millv1.CommitLease{
393
+
LeaseId: lease.id,
394
+
Secrets: pbSecrets,
395
+
}}
396
+
397
+
for {
398
+
if res, ok := pollTerminal(lease); ok {
399
+
return terminalError(res.GetTerminalStatus())
400
+
}
401
+
if !lease.markCommitting() {
402
+
return engine.ErrWorkflowFailed
403
+
}
404
+
405
+
sess := m.sessionForNode(lease.nodeID)
406
+
if sess == nil {
407
+
if done, err := m.waitCommitRetry(ctx, lease); done || err != nil {
408
+
return err
409
+
}
410
+
continue
411
+
}
412
+
413
+
reqCtx, cancel := context.WithTimeout(ctx, m.cfg.BidTimeout)
414
+
resp, err := sess.request(reqCtx, lease.id, commit)
415
+
cancel()
416
+
if err != nil {
417
+
switch {
418
+
case errors.Is(err, errSessionClosed):
419
+
if done, err := m.waitCommitRetry(ctx, lease); done || err != nil {
420
+
return err
421
+
}
422
+
continue
423
+
case errors.Is(err, context.DeadlineExceeded) && ctx.Err() == nil:
424
+
continue
425
+
case errors.Is(err, context.DeadlineExceeded):
426
+
return engine.ErrTimedOut
427
+
case errors.Is(err, context.Canceled):
428
+
return err
429
+
default:
430
+
m.l.Warn("commit lease send failed; waiting for reconnect", "lease", lease.id, "node", lease.nodeID, "err", err)
431
+
if done, err := m.waitCommitRetry(ctx, lease); done || err != nil {
432
+
return err
433
+
}
434
+
continue
435
+
}
436
+
}
437
+
if resp.GetCommitted() == nil {
438
+
return engine.ErrWorkflowFailed
439
+
}
440
+
lease.markRunning()
441
+
if cancelled, reason := lease.cancelRequested(); cancelled {
442
+
m.sendCancel(sess, lease, reason)
443
+
}
444
+
break
445
+
}
446
+
447
+
select {
448
+
case res := <-lease.terminal:
449
+
return terminalError(res.GetTerminalStatus())
450
+
case <-lease.dead:
451
+
m.l.Warn("executor lost while running job", "lease", lease.id, "node", lease.nodeID)
452
+
return engine.ErrWorkflowFailed
453
+
case <-ctx.Done():
454
+
if ctx.Err() == context.DeadlineExceeded {
455
+
return engine.ErrTimedOut
456
+
}
457
+
return ctx.Err()
458
+
}
459
+
}
460
+
461
+
func (m *Mill) waitCommitRetry(ctx context.Context, lease *RemoteLease) (bool, error) {
462
+
res, err := m.waitLeaseChange(ctx, lease)
463
+
if err != nil {
464
+
return true, err
465
+
}
466
+
if res != nil {
467
+
return true, terminalError(res.GetTerminalStatus())
468
+
}
469
+
return false, nil
470
+
}
471
+
472
+
func terminalError(status string) error {
473
+
switch models.StatusKind(status) {
474
+
case models.StatusKindSuccess:
475
+
return nil
476
+
case models.StatusKindTimeout:
477
+
return engine.ErrTimedOut
478
+
case models.StatusKindCancelled:
479
+
return engine.ErrCancelled
480
+
default:
481
+
return engine.ErrWorkflowFailed
482
+
}
483
+
}
484
+
485
+
func pollTerminal(lease *RemoteLease) (*millv1.AttemptResult, bool) {
486
+
select {
487
+
case res := <-lease.terminal:
488
+
return res, true
489
+
default:
490
+
return nil, false
491
+
}
492
+
}
493
+
494
+
func (m *Mill) waitLeaseChange(ctx context.Context, lease *RemoteLease) (*millv1.AttemptResult, error) {
495
+
ch := m.currentChangeCh()
496
+
select {
497
+
case res := <-lease.terminal:
498
+
return res, nil
499
+
case <-lease.dead:
500
+
return nil, engine.ErrWorkflowFailed
501
+
case <-ctx.Done():
502
+
if ctx.Err() == context.DeadlineExceeded {
503
+
return nil, engine.ErrTimedOut
504
+
}
505
+
return nil, ctx.Err()
506
+
case <-ch:
507
+
return nil, nil
508
+
case <-time.After(placementPoll):
509
+
return nil, nil
510
+
}
511
+
}
512
+
513
+
// --- cleanup ---------------------------------------------------------------
514
+
515
+
func (m *Mill) destroy(wid models.WorkflowId) {
516
+
m.mu.Lock()
517
+
var lease *RemoteLease
518
+
for _, l := range m.leases {
519
+
if l.wid == wid {
520
+
lease = l
521
+
break
522
+
}
523
+
}
524
+
m.mu.Unlock()
525
+
if lease == nil {
526
+
return
527
+
}
528
+
reason := "workflow destroyed"
529
+
switch lease.requestCancel(reason) {
530
+
case cancelLocal:
531
+
if sess := m.sessionForNode(lease.nodeID); sess != nil {
532
+
_ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}})
533
+
}
534
+
lease.deliverCancelled(reason)
535
+
case cancelRemote:
536
+
if sess := m.sessionForNode(lease.nodeID); sess != nil {
537
+
m.sendCancel(sess, lease, reason)
538
+
}
539
+
}
540
+
}
541
+
542
+
func (m *Mill) releaseSlot(s *millSlot) {
543
+
lease := s.lease
544
+
if lease.getState() == leaseReserved {
545
+
m.releaseRemote(lease)
546
+
}
547
+
lease.closeLog()
548
+
549
+
m.mu.Lock()
550
+
delete(m.leases, lease.id)
551
+
m.mu.Unlock()
552
+
553
+
m.notifyChange()
554
+
}
555
+
556
+
func (m *Mill) releaseRemote(lease *RemoteLease) {
557
+
lease.setState(leaseDone)
558
+
if sess := m.sessionForNode(lease.nodeID); sess != nil {
559
+
_ = sess.send(&millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: lease.id}})
560
+
}
561
+
}
562
+
563
+
func (m *Mill) sendCancel(sess *millSession, lease *RemoteLease, reason string) {
564
+
_ = sess.send(&millproto.Message{CancelAttempt: &millv1.CancelAttempt{
565
+
LeaseId: lease.id,
566
+
Reason: reason,
567
+
}})
568
+
}
569
+
570
+
func (m *Mill) sessionForNode(nodeID string) *millSession {
571
+
m.mu.Lock()
572
+
defer m.mu.Unlock()
573
+
s := m.sessions[nodeID]
574
+
if s == nil || s.disconnected {
575
+
return nil
576
+
}
577
+
return s
578
+
}
579
+
580
+
// --- relay handlers --------------------------------------------------------
581
+
582
+
func (m *Mill) shouldAcceptOffset(sess *millSession, offset uint64) bool {
583
+
m.mu.Lock()
584
+
defer m.mu.Unlock()
585
+
current := m.nodeOffset[sess.nodeID]
586
+
return offset == current+1
587
+
}
588
+
589
+
func (m *Mill) ackOffset(sess *millSession, offset uint64) {
590
+
m.mu.Lock()
591
+
if offset > m.nodeOffset[sess.nodeID] {
592
+
m.nodeOffset[sess.nodeID] = offset
593
+
}
594
+
m.mu.Unlock()
595
+
596
+
_ = sess.send(&millproto.Message{Ack: &millv1.Ack{UpToOffset: offset}})
597
+
}
598
+
599
+
func (m *Mill) processRelay(sess *millSession, offset uint64, kind string, apply func() error) {
600
+
if !m.shouldAcceptOffset(sess, offset) {
601
+
return
602
+
}
603
+
// ack only after the side effect lands, or a transient write error makes
604
+
// the message unreplayable.
605
+
if err := apply(); err != nil {
606
+
m.l.Error("process relayed message failed", "kind", kind, "node", sess.nodeID, "offset", offset, "err", err)
607
+
return
608
+
}
609
+
m.ackOffset(sess, offset)
610
+
}
611
+
612
+
func (m *Mill) onSnapshot(sess *millSession, snap *millv1.NodeSnapshot) {
613
+
m.mu.Lock()
614
+
sess.snapshot = snap
615
+
m.mu.Unlock()
616
+
m.notifyChange()
617
+
}
618
+
619
+
func (m *Mill) onStatusRelay(sess *millSession, ev *millv1.StatusEvent) {
620
+
m.processRelay(sess, ev.GetOffset(), "status", func() error {
621
+
if m.db == nil {
622
+
return fmt.Errorf("mill db not attached")
623
+
}
624
+
m.mu.Lock()
625
+
lease := m.leases[ev.GetLeaseId()]
626
+
m.mu.Unlock()
627
+
if lease == nil || lease.nodeID != sess.nodeID {
628
+
// not a lease this executor holds; don't let it author status rows for
629
+
// arbitrary pipelines.
630
+
return nil
631
+
}
632
+
var errMsg *string
633
+
if e := ev.GetError(); e != "" {
634
+
errMsg = &e
635
+
}
636
+
var exit *int64
637
+
if c := ev.GetExitCode(); c != 0 {
638
+
exit = &c
639
+
}
640
+
// authoritative pipeline/workflow come from the lease, not whatever the
641
+
// executor sent.
642
+
return m.db.InsertRelayedStatus(string(lease.wid.PipelineId.AtUri()), lease.wid.Name, ev.GetStatus(), errMsg, exit, m.n)
643
+
})
644
+
}
645
+
646
+
func (m *Mill) onLogRelay(sess *millSession, ll *millv1.LogLine) {
647
+
m.processRelay(sess, ll.GetOffset(), "log", func() error {
648
+
m.mu.Lock()
649
+
lease := m.leases[ll.GetLeaseId()]
650
+
m.mu.Unlock()
651
+
if lease == nil || lease.nodeID != sess.nodeID {
652
+
return nil
653
+
}
654
+
path := models.LogFilePath(m.cfg.LogDir, lease.wid)
655
+
line := append(ll.GetRawJson(), '\n')
656
+
return lease.writeLog(path, line)
657
+
})
658
+
}
659
+
660
+
func (m *Mill) onAttemptResult(sess *millSession, ar *millv1.AttemptResult) {
661
+
m.processRelay(sess, ar.GetOffset(), "attempt-result", func() error {
662
+
m.mu.Lock()
663
+
lease := m.leases[ar.GetLeaseId()]
664
+
m.mu.Unlock()
665
+
if lease == nil || lease.nodeID != sess.nodeID {
666
+
return nil
667
+
}
668
+
lease.deliverTerminal(ar)
669
+
return nil
670
+
})
671
+
}
672
+
673
+
// --- helpers ---------------------------------------------------------------
674
+
675
+
func marshalJob(wf *models.Workflow) (pipeline string, workflow string, err error) {
676
+
st, ok := wf.Data.(*millWorkflowState)
677
+
if !ok || st == nil {
678
+
return "", "", fmt.Errorf("mill workflow state missing")
679
+
}
680
+
p, err := json.Marshal(st.RawPipeline)
681
+
if err != nil {
682
+
return "", "", fmt.Errorf("marshal pipeline: %w", err)
683
+
}
684
+
w, err := json.Marshal(st.RawWorkflow)
685
+
if err != nil {
686
+
return "", "", fmt.Errorf("marshal workflow: %w", err)
687
+
}
688
+
return string(p), string(w), nil
689
+
}
690
+
691
+
func (m *Mill) jobWid(wf *models.Workflow) models.WorkflowId {
692
+
st, ok := wf.Data.(*millWorkflowState)
693
+
if !ok || st == nil {
694
+
return models.WorkflowId{}
695
+
}
696
+
return st.Wid
697
+
}
+154
spindle/mill/mill_test.go
+154
spindle/mill/mill_test.go
···
1
+
package mill
2
+
3
+
import (
4
+
"context"
5
+
"errors"
6
+
"io"
7
+
"log/slog"
8
+
"testing"
9
+
"time"
10
+
11
+
"tangled.org/core/api/tangled"
12
+
"tangled.org/core/spindle/engine"
13
+
"tangled.org/core/spindle/models"
14
+
15
+
millproto "tangled.org/core/spindle/mill/proto"
16
+
millv1 "tangled.org/core/spindle/mill/proto/gen"
17
+
)
18
+
19
+
type scriptedEncoder func(*millproto.Message) error
20
+
21
+
func (e scriptedEncoder) Encode(msg *millproto.Message) error { return e(msg) }
22
+
23
+
func testWorkflow(name string) *models.Workflow {
24
+
return &models.Workflow{
25
+
Name: name,
26
+
Environment: map[string]string{},
27
+
Steps: []models.Step{remoteStep{}},
28
+
Data: &millWorkflowState{
29
+
TargetEngine: "dummy",
30
+
RawWorkflow: tangled.Pipeline_Workflow{Name: name},
31
+
RawPipeline: tangled.Pipeline{},
32
+
},
33
+
}
34
+
}
35
+
36
+
func TestCommitRetriesAfterSessionCloseBeforeCommitted(t *testing.T) {
37
+
l := slog.New(slog.NewTextHandler(io.Discard, nil))
38
+
m := New(l, Config{BidTimeout: 25 * time.Millisecond, ReconnectGrace: time.Second})
39
+
wf := testWorkflow("build")
40
+
wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
41
+
lease := newLease("lease-1", "node-1", "dummy")
42
+
lease.wid = wid
43
+
wf.Data.(*millWorkflowState).Lease = lease
44
+
45
+
m.mu.Lock()
46
+
m.leases[lease.id] = lease
47
+
m.mu.Unlock()
48
+
49
+
var sess1 *millSession
50
+
firstCommit := make(chan struct{})
51
+
sess1 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error {
52
+
if msg.GetCommitLease() != nil {
53
+
close(firstCommit)
54
+
m.detachSession(sess1)
55
+
}
56
+
return nil
57
+
}), l)
58
+
m.attachSession(sess1)
59
+
60
+
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
61
+
defer cancel()
62
+
done := make(chan error, 1)
63
+
go func() { done <- m.commitAndWait(ctx, wf, nil) }()
64
+
65
+
select {
66
+
case <-firstCommit:
67
+
case <-ctx.Done():
68
+
t.Fatal("first commit was not sent")
69
+
}
70
+
71
+
var sess2 *millSession
72
+
sess2 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error {
73
+
if msg.GetCommitLease() == nil {
74
+
return nil
75
+
}
76
+
leaseID := msg.GetCommitLease().GetLeaseId()
77
+
sess2.deliver(leaseID, &millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}})
78
+
m.onAttemptResult(sess2, &millv1.AttemptResult{
79
+
Offset: 1,
80
+
LeaseId: leaseID,
81
+
TerminalStatus: string(models.StatusKindSuccess),
82
+
})
83
+
return nil
84
+
}), l)
85
+
m.attachSession(sess2)
86
+
m.sessionReady(sess2)
87
+
88
+
select {
89
+
case err := <-done:
90
+
if err != nil {
91
+
t.Fatalf("commitAndWait() error = %v, want success after reconnect", err)
92
+
}
93
+
case <-ctx.Done():
94
+
t.Fatal("commitAndWait() did not finish after reconnect")
95
+
}
96
+
}
97
+
98
+
func TestDestroyRunningLeaseDoesNotDropCancelledTerminal(t *testing.T) {
99
+
l := slog.New(slog.NewTextHandler(io.Discard, nil))
100
+
m := New(l, Config{})
101
+
wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
102
+
lease := newLease("lease-1", "node-1", "dummy")
103
+
lease.wid = wid
104
+
lease.setState(leaseRunning)
105
+
106
+
m.mu.Lock()
107
+
m.leases[lease.id] = lease
108
+
m.mu.Unlock()
109
+
110
+
m.destroy(wid)
111
+
if lease.getState() == leaseDone {
112
+
t.Fatal("destroy sealed the lease before the terminal result")
113
+
}
114
+
115
+
lease.deliverTerminal(&millv1.AttemptResult{
116
+
LeaseId: lease.id,
117
+
TerminalStatus: string(models.StatusKindCancelled),
118
+
})
119
+
res := <-lease.terminal
120
+
if err := terminalError(res.GetTerminalStatus()); !errors.Is(err, engine.ErrCancelled) {
121
+
t.Fatalf("terminalError() = %v, want ErrCancelled", err)
122
+
}
123
+
}
124
+
125
+
func TestPlaceBlocksWhenNoCapacity(t *testing.T) {
126
+
l := slog.New(slog.NewTextHandler(io.Discard, nil))
127
+
m := New(l, Config{})
128
+
wf := testWorkflow("build")
129
+
wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
130
+
131
+
// no executors at all: place must block until ctx expires (user sees pending).
132
+
ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
133
+
defer cancel()
134
+
135
+
_, err := m.place(ctx, "dummy", wid, wf)
136
+
if err != context.DeadlineExceeded {
137
+
t.Fatalf("place() error = %v, want DeadlineExceeded", err)
138
+
}
139
+
}
140
+
141
+
func TestMaxPendingRejects(t *testing.T) {
142
+
l := slog.New(slog.NewTextHandler(io.Discard, nil))
143
+
m := New(l, Config{MaxPending: 1})
144
+
145
+
m.mu.Lock()
146
+
m.pending = 1
147
+
m.mu.Unlock()
148
+
149
+
wf2 := testWorkflow("b")
150
+
_, err := m.place(context.Background(), "dummy", models.WorkflowId{Name: "b"}, wf2)
151
+
if err == nil {
152
+
t.Fatal("place() past maxPending should error")
153
+
}
154
+
}
+139
spindle/mill/session.go
+139
spindle/mill/session.go
···
1
+
package mill
2
+
3
+
import (
4
+
"context"
5
+
"errors"
6
+
"log/slog"
7
+
"sync"
8
+
"time"
9
+
10
+
millproto "tangled.org/core/spindle/mill/proto"
11
+
millv1 "tangled.org/core/spindle/mill/proto/gen"
12
+
)
13
+
14
+
var errSessionClosed = errors.New("mill: executor session closed")
15
+
16
+
// millSession is one live websocket to an executor. Unlike the microVM
17
+
// AgentSession it is multiplexed: many leases plus async status/log relay share
18
+
// it, so a single reader goroutine demuxes by message type and correlates
19
+
// request/response by lease id. We never hold a lock across a decode.
20
+
type millSession struct {
21
+
nodeID string
22
+
enc messageEncoder
23
+
l *slog.Logger
24
+
25
+
// snapshot, disconnected and graceTimer are guarded by Mill.mu, not the
26
+
// session mutex below (the fleet ranks across sessions under its own lock).
27
+
snapshot *millv1.NodeSnapshot
28
+
disconnected bool
29
+
graceTimer *time.Timer
30
+
31
+
mu sync.Mutex
32
+
pending map[string]chan *millproto.Message // lease_id -> response waiter
33
+
34
+
closeOnce sync.Once
35
+
closed chan struct{}
36
+
}
37
+
38
+
type messageEncoder interface {
39
+
Encode(*millproto.Message) error
40
+
}
41
+
42
+
func newSession(nodeID string, enc messageEncoder, l *slog.Logger) *millSession {
43
+
return &millSession{
44
+
nodeID: nodeID,
45
+
enc: enc,
46
+
l: l,
47
+
pending: make(map[string]chan *millproto.Message),
48
+
closed: make(chan struct{}),
49
+
}
50
+
}
51
+
52
+
func (s *millSession) send(msg *millproto.Message) error {
53
+
return s.enc.Encode(msg)
54
+
}
55
+
56
+
func (s *millSession) close() {
57
+
s.closeOnce.Do(func() { close(s.closed) })
58
+
}
59
+
60
+
// await registers a one-shot waiter for the next response addressed to leaseID,
61
+
// returning the channel and a cancel that unregisters it.
62
+
func (s *millSession) await(leaseID string) (<-chan *millproto.Message, func()) {
63
+
ch := make(chan *millproto.Message, 1)
64
+
s.mu.Lock()
65
+
s.pending[leaseID] = ch
66
+
s.mu.Unlock()
67
+
return ch, func() {
68
+
s.mu.Lock()
69
+
if s.pending[leaseID] == ch {
70
+
delete(s.pending, leaseID)
71
+
}
72
+
s.mu.Unlock()
73
+
}
74
+
}
75
+
76
+
// deliver routes a correlated response to a waiter, if any is registered.
77
+
func (s *millSession) deliver(leaseID string, msg *millproto.Message) {
78
+
s.mu.Lock()
79
+
ch := s.pending[leaseID]
80
+
delete(s.pending, leaseID)
81
+
s.mu.Unlock()
82
+
if ch != nil {
83
+
select {
84
+
case ch <- msg:
85
+
default:
86
+
}
87
+
}
88
+
}
89
+
90
+
// request sends msg and waits for the correlated response addressed to leaseID,
91
+
// honouring ctx and session close.
92
+
func (s *millSession) request(ctx context.Context, leaseID string, msg *millproto.Message) (*millproto.Message, error) {
93
+
ch, cancel := s.await(leaseID)
94
+
defer cancel()
95
+
96
+
if err := s.send(msg); err != nil {
97
+
return nil, err
98
+
}
99
+
100
+
select {
101
+
case resp := <-ch:
102
+
return resp, nil
103
+
case <-ctx.Done():
104
+
return nil, ctx.Err()
105
+
case <-s.closed:
106
+
return nil, errSessionClosed
107
+
}
108
+
}
109
+
110
+
// readLoop demuxes incoming frames until the decoder errors (connection gone).
111
+
func (s *millSession) readLoop(m *Mill, dec *millproto.Decoder) {
112
+
for {
113
+
msg, err := dec.Decode()
114
+
if err != nil {
115
+
s.l.Debug("session read ended", "node", s.nodeID, "err", err)
116
+
return
117
+
}
118
+
s.dispatch(m, msg)
119
+
}
120
+
}
121
+
122
+
func (s *millSession) dispatch(m *Mill, msg *millproto.Message) {
123
+
switch {
124
+
case msg.GetNodeSnapshot() != nil:
125
+
m.onSnapshot(s, msg.GetNodeSnapshot())
126
+
case msg.GetReserveResult() != nil:
127
+
s.deliver(msg.GetReserveResult().GetLeaseId(), msg)
128
+
case msg.GetCommitted() != nil:
129
+
s.deliver(msg.GetCommitted().GetLeaseId(), msg)
130
+
case msg.GetStatusEvent() != nil:
131
+
m.onStatusRelay(s, msg.GetStatusEvent())
132
+
case msg.GetLogLine() != nil:
133
+
m.onLogRelay(s, msg.GetLogLine())
134
+
case msg.GetAttemptResult() != nil:
135
+
m.onAttemptResult(s, msg.GetAttemptResult())
136
+
default:
137
+
s.l.Warn("session received unexpected message", "node", s.nodeID)
138
+
}
139
+
}