package executor
import (
"context"
"encoding/json"
"errors"
"fmt"
"github.com/bluesky-social/indigo/atproto/syntax"
"github.com/gorilla/websocket"
"google.golang.org/protobuf/proto"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"slices"
"strings"
"testing"
"time"
"tangled.org/core/notifier"
"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/quota"
"tangled.org/core/spindle/secrets"
"tangled.org/core/spindle/storage"
"tangled.org/core/workflow"
)
type captureEncoder struct {
messages chan *millproto.Message
}
func newCaptureEncoder() *captureEncoder {
return &captureEncoder{messages: make(chan *millproto.Message, 4)}
}
func (e *captureEncoder) Encode(msg *millproto.Message) error {
e.messages <- msg
return nil
}
type blockingEncoder struct {
started chan struct{}
release chan struct{}
}
type artifactWriterFunc func(context.Context, string, io.Reader) error
func (f artifactWriterFunc) Put(ctx context.Context, ref string, r io.Reader) error {
return f(ctx, ref, r)
}
func (e *blockingEncoder) Encode(*millproto.Message) error {
select {
case e.started <- struct{}{}:
default:
}
<-e.release
return nil
}
type fakeSlot struct {
released int
release func()
}
func (s *fakeSlot) Release() {
if s.release != nil {
s.release()
}
s.released++
}
type fakeEngine struct {
setupCalled bool
setupTrusted bool
setupNixCache *models.NixCacheGrant
runCalled bool
destroyCalled bool
acquireCalled bool
acquireStarted chan struct{}
releaseAcquire chan struct{}
slot engine.WorkflowSlot
secrets chan []secrets.UnlockedSecret
done chan struct{}
initErr error
}
func (e *fakeEngine) InitWorkflow(twf workflow.WorkflowSpec, tpl workflow.PipelineSpec) (*models.Workflow, error) {
if e.initErr != nil {
return nil, e.initErr
}
return &models.Workflow{Name: twf.Name}, nil
}
func (e *fakeEngine) SetupWorkflow(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, l models.WorkflowLogger) error {
e.setupCalled = true
e.setupTrusted = wf.TrustedSource
e.setupNixCache = wf.NixCache
return nil
}
func (e *fakeEngine) WorkflowTimeout() time.Duration { return 7 * time.Minute }
func (e *fakeEngine) DestroyWorkflow(ctx context.Context, wid models.WorkflowId) error {
e.destroyCalled = true
return nil
}
func (e *fakeEngine) RunStep(ctx context.Context, wid models.WorkflowId, w *models.Workflow, idx int, s []secrets.UnlockedSecret, l models.WorkflowLogger) error {
e.runCalled = true
if e.secrets != nil {
e.secrets <- s
}
if e.done != nil {
close(e.done)
}
return nil
}
func (e *fakeEngine) AcquireWorkflowSlot(ctx context.Context, wid models.WorkflowId, wf *models.Workflow, mode engine.AcquireMode) (engine.WorkflowSlot, error) {
e.acquireCalled = true
if e.acquireStarted != nil {
close(e.acquireStarted)
}
if e.releaseAcquire != nil {
select {
case <-e.releaseAcquire:
case <-ctx.Done():
return nil, ctx.Err()
}
}
if e.slot != nil {
return e.slot, nil
}
return &fakeSlot{}, nil
}
type resourceReportingEngine struct {
*fakeEngine
resources quota.Resources
usage engine.WorkflowResourceUsage
}
func (e *resourceReportingEngine) QuotaResources(*models.Workflow) quota.Resources {
return e.resources
}
func (e *resourceReportingEngine) WorkflowResourceUsage(*models.Workflow) (engine.WorkflowResourceUsage, bool) {
return e.usage, true
}
type fakeCacheEngine struct {
*fakeEngine
restored bool
saved bool
}
func (e *fakeCacheEngine) RestoreCache(context.Context, models.WorkflowId, *models.Workflow, storage.Storage, []models.CacheBinding, models.WorkflowLogger) error {
e.restored = true
return nil
}
func (e *fakeCacheEngine) SaveCache(context.Context, models.WorkflowId, *models.Workflow, storage.Storage, []models.CacheBinding, models.WorkflowLogger) error {
e.saved = true
return nil
}
type fakeStep struct{}
func (fakeStep) Name() string { return "test" }
func (fakeStep) Command() string { return "true" }
func (fakeStep) Kind() models.StepKind { return models.StepKindUser }
func testReserveSeat(t *testing.T, leaseID, engineName string) *millv1.ReserveSeat {
t.Helper()
twf, err := json.Marshal(workflow.WorkflowSpec{Name: "build"})
if err != nil {
t.Fatal(err)
}
repoDID := "did:web:example.com"
tpl, err := json.Marshal(workflow.PipelineSpec{TriggerMetadata: &workflow.TriggerMetadata{
Repo: &workflow.TriggerRepo{RepoDid: &repoDID},
}})
if err != nil {
t.Fatal(err)
}
return &millv1.ReserveSeat{
LeaseId: leaseID,
TargetEngine: engineName,
RawWorkflowJson: string(twf),
RawPipelineJson: string(tpl),
PipelineId: "pipeline",
RepoDid: repoDID,
}
}
func TestNewFailsWhenOutboxCannotInitialize(t *testing.T) {
d := testDB(t)
if err := d.Close(); err != nil {
t.Fatal(err)
}
n := notifier.New()
cfg := &config.Config{}
if _, err := New(cfg, nil, d, &n, slog.New(slog.NewTextHandler(io.Discard, nil)), nil, nil); err == nil {
t.Fatal("New succeeded with an unavailable outbox database")
}
}
func TestReservedEngineHandsBackHeldSlotOnce(t *testing.T) {
inner := &fakeEngine{}
slot := &fakeSlot{}
re := newReservedEngine(inner, slot)
got, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait)
if err != nil {
t.Fatal(err)
}
if got != engine.WorkflowSlot(slot) {
t.Fatal("AcquireWorkflowSlot() did not return the held slot")
}
if inner.acquireCalled {
t.Fatal("wrapper must not call the inner engine's AcquireWorkflowSlot")
}
if _, err := re.(engine.WorkflowSlotter).AcquireWorkflowSlot(context.Background(), models.WorkflowId{}, nil, engine.Wait); err == nil {
t.Fatal("second AcquireWorkflowSlot() should error")
}
}
func TestReservedEngineForwardsResourceReporting(t *testing.T) {
inner := &resourceReportingEngine{
fakeEngine: &fakeEngine{},
resources: quota.Resources{
quota.ResourceWorkflows: 1,
quota.ResourceVCPUs: 4,
},
usage: engine.WorkflowResourceUsage{
CPUUsec: 1234,
MemoryPeakBytes: 4096,
CgroupAvailable: true,
VolumeAllocatedBytes: 8192,
VolumeAvailable: true,
},
}
re := newReservedEngine(inner, &fakeSlot{})
wf := &models.Workflow{}
resources := re.(engine.WorkflowQuotaReporter).QuotaResources(wf)
if resources[quota.ResourceWorkflows] != 1 || resources[quota.ResourceVCPUs] != 4 {
t.Fatalf("QuotaResources() = %v, want workflow and vCPU resources", resources)
}
usage, ok := re.(engine.WorkflowResourceUsageReporter).WorkflowResourceUsage(wf)
if !ok {
t.Fatal("WorkflowResourceUsage() unavailable")
}
if usage != inner.usage {
t.Fatalf("WorkflowResourceUsage() = %+v, want %+v", usage, inner.usage)
}
}
func TestReservedEngineForwardsCacheRunner(t *testing.T) {
inner := &fakeCacheEngine{fakeEngine: &fakeEngine{}}
re := newReservedEngine(inner, &fakeSlot{})
runner, ok := re.(engine.CacheRunner)
if !ok {
t.Fatal("reserved engine dropped CacheRunner")
}
ctx := context.Background()
if err := runner.RestoreCache(ctx, models.WorkflowId{}, nil, nil, nil, nil); err != nil {
t.Fatal(err)
}
if err := runner.SaveCache(ctx, models.WorkflowId{}, nil, nil, nil, nil); err != nil {
t.Fatal(err)
}
if !inner.restored || !inner.saved {
t.Fatalf("cache calls were not forwarded: restored=%t saved=%t", inner.restored, inner.saved)
}
}
func TestHandleCommitIsIdempotent(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
e.active["lease-1"] = &reservation{leaseID: "lease-1", committed: true}
e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "lease-1"})
msg := <-enc.messages
if got := msg.GetCommitted().GetLeaseId(); got != "lease-1" {
t.Fatalf("Committed lease = %q, want lease-1", got)
}
}
func TestHandleCommitRejectsMissingReservation(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: "expired"})
result := (<-enc.messages).GetReserveResult()
if result == nil {
t.Fatal("missing reservation commit did not receive a ReserveResult")
}
if result.GetLeaseId() != "expired" || result.GetAccepted() {
t.Fatalf("ReserveResult = %+v, want correlated rejection", result)
}
}
func TestHandleCancelFinalizesExpiredReservation(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
e.handleCancel("lease-expired")
var ack *millv1.CancelAck
for ack == nil {
select {
case msg := <-enc.messages:
ack = msg.GetCancelAck()
case <-time.After(time.Second):
t.Fatal("cancel acknowledgement timed out")
}
}
if ack.GetLeaseId() != "lease-expired" {
t.Fatalf("CancelAck = %+v, want lease-expired", ack)
}
rows, err := e.db.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("cancel terminal outbox rows = %d, want 1", len(rows))
}
var entry millv1.Event
if err := proto.Unmarshal(rows[0].Payload, &entry); err != nil {
t.Fatal(err)
}
if got := entry.GetAttemptResult().GetStatus(); got != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED {
t.Fatalf("cancel terminal = %v, want CANCELLED", got)
}
}
func TestTerminalOutboxEntrySuppressesLaterEvents(t *testing.T) {
e := testExecutor(t)
if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
t.Fatal(err)
}
e.handleCancel("lease-1")
if err := e.appendStatus("lease-1", &db.StatusRow{Status: string(models.StatusKindRunning)}); err != nil {
t.Fatal(err)
}
rows, err := e.db.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("outbox rows after terminal = %d, want 1", len(rows))
}
var entry millv1.Event
if err := proto.Unmarshal(rows[0].Payload, &entry); err != nil {
t.Fatal(err)
}
if got := entry.GetAttemptResult().GetStatus(); got != millv1.TerminalStatus_TERMINAL_STATUS_SUCCESS {
t.Fatalf("terminal status = %v, want SUCCESS", got)
}
}
func TestTerminalGuardSurvivesExecutorRestart(t *testing.T) {
d := testDB(t)
first := &Executor{
db: d,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := first.initOutbox(); err != nil {
t.Fatal(err)
}
if err := first.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
t.Fatal(err)
}
restarted := &Executor{
db: d,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := restarted.initOutbox(); err != nil {
t.Fatal(err)
}
restarted.handleCancel("lease-1")
rows, err := d.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("outbox rows after restart and duplicate cancel = %d, want 1", len(rows))
}
}
func TestPendingArtifactRecoveryCompletesWithoutDeadlock(t *testing.T) {
d := testDB(t)
logDir := t.TempDir()
wid := models.WorkflowId{
PipelineId: models.PipelineId("3abc"),
Name: "build",
}
logPath := models.LogFilePath(logDir, wid)
if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(logPath, []byte("finished"), 0600); err != nil {
t.Fatal(err)
}
if err := d.SavePendingArtifact(
"lease-1", wid, string(models.StatusKindSuccess), "", 0,
"logs/lease-1.log", "sha256:test", string(engine.FailureClassNone),
string(engine.FailureReasonSuccess), true,
); err != nil {
t.Fatal(err)
}
uploaded := make(chan string, 1)
e := &Executor{
db: d,
cfg: &config.Config{Server: config.Server{LogDir: logDir}},
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
writer: artifactWriterFunc(func(_ context.Context, ref string, r io.Reader) error {
if _, err := io.ReadAll(r); err != nil {
return err
}
uploaded <- ref
return nil
}),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatalf("initOutbox: %v", err)
}
done := make(chan error, 1)
go func() { done <- e.recoverPendingArtifacts(context.Background()) }()
select {
case err := <-done:
if err != nil {
t.Fatalf("recoverPendingArtifacts: %v", err)
}
case <-time.After(time.Second):
t.Fatal("pending artifact recovery deadlocked")
}
select {
case ref := <-uploaded:
if ref != "logs/lease-1.log" {
t.Fatalf("uploaded ref = %q", ref)
}
default:
t.Fatal("pending artifact was not uploaded")
}
pending, err := d.ListPendingArtifacts()
if err != nil {
t.Fatal(err)
}
if len(pending) != 0 {
t.Fatalf("pending artifacts after recovery = %d, want 0", len(pending))
}
rows, err := d.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("outbox rows after recovery = %d, want 1", len(rows))
}
}
func TestPendingArtifactRecoveryKeepsRowWhenLogIsMissing(t *testing.T) {
d := testDB(t)
wid := models.WorkflowId{
PipelineId: models.PipelineId("3abc"),
Name: "build",
}
if err := d.SavePendingArtifact(
"lease-1", wid, string(models.StatusKindSuccess), "", 0,
"logs/lease-1.log", "sha256:test", string(engine.FailureClassNone),
string(engine.FailureReasonSuccess), true,
); err != nil {
t.Fatal(err)
}
e := &Executor{
db: d,
cfg: &config.Config{Server: config.Server{LogDir: t.TempDir()}},
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
writer: artifactWriterFunc(func(context.Context, string, io.Reader) error {
t.Fatal("writer called without a log file")
return nil
}),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatalf("initOutbox: %v", err)
}
if err := e.recoverPendingArtifacts(context.Background()); err != nil {
t.Fatalf("recoverPendingArtifacts: %v", err)
}
pending, err := d.ListPendingArtifacts()
if err != nil {
t.Fatal(err)
}
if len(pending) != 1 {
t.Fatalf("pending artifacts with missing log = %d, want 1", len(pending))
}
rows, err := d.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 0 {
t.Fatalf("outbox rows with missing log = %d, want 0", len(rows))
}
}
func TestPendingArtifactRecoveryUsesAggregateDeadline(t *testing.T) {
d := testDB(t)
logDir := t.TempDir()
wids := []models.WorkflowId{
{PipelineId: models.PipelineId("3abc"), Name: "build"},
{PipelineId: models.PipelineId("3def"), Name: "test"},
}
for i, wid := range wids {
logPath := models.LogFilePath(logDir, wid)
if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(logPath, []byte("finished"), 0600); err != nil {
t.Fatal(err)
}
leaseID := fmt.Sprintf("lease-%d", i+1)
if err := d.SavePendingArtifact(
leaseID, wid, string(models.StatusKindSuccess), "", 0,
"logs/"+leaseID+".log", "sha256:test", string(engine.FailureClassNone),
string(engine.FailureReasonSuccess), true,
); err != nil {
t.Fatal(err)
}
}
var deadlines []time.Time
e := &Executor{
db: d,
cfg: &config.Config{Server: config.Server{LogDir: logDir}},
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
writer: artifactWriterFunc(func(ctx context.Context, _ string, _ io.Reader) error {
deadline, ok := ctx.Deadline()
if !ok {
t.Fatal("artifact recovery context has no deadline")
}
deadlines = append(deadlines, deadline)
return errors.New("unavailable")
}),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatalf("initOutbox: %v", err)
}
started := time.Now()
if err := e.recoverPendingArtifacts(context.Background()); err != nil {
t.Fatalf("recoverPendingArtifacts: %v", err)
}
finished := time.Now()
if len(deadlines) != len(wids) {
t.Fatalf("upload attempts = %d, want %d", len(deadlines), len(wids))
}
if !deadlines[0].Equal(deadlines[1]) {
t.Fatalf("artifact recovery deadlines differ: %v, %v", deadlines[0], deadlines[1])
}
wantTimeout := 2 * time.Minute
if deadlines[0].Before(started.Add(wantTimeout)) || deadlines[0].After(finished.Add(wantTimeout)) {
t.Fatalf("artifact recovery deadline = %v, want %v after startup", deadlines[0], wantTimeout)
}
}
func TestHandleCommitPreservesPreauthorizedSecrets(t *testing.T) {
d := testDB(t)
n := notifier.New()
enc := newCaptureEncoder()
e := &Executor{
cfg: &config.Config{Server: config.Server{LogDir: t.TempDir()}},
db: d,
n: &n,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
enc: enc,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
e.lifecycleCtx = context.Background()
inner := &fakeEngine{secrets: make(chan []secrets.UnlockedSecret, 1), done: make(chan struct{})}
slot := &fakeSlot{}
repoDid, err := syntax.ParseDID("did:web:example.com")
if err != nil {
t.Fatal(err)
}
repoDidString := repoDid.String()
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
realEngine: inner,
slot: slot,
wf: &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
pipeline: &models.Pipeline{
RepoDid: repoDid,
TriggerMetadata: &workflow.TriggerMetadata{
Repo: &workflow.TriggerRepo{RepoDid: &repoDidString},
},
TrustedSource: true,
},
}
e.active[res.leaseID] = res
e.handleCommit(context.Background(), &millv1.CommitLease{
LeaseId: res.leaseID,
Secrets: []*millv1.Secret{{Key: "TOKEN", Value: "secret-value"}},
})
if got := (<-enc.messages).GetCommitted().GetLeaseId(); got != res.leaseID {
t.Fatalf("Committed lease = %q, want %q", got, res.leaseID)
}
if got := e.maskSecrets(res, "value=secret-value"); got != "value=***" {
t.Fatalf("masked log = %q", got)
}
select {
case got := <-inner.secrets:
if len(got) != 1 || got[0].Key != "TOKEN" || got[0].Value != "secret-value" {
t.Fatalf("RunStep secrets = %+v", got)
}
case <-time.After(2 * time.Second):
t.Fatal("RunStep did not receive CommitLease secrets")
}
select {
case <-inner.done:
case <-time.After(2 * time.Second):
t.Fatal("workflow did not finish")
}
if res.stopTail != nil {
res.stopTail()
}
e.jobsWG.Wait()
}
func TestConnectHandshakesWhilePendingArtifactRecoveryIsBlocked(t *testing.T) {
handshake := make(chan error, 1)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
if err != nil {
handshake <- err
return
}
defer conn.Close()
stream := millproto.NewWSStream(conn)
dec := millproto.NewDecoder(stream)
enc := millproto.NewEncoder(stream)
helloMsg, err := dec.Decode()
if err != nil {
handshake <- err
return
}
hello := helloMsg.GetHello()
if hello == nil {
handshake <- errors.New("first executor message was not Hello")
return
}
if err := enc.Encode(&millproto.Message{Resume: &millv1.Resume{Epoch: hello.GetEpoch()}}); err != nil {
handshake <- err
return
}
snapshotMsg, err := dec.Decode()
if err != nil {
handshake <- err
return
}
if snapshotMsg.GetNodeSnapshot() == nil {
handshake <- errors.New("executor did not send NodeSnapshot after Resume")
return
}
handshake <- nil
for {
if _, err := dec.Decode(); err != nil {
return
}
}
}))
t.Cleanup(srv.Close)
e := testSessionExecutor(t, "ws"+strings.TrimPrefix(srv.URL, "http"))
n := notifier.New()
e.n = &n
logDir := t.TempDir()
e.cfg.Server.LogDir = logDir
wid := models.WorkflowId{
PipelineId: models.PipelineId("3abc"),
Name: "build",
}
logPath := models.LogFilePath(logDir, wid)
if err := os.MkdirAll(filepath.Dir(logPath), 0755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(logPath, []byte("finished"), 0600); err != nil {
t.Fatal(err)
}
if err := e.db.SavePendingArtifact(
"lease-1", wid, string(models.StatusKindSuccess), "", 0,
"logs/lease-1.log", "sha256:test", string(engine.FailureClassNone),
string(engine.FailureReasonSuccess), true,
); err != nil {
t.Fatal(err)
}
recoveryStarted := make(chan struct{})
e.writer = artifactWriterFunc(func(ctx context.Context, _ string, _ io.Reader) error {
close(recoveryStarted)
<-ctx.Done()
return ctx.Err()
})
ctx, cancel := context.WithCancel(context.Background())
done := make(chan struct{})
go func() {
e.Connect(ctx)
close(done)
}()
select {
case <-recoveryStarted:
case <-time.After(time.Second):
cancel()
t.Fatal("pending artifact recovery did not start")
}
select {
case err := <-handshake:
if err != nil {
cancel()
t.Fatalf("mill handshake: %v", err)
}
case <-time.After(time.Second):
cancel()
t.Fatal("blocked artifact recovery delayed the mill handshake")
}
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("Connect did not stop after cancellation")
}
}
func TestHandleCommitRejectsCachesAndSecretsForForkSource(t *testing.T) {
d := testDB(t)
n := notifier.New()
enc := newCaptureEncoder()
cache, err := storage.NewDisk(t.TempDir())
if err != nil {
t.Fatal(err)
}
e := &Executor{
db: d,
n: &n,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
cfg: &config.Config{Server: config.Server{LogDir: t.TempDir()}},
cache: cache,
enc: enc,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
e.lifecycleCtx = context.Background()
inner := &fakeCacheEngine{fakeEngine: &fakeEngine{
secrets: make(chan []secrets.UnlockedSecret, 1),
done: make(chan struct{}),
}}
targetRepoDid, err := syntax.ParseDID("did:plc:target")
if err != nil {
t.Fatal(err)
}
sourceRepoDid := "did:plc:fork"
cacheObjectID := "11111111-1111-4111-8111-111111111111"
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
realEngine: inner,
slot: &fakeSlot{},
wf: &models.Workflow{
Name: "build",
Steps: []models.Step{fakeStep{}},
Caches: []models.CacheEntry{{Key: "deps", Paths: []string{"deps"}}},
},
pipeline: &models.Pipeline{
RepoDid: targetRepoDid,
TriggerMetadata: &workflow.TriggerMetadata{
SourceRepo: &sourceRepoDid,
Repo: &workflow.TriggerRepo{RepoDid: &sourceRepoDid},
},
},
}
e.active[res.leaseID] = res
e.handleCommit(context.Background(), &millv1.CommitLease{
LeaseId: res.leaseID,
Secrets: []*millv1.Secret{{Key: "TOKEN", Value: "secret-value"}},
CacheBindings: []*millv1.CacheBinding{{
EntryIndex: 0,
RestoreId: cacheObjectID,
RestoreKey: "objects/" + targetRepoDid.String() + "/" + cacheObjectID,
SaveId: cacheObjectID,
SaveKey: "objects/" + targetRepoDid.String() + "/" + cacheObjectID,
}},
})
if got := (<-enc.messages).GetCommitted().GetLeaseId(); got != res.leaseID {
t.Fatalf("Committed lease = %q, want %q", got, res.leaseID)
}
e.jobsWG.Wait()
if inner.restored || inner.saved {
t.Fatalf("fork cache access = restored %t, saved %t", inner.restored, inner.saved)
}
if got := <-inner.secrets; len(got) != 0 {
t.Fatalf("fork secrets = %+v, want none", got)
}
if !inner.setupCalled || inner.setupTrusted {
t.Fatalf("fork setup = called %t, trusted %t, want called and untrusted", inner.setupCalled, inner.setupTrusted)
}
}
type timeoutCacheEngine struct {
*fakeEngine
events []string
}
func (e *timeoutCacheEngine) DestroyWorkflow(context.Context, models.WorkflowId) error {
e.events = append(e.events, "destroy")
return nil
}
func (e *timeoutCacheEngine) WorkflowTimeout() time.Duration { return 100 * time.Millisecond }
func (e *timeoutCacheEngine) RunStep(ctx context.Context, _ models.WorkflowId, _ *models.Workflow, _ int, _ []secrets.UnlockedSecret, _ models.WorkflowLogger) error {
<-ctx.Done()
return engine.ErrTimedOut
}
func (e *timeoutCacheEngine) RestoreCache(context.Context, models.WorkflowId, *models.Workflow, storage.Storage, []models.CacheBinding, models.WorkflowLogger) error {
return nil
}
func (e *timeoutCacheEngine) SaveCache(ctx context.Context, _ models.WorkflowId, _ *models.Workflow, store storage.Storage, caches []models.CacheBinding, _ models.WorkflowLogger) error {
for _, c := range caches {
if err := store.Put(ctx, c.SaveKey, strings.NewReader("archive")); err != nil {
return err
}
}
e.events = append(e.events, "save")
return nil
}
func TestCommittedWorkflowSavesAlwaysCacheAfterTimeout(t *testing.T) {
d := testDB(t)
n := notifier.New()
cache, err := storage.NewDisk(t.TempDir())
if err != nil {
t.Fatal(err)
}
enc := newCaptureEncoder()
e := &Executor{
db: d,
n: &n,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
cfg: &config.Config{Server: config.Server{LogDir: t.TempDir()}},
cache: cache,
enc: enc,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
e.lifecycleCtx = context.Background()
repoDid, err := syntax.ParseDID("did:plc:target")
if err != nil {
t.Fatal(err)
}
alwaysID := "11111111-1111-4111-8111-111111111111"
onSuccessID := "22222222-2222-4222-8222-222222222222"
inner := &timeoutCacheEngine{fakeEngine: &fakeEngine{}}
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
realEngine: inner,
slot: &fakeSlot{},
wf: &models.Workflow{
Name: "build",
Steps: []models.Step{fakeStep{}},
Caches: []models.CacheEntry{
{Key: "always", Paths: []string{"a"}, When: "always"},
{Key: "on-success", Paths: []string{"b"}},
},
},
pipeline: &models.Pipeline{RepoDid: repoDid, TrustedSource: true},
}
e.active[res.leaseID] = res
e.handleCommit(context.Background(), &millv1.CommitLease{
LeaseId: res.leaseID,
CacheBindings: []*millv1.CacheBinding{
{EntryIndex: 0, SaveId: alwaysID, SaveKey: "objects/" + repoDid.String() + "/" + alwaysID},
{EntryIndex: 1, SaveId: onSuccessID, SaveKey: "objects/" + repoDid.String() + "/" + onSuccessID},
},
})
if got := (<-enc.messages).GetCommitted().GetLeaseId(); got != res.leaseID {
t.Fatalf("Committed lease = %q, want %q", got, res.leaseID)
}
go func() {
for range enc.messages {
}
}()
e.jobsWG.Wait()
if status, err := d.GetStatus(res.wid); err != nil || status != models.StatusKindTimeout {
t.Fatalf("status = %s, %v, want timeout", status, err)
}
if !inner.setupCalled || !inner.setupTrusted {
t.Fatalf("setup = called %t, trusted %t, want called and trusted", inner.setupCalled, inner.setupTrusted)
}
if !slices.Equal(inner.events, []string{"save", "destroy"}) {
t.Fatalf("events = %v, want save then a single destroy", inner.events)
}
rows, err := d.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
updates := map[string]millv1.CacheUpdateAction{}
for _, row := range rows {
ev, err := decodeEvent(row.Payload, row.Seqno)
if err != nil {
t.Fatal(err)
}
if u := ev.GetCacheUpdate(); u != nil {
updates[u.GetId()] = u.GetAction()
}
}
if updates[alwaysID] != millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED {
t.Fatalf("always cache update = %v, want stored", updates[alwaysID])
}
if updates[onSuccessID] == millv1.CacheUpdateAction_CACHE_UPDATE_ACTION_STORED {
t.Fatal("on-success cache was stored after a timeout")
}
}
func TestRunSessionCancellationClosesStalledWebsocket(t *testing.T) {
connected := make(chan struct{})
release := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
conn, err := websocket.Upgrade(w, r, nil, 1024, 1024)
if err != nil {
return
}
defer conn.Close()
close(connected)
<-release
}))
t.Cleanup(func() {
close(release)
srv.Close()
})
e := testSessionExecutor(t, "ws"+strings.TrimPrefix(srv.URL, "http"))
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() { done <- e.runSession(ctx) }()
<-connected
cancel()
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("runSession did not return after context cancellation")
}
}
func testSessionExecutor(t *testing.T, url string) *Executor {
d := testDB(t)
e := &Executor{
millURL: url,
seats: 1,
engines: make(map[string]models.Engine),
db: d,
cfg: &config.Config{Server: config.Server{Dev: true}},
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
return e
}
func testExecutor(t *testing.T) *Executor {
d := testDB(t)
e := &Executor{
db: d,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
return e
}
func TestReserveHonorsExecutorSeatCapacity(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
e.seats = 1
e.engines = map[string]models.Engine{"dummy": &fakeEngine{}}
e.handleReserve(context.Background(), testReserveSeat(t, "lease-1", "dummy"))
first := (<-enc.messages).GetReserveResult()
if first == nil || !first.GetAccepted() {
t.Fatalf("first reserve result = %+v, want accepted", first)
}
if snapshot := (<-enc.messages).GetNodeSnapshot(); snapshot == nil {
t.Fatal("accepted reservation did not publish a snapshot")
}
e.handleReserve(context.Background(), testReserveSeat(t, "lease-2", "dummy"))
second := (<-enc.messages).GetReserveResult()
if second == nil || second.GetAccepted() {
t.Fatalf("second reserve result = %+v, want transient rejection", second)
}
if second.GetRejectClass() != millv1.RejectClass_REJECT_CLASS_TRANSIENT {
t.Fatalf("second reject class = %v, want transient", second.GetRejectClass())
}
if second.GetRejectReason() != "no executor seats available" {
t.Fatalf("second reject reason = %q", second.GetRejectReason())
}
cleanup, ok := e.takeUncommittedReservation("lease-1", true)
if !ok {
t.Fatal("first reservation missing during cleanup")
}
cleanup()
}
func testDB(t *testing.T) *db.DB {
d, err := db.Make(context.Background(), filepath.Join(t.TempDir(), "spindle.db"))
if err != nil {
t.Fatal(err)
}
return d
}
func TestFinishJobReportsCancelledReservationAsCancelled(t *testing.T) {
d := testDB(t)
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
cancelled: true,
}
e := &Executor{
db: d,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: map[string]*reservation{res.leaseID: res},
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
e.finishJob(res, &db.StatusRow{
Pipeline: string(res.wid.PipelineId),
Workflow: res.wid.Name,
Status: string(models.StatusKindFailed),
})
rows, err := d.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("outbox rows = %d, want 1", len(rows))
}
var entry millv1.Event
if err := proto.Unmarshal(rows[0].Payload, &entry); err != nil {
t.Fatal(err)
}
got := entry.GetAttemptResult().GetStatus()
if got != millv1.TerminalStatus_TERMINAL_STATUS_CANCELLED {
t.Fatalf("terminal status = %v, want CANCELLED", got)
}
}
func TestFinishJobWaitsForEngineCleanup(t *testing.T) {
d := testDB(t)
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{
PipelineId: models.PipelineId("r"),
Name: "build",
},
runDone: make(chan struct{}),
}
e := &Executor{
db: d,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: map[string]*reservation{res.leaseID: res},
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
done := make(chan error, 1)
go func() {
done <- e.finishJob(res, &db.StatusRow{
Pipeline: string(res.wid.PipelineId),
Workflow: res.wid.Name,
Status: string(models.StatusKindSuccess),
})
}()
select {
case err := <-done:
t.Fatalf("finishJob returned before engine cleanup: %v", err)
case <-time.After(50 * time.Millisecond):
}
rows, err := d.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 0 {
t.Fatalf("terminal rows before engine cleanup = %d, want 0", len(rows))
}
close(res.runDone)
select {
case err := <-done:
if err != nil {
t.Fatal(err)
}
case <-time.After(time.Second):
t.Fatal("finishJob did not resume after engine cleanup")
}
}
func outboxEvents(t *testing.T, d *db.DB) []*millv1.Event {
t.Helper()
rows, err := d.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
events := make([]*millv1.Event, 0, len(rows))
for _, row := range rows {
var entry millv1.Event
if err := proto.Unmarshal(row.Payload, &entry); err != nil {
t.Fatal(err)
}
events = append(events, &entry)
}
return events
}
func TestFinishedStatusGoesOutBeforeEngineCleanup(t *testing.T) {
e := testExecutor(t)
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
runDone: make(chan struct{}),
millAcceptsOutcome: true,
}
e.active[res.leaseID] = res
returned := make(chan error, 1)
go func() {
returned <- e.onStatusRow(&db.StatusRow{
Pipeline: string(res.wid.PipelineId),
Workflow: res.wid.Name,
Status: string(models.StatusKindFailed),
})
}()
select {
case err := <-returned:
if err != nil {
t.Fatal(err)
}
case <-time.After(time.Second):
t.Fatal("the status loop waited on engine cleanup")
}
events := outboxEvents(t, e.db)
if len(events) != 1 || events[0].GetAttemptOutcome().GetStatus() != millv1.TerminalStatus_TERMINAL_STATUS_FAILED {
t.Fatalf("events before cleanup = %v, want one failed outcome", events)
}
e.mu.Lock()
res.cancelled = true
e.mu.Unlock()
close(res.runDone)
e.finishWG.Wait()
events = outboxEvents(t, e.db)
if len(events) != 2 || events[1].GetAttemptResult().GetStatus() != millv1.TerminalStatus_TERMINAL_STATUS_FAILED {
t.Fatalf("events after cleanup = %v, want the outcome then a failed result", events)
}
}
func TestOlderMillGetsNoOutcome(t *testing.T) {
e := testExecutor(t)
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
runDone: make(chan struct{}),
}
e.active[res.leaseID] = res
returned := make(chan error, 1)
go func() {
returned <- e.onStatusRow(&db.StatusRow{
Pipeline: string(res.wid.PipelineId),
Workflow: res.wid.Name,
Status: string(models.StatusKindFailed),
})
}()
time.Sleep(50 * time.Millisecond)
if events := outboxEvents(t, e.db); len(events) != 0 {
t.Fatalf("events before cleanup = %v, an older mill rejects outcomes", events)
}
close(res.runDone)
if err := <-returned; err != nil {
t.Fatal(err)
}
events := outboxEvents(t, e.db)
if len(events) != 1 || events[0].GetAttemptResult() == nil {
t.Fatalf("events = %v, want only the result", events)
}
}
func TestReplayRejectsMalformedOutboxRow(t *testing.T) {
d := testDB(t)
e := &Executor{
db: d,
enc: newCaptureEncoder(),
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
if _, err := d.AppendOutboxRow([]byte("not protobuf"), true); err != nil {
t.Fatal(err)
}
if err := e.replay(0); err == nil {
t.Fatal("replay accepted a malformed row and would leave a permanent seqno gap")
}
}
func TestReplayDropsRowsAlreadyAcknowledgedByResume(t *testing.T) {
e := testExecutor(t)
if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
t.Fatal(err)
}
rows, err := e.db.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("outbox rows = %d, want 1", len(rows))
}
e.enc = newCaptureEncoder()
if err := e.replay(rows[0].Seqno); err != nil {
t.Fatalf("replay: %v", err)
}
rows, err = e.db.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 0 {
t.Fatalf("outbox rows after resume = %d, want 0", len(rows))
}
if len(e.terminalSeqnos) != 0 {
t.Fatalf("terminal guards after resume = %d, want 0", len(e.terminalSeqnos))
}
select {
case <-e.outboxIdle():
default:
t.Fatal("resumed acknowledgement left the outbox busy")
}
}
func TestSocketCancellationIndependence(t *testing.T) {
d := testDB(t)
n := notifier.New()
e := &Executor{
db: d,
n: &n,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
cfg: &config.Config{Server: config.Server{LogDir: t.TempDir()}},
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
lifecycleCtx, cancelLifecycle := context.WithCancel(context.Background())
defer cancelLifecycle()
e.lifecycleCtx = lifecycleCtx
inner := &fakeEngine{secrets: make(chan []secrets.UnlockedSecret, 1), done: make(chan struct{})}
slot := &fakeSlot{}
repoDid, err := syntax.ParseDID("did:web:example.com")
if err != nil {
t.Fatal(err)
}
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
realEngine: inner,
slot: slot,
wf: &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
pipeline: &models.Pipeline{RepoDid: repoDid},
}
e.active[res.leaseID] = res
sessionCtx, cancelSession := context.WithCancel(lifecycleCtx)
e.handleCommit(sessionCtx, &millv1.CommitLease{
LeaseId: "lease-1",
})
cancelSession()
// session disconnect must not cancel the running job
select {
case <-inner.done:
case <-time.After(2 * time.Second):
t.Fatal("workflow did not complete even though websocket session was cancelled")
}
e.jobsWG.Wait()
}
func TestMonotonicSnapshots(t *testing.T) {
enc := newCaptureEncoder()
e := &Executor{
enc: enc,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
}
e.pushSnapshot()
msg1 := <-enc.messages
seq1 := msg1.GetNodeSnapshot().GetSeqno()
if seq1 != 1 {
t.Fatalf("first seq = %d, want 1", seq1)
}
e.pushSnapshot()
msg2 := <-enc.messages
seq2 := msg2.GetNodeSnapshot().GetSeqno()
if seq2 != 2 {
t.Fatalf("second seq = %d, want 2", seq2)
}
}
func TestDrainWaitsForActiveReservation(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
e.engines = map[string]models.Engine{"dummy": &fakeEngine{}}
res := &reservation{leaseID: "lease-1", committed: true}
e.active[res.leaseID] = res
done := make(chan error, 1)
go func() {
done <- e.Drain(context.Background())
}()
snapshot := (<-enc.messages).GetNodeSnapshot()
if snapshot == nil || snapshot.GetEngines()["dummy"].GetAvailable() {
t.Fatalf("drain snapshot = %+v, want unavailable engine", snapshot)
}
if got := snapshot.GetActiveLeaseIds(); len(got) != 1 || got[0] != res.leaseID {
t.Fatalf("active leases = %v, want [%s]", got, res.leaseID)
}
select {
case err := <-done:
t.Fatalf("Drain returned before reservation cleanup: %v", err)
default:
}
e.mu.Lock()
cleanup := e.removeReservationLocked(res, false)
e.mu.Unlock()
cleanup()
select {
case err := <-done:
if err != nil {
t.Fatalf("Drain: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("Drain did not return after reservation cleanup")
}
}
func TestDrainWaitsForReservationCleanup(t *testing.T) {
releaseStarted := make(chan struct{})
releaseSlot := make(chan struct{})
slot := &fakeSlot{release: func() {
close(releaseStarted)
<-releaseSlot
}}
e := testExecutor(t)
e.enc = newCaptureEncoder()
res := &reservation{leaseID: "lease-1", slot: slot}
e.active[res.leaseID] = res
cleanup, ok := e.takeUncommittedReservation(res.leaseID, true)
if !ok {
t.Fatal("failed to take reservation")
}
cleanupDone := make(chan struct{})
go func() {
cleanup()
close(cleanupDone)
}()
<-releaseStarted
done := make(chan error, 1)
go func() {
done <- e.Drain(context.Background())
}()
<-e.enc.(*captureEncoder).messages
select {
case err := <-done:
t.Fatalf("Drain returned before slot release completed: %v", err)
default:
}
close(releaseSlot)
<-cleanupDone
select {
case err := <-done:
if err != nil {
t.Fatalf("Drain: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("Drain did not return after slot release")
}
if slot.released != 1 {
t.Fatalf("slot releases = %d, want 1", slot.released)
}
}
func TestDrainWaitsForWorkflowCleanup(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
e.jobsWG.Add(1)
done := make(chan error, 1)
go func() {
done <- e.Drain(context.Background())
}()
<-enc.messages
select {
case err := <-done:
t.Fatalf("Drain returned before workflow cleanup: %v", err)
default:
}
e.jobsWG.Done()
select {
case err := <-done:
if err != nil {
t.Fatalf("Drain: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("Drain did not return after workflow cleanup")
}
}
func TestDrainWaitsForTerminalAcknowledgement(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
t.Fatal(err)
}
batch := (<-enc.messages).GetEventBatch()
if batch == nil || len(batch.GetEvents()) != 1 {
t.Fatalf("terminal batch = %+v, want one event", batch)
}
done := make(chan error, 1)
go func() {
done <- e.Drain(context.Background())
}()
<-enc.messages
select {
case err := <-done:
t.Fatalf("Drain returned before terminal acknowledgement: %v", err)
default:
}
e.handleAck(&millv1.Ack{Epoch: e.epoch, UpToSeqno: batch.GetEvents()[0].GetSeqno()})
select {
case err := <-done:
if err != nil {
t.Fatalf("Drain: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("Drain did not return after terminal acknowledgement")
}
}
func TestDrainReturnsWithDurableOutboxAfterReconnectGrace(t *testing.T) {
e := testExecutor(t)
e.cfg = &config.Config{}
e.cfg.Mill.ReconnectGrace = 10 * time.Millisecond
if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
t.Fatal(err)
}
if err := e.Drain(context.Background()); err != nil {
t.Fatalf("Drain with disconnected Mill: %v", err)
}
rows, err := e.db.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 1 {
t.Fatalf("durable outbox rows = %d, want 1", len(rows))
}
}
func TestDrainAllowsReconnectToAcknowledgeOutbox(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.cfg = &config.Config{}
e.cfg.Mill.ReconnectGrace = 2 * time.Second
e.enc = enc
if err := e.appendTerminal("lease-1", string(models.StatusKindSuccess), nil); err != nil {
t.Fatal(err)
}
batch := (<-enc.messages).GetEventBatch()
if batch == nil || len(batch.GetEvents()) != 1 {
t.Fatalf("terminal batch = %+v, want one event", batch)
}
done := make(chan error, 1)
go func() {
done <- e.Drain(context.Background())
}()
<-enc.messages
e.connMu.Lock()
e.enc = nil
e.connMu.Unlock()
select {
case err := <-done:
t.Fatalf("Drain returned before reconnect grace elapsed: %v", err)
case <-time.After(50 * time.Millisecond):
}
e.connMu.Lock()
e.enc = enc
e.connMu.Unlock()
e.handleAck(&millv1.Ack{Epoch: e.epoch, UpToSeqno: batch.GetEvents()[0].GetSeqno()})
select {
case err := <-done:
if err != nil {
t.Fatalf("Drain after reconnect acknowledgement: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("Drain did not accept acknowledgement during reconnect grace")
}
rows, err := e.db.ListOutboxRows()
if err != nil {
t.Fatal(err)
}
if len(rows) != 0 {
t.Fatalf("durable outbox rows = %d, want 0", len(rows))
}
}
func TestReserveCannotRacePastDrain(t *testing.T) {
enc := newCaptureEncoder()
e := testExecutor(t)
e.enc = enc
slot := &fakeSlot{}
inner := &fakeEngine{
acquireStarted: make(chan struct{}),
releaseAcquire: make(chan struct{}),
slot: slot,
}
e.engines = map[string]models.Engine{"dummy": inner}
reserve := testReserveSeat(t, "lease-1", "dummy")
reserveDone := make(chan struct{})
go func() {
e.handleReserve(context.Background(), reserve)
close(reserveDone)
}()
<-inner.acquireStarted
if err := e.Drain(context.Background()); err != nil {
t.Fatalf("Drain: %v", err)
}
snapshot := (<-enc.messages).GetNodeSnapshot()
if snapshot == nil || snapshot.GetEngines()["dummy"].GetAvailable() {
t.Fatalf("drain snapshot = %+v, want unavailable engine", snapshot)
}
close(inner.releaseAcquire)
<-reserveDone
rejected := (<-enc.messages).GetReserveResult()
if rejected == nil || rejected.GetAccepted() || rejected.GetRejectReason() != "draining" {
t.Fatalf("reserve result = %+v, want draining rejection", rejected)
}
if slot.released != 1 {
t.Fatalf("slot releases = %d, want 1", slot.released)
}
if len(e.active) != 0 {
t.Fatalf("active reservations = %d, want 0", len(e.active))
}
}
func TestDrainCancellationIncludesUnavailableSnapshot(t *testing.T) {
enc := &blockingEncoder{
started: make(chan struct{}, 1),
release: make(chan struct{}),
}
e := testExecutor(t)
e.enc = enc
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
go func() {
done <- e.Drain(ctx)
}()
<-enc.started
cancel()
select {
case err := <-done:
if !errors.Is(err, context.Canceled) {
t.Fatalf("Drain error = %v, want context canceled", err)
}
case <-time.After(2 * time.Second):
t.Fatal("Drain ignored cancellation while sending the unavailable snapshot")
}
close(enc.release)
}
func TestTimerRace(t *testing.T) {
d := testDB(t)
e := &Executor{
db: d,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
twf, _ := json.Marshal(workflow.WorkflowSpec{Name: "build"})
tpl, _ := json.Marshal(workflow.PipelineSpec{TriggerMetadata: &workflow.TriggerMetadata{}})
inner := &fakeEngine{}
e.engines = map[string]models.Engine{"microvm": inner}
e.handleReserve(context.Background(), &millv1.ReserveSeat{
LeaseId: "lease-1",
TargetEngine: "microvm",
RawWorkflowJson: string(twf),
RawPipelineJson: string(tpl),
PipelineId: "pipeline",
TtlSeconds: 1,
RepoDid: "did:web:example.com",
})
e.mu.Lock()
res := e.active["lease-1"]
e.mu.Unlock()
if res == nil {
t.Fatal("reservation was not added")
}
deadline := time.Now().Add(5 * time.Second)
for {
e.mu.Lock()
activeLen := len(e.active)
e.mu.Unlock()
if activeLen == 0 {
break
}
if time.Now().After(deadline) {
t.Fatal("reservation was leaked and never expired")
}
time.Sleep(10 * time.Millisecond)
}
}
func TestStructuredShutdown(t *testing.T) {
d := testDB(t)
n := notifier.New()
e := &Executor{
db: d,
n: &n,
l: slog.New(slog.NewTextHandler(io.Discard, nil)),
active: make(map[string]*reservation),
maxOutboxBytes: 10 * 1024 * 1024,
cfg: &config.Config{Server: config.Server{LogDir: t.TempDir()}},
}
if err := e.initOutbox(); err != nil {
t.Fatal(err)
}
ctx, cancel := context.WithCancel(context.Background())
e.lifecycleCtx = ctx
inner := &fakeEngine{secrets: make(chan []secrets.UnlockedSecret, 1), done: make(chan struct{})}
slot := &fakeSlot{}
repoDid, err := syntax.ParseDID("did:web:example.com")
if err != nil {
t.Fatal(err)
}
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
realEngine: inner,
slot: slot,
wf: &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
pipeline: &models.Pipeline{RepoDid: repoDid},
}
e.active[res.leaseID] = res
e.handleCommit(ctx, &millv1.CommitLease{
LeaseId: "lease-1",
})
cancel()
e.jobsWG.Wait()
select {
case <-inner.done:
default:
t.Fatal("shutdown returned but running job did not finish")
}
}
func TestHandleCommitHandsTheRunItsRepoCacheTokens(t *testing.T) {
lease := &millv1.NixCacheGrant{Cache: "plc-target", PublicKey: "plc-target:AAAA=", PullToken: "pull", PushToken: "push", MaxUploadBytes: 200}
for _, tt := range []struct {
name string
trusted bool
wantPush string
}{
{"trusted run gets pull and push", true, "push"},
{"fork run gets pull only", false, ""},
} {
t.Run(tt.name, func(t *testing.T) {
n := notifier.New()
enc := newCaptureEncoder()
e := testExecutor(t)
e.cfg = &config.Config{Server: config.Server{LogDir: t.TempDir()}}
e.n = &n
e.enc = enc
e.lifecycleCtx = context.Background()
inner := &fakeEngine{done: make(chan struct{})}
res := &reservation{
leaseID: "lease-1",
wid: models.WorkflowId{PipelineId: models.PipelineId("r"), Name: "build"},
realEngine: inner,
slot: &fakeSlot{},
wf: &models.Workflow{Name: "build", Steps: []models.Step{fakeStep{}}},
pipeline: &models.Pipeline{RepoDid: "did:plc:target", TrustedSource: tt.trusted},
}
e.active[res.leaseID] = res
e.handleCommit(context.Background(), &millv1.CommitLease{LeaseId: res.leaseID, NixCache: lease})
<-enc.messages
select {
case <-inner.done:
case <-time.After(2 * time.Second):
t.Fatal("workflow did not run")
}
if res.stopTail != nil {
res.stopTail()
}
e.jobsWG.Wait()
got := inner.setupNixCache
if got == nil || got.Cache != "plc-target" || got.PublicKey != lease.PublicKey || got.PullToken != "pull" || got.PushToken != tt.wantPush || got.MaxUploadBytes != 200 {
t.Fatalf("engine's cache = %+v, want plc-target with pull, push token %q and the upload cap", got, tt.wantPush)
}
})
}
}