package mill

import (
	"bytes"
	"context"
	"errors"
	"io"
	"log/slog"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"tangled.org/core/api/tangled"
	"tangled.org/core/spindle/engine"
	"tangled.org/core/spindle/models"

	millproto "tangled.org/core/spindle/mill/proto"
	millv1 "tangled.org/core/spindle/mill/proto/gen"
)

type scriptedEncoder func(*millproto.Message) error

func (e scriptedEncoder) Encode(msg *millproto.Message) error { return e(msg) }

func testWorkflow(name string) *models.Workflow {
	return &models.Workflow{
		Name:        name,
		Environment: map[string]string{},
		Steps:       []models.Step{remoteStep{}},
		Data: &millWorkflowState{
			TargetEngine: "dummy",
			RawWorkflow:  tangled.Pipeline_Workflow{Name: name},
			RawPipeline:  tangled.Pipeline{},
		},
	}
}

func testWorkflowWithRunsOn(name string, runsOn []string) *models.Workflow {
	wf := testWorkflow(name)
	wf.Data.(*millWorkflowState).RawWorkflow.RunsOn = runsOn
	return wf
}

func addCandidateSession(t *testing.T, m *Mill, nodeID string, labels []string, load float64, enc messageEncoder) *millSession {
	t.Helper()
	if enc == nil {
		enc = scriptedEncoder(func(*millproto.Message) error { return nil })
	}
	sess := newSession(nodeID, enc, slog.New(slog.NewTextHandler(io.Discard, nil)))
	sess.labels = labels
	sess.snapshot = &millv1.NodeSnapshot{
		NodeId: nodeID,
		Engines: map[string]*millv1.EngineAvailability{
			"dummy": {Available: load < 1.0, Load: map[string]float64{"slots": load}},
		},
	}
	m.mu.Lock()
	m.sessions[nodeID] = sess
	m.mu.Unlock()
	return sess
}

func assertRankedNodes(t *testing.T, got []*millSession, want []string) {
	t.Helper()
	if len(got) != len(want) {
		t.Fatalf("rankCandidates() returned %d candidates, want %d: got %v want %v", len(got), len(want), sessionIDs(got), want)
	}
	for i := range want {
		if got[i].nodeID != want[i] {
			t.Fatalf("rankCandidates()[%d] = %q, want %q; full order got %v want %v", i, got[i].nodeID, want[i], sessionIDs(got), want)
		}
	}
}

func sessionIDs(sessions []*millSession) []string {
	out := make([]string, len(sessions))
	for i, sess := range sessions {
		out[i] = sess.nodeID
	}
	return out
}

func sameStringMultiset(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	counts := make(map[string]int, len(a))
	for _, s := range a {
		counts[s]++
	}
	for _, s := range b {
		if counts[s] == 0 {
			return false
		}
		counts[s]--
	}
	return true
}

type reserveReply struct {
	accepted    bool
	rejectClass millv1.RejectClass
	reason      string
}

func addReplyingCandidateSession(t *testing.T, m *Mill, nodeID string, labels []string, load float64, asked chan<- string, reply reserveReply) *millSession {
	t.Helper()
	var sess *millSession
	sess = addCandidateSession(t, m, nodeID, labels, load, scriptedEncoder(func(msg *millproto.Message) error {
		rs := msg.GetReserveSeat()
		if rs == nil {
			return nil
		}
		if asked != nil {
			asked <- nodeID
		}
		sess.deliver(rs.GetLeaseId(), &millproto.Message{ReserveResult: &millv1.ReserveResult{
			LeaseId:      rs.GetLeaseId(),
			Accepted:     reply.accepted,
			RejectReason: reply.reason,
			RejectClass:  reply.rejectClass,
		}})
		return nil
	}))
	return sess
}

func drainAsked(ch <-chan string) []string {
	var out []string
	for {
		select {
		case nodeID := <-ch:
			out = append(out, nodeID)
		default:
			return out
		}
	}
}

func TestCommitRetriesAfterSessionCloseBeforeCommitted(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{BidTimeout: 25 * time.Millisecond, ReconnectGrace: time.Second})
	wf := testWorkflow("build")
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "dummy")
	lease.wid = wid
	wf.Data.(*millWorkflowState).Lease = lease

	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	var sess1 *millSession
	firstCommit := make(chan struct{})
	sess1 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetCommitLease() != nil {
			close(firstCommit)
			m.detachSession(sess1)
		}
		return nil
	}), l)
	m.attachSession(sess1)

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	done := make(chan error, 1)
	go func() { done <- m.commitAndWait(ctx, wf, nil) }()

	select {
	case <-firstCommit:
	case <-ctx.Done():
		t.Fatal("first commit was not sent")
	}

	var sess2 *millSession
	sess2 = newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetCommitLease() == nil {
			return nil
		}
		leaseID := msg.GetCommitLease().GetLeaseId()
		sess2.deliver(leaseID, &millproto.Message{Committed: &millv1.Committed{LeaseId: leaseID}})
		m.onAttemptResult(sess2, &millv1.AttemptResult{
			Offset:         1,
			LeaseId:        leaseID,
			TerminalStatus: string(models.StatusKindSuccess),
		})
		return nil
	}), l)
	m.attachSession(sess2)
	m.sessionReady(sess2)

	select {
	case err := <-done:
		if err != nil {
			t.Fatalf("commitAndWait() error = %v, want success after reconnect", err)
		}
	case <-ctx.Done():
		t.Fatal("commitAndWait() did not finish after reconnect")
	}
}

func TestDestroyRunningLeaseDoesNotDropCancelledTerminal(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "dummy")
	lease.wid = wid
	lease.setState(leaseRunning)

	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(wid)
	if lease.getState() == leaseDone {
		t.Fatal("destroy sealed the lease before the terminal result")
	}

	lease.deliverTerminal(&millv1.AttemptResult{
		LeaseId:        lease.id,
		TerminalStatus: string(models.StatusKindCancelled),
	})
	res := <-lease.terminal
	if err := terminalError(res.GetTerminalStatus()); !errors.Is(err, engine.ErrCancelled) {
		t.Fatalf("terminalError() = %v, want ErrCancelled", err)
	}
}

func TestPlaceBlocksWhenNoCapacity(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{})
	wf := testWorkflow("build")
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}

	// no executors at all: place must block until ctx expires (user sees pending).
	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
	defer cancel()

	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded", err)
	}
}

func TestRankCandidatesFiltersRequiredLabelsWithANDSemantics(t *testing.T) {
	m := New(slog.New(slog.NewTextHandler(io.Discard, nil)), Config{})
	addCandidateSession(t, m, "linux-high", []string{"linux"}, 0.0, nil)
	addCandidateSession(t, m, "linux-arm", []string{"linux", "arm64"}, 0.25, nil)
	addCandidateSession(t, m, "unlabeled", nil, 0.5, nil)
	addCandidateSession(t, m, "linux-arm-gpu", []string{"linux", "arm64", "gpu"}, 0.75, nil)
	addCandidateSession(t, m, "linux-arm-full", []string{"linux", "arm64"}, 1.0, nil)

	tests := []struct {
		name           string
		requiredLabels []string
		want           []string
	}{
		{
			name: "no required labels keeps old capacity ranking",
			want: []string{"linux-high", "linux-arm", "unlabeled", "linux-arm-gpu"},
		},
		{
			name:           "single required label includes every candidate carrying it",
			requiredLabels: []string{"linux"},
			want:           []string{"linux-high", "linux-arm", "linux-arm-gpu"},
		},
		{
			name:           "all required labels must be present",
			requiredLabels: []string{"linux", "arm64"},
			want:           []string{"linux-arm", "linux-arm-gpu"},
		},
		{
			name:           "one missing required label excludes the candidate",
			requiredLabels: []string{"linux", "arm64", "gpu"},
			want:           []string{"linux-arm-gpu"},
		},
		{
			name:           "unknown required label leaves no candidate",
			requiredLabels: []string{"linux", "arm64", "metal"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			assertRankedNodes(t, m.rankCandidates("dummy", tt.requiredLabels, nil), tt.want)
		})
	}
}

func TestPlaceWithMissingRequiredLabelsStaysPendingWithoutReserve(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{BidTimeout: 10 * time.Millisecond})
	reserveSent := make(chan struct{}, 1)
	addCandidateSession(t, m, "linux-only", []string{"linux"}, 0.75, scriptedEncoder(func(msg *millproto.Message) error {
		if msg.GetReserveSeat() != nil {
			select {
			case reserveSent <- struct{}{}:
			default:
			}
		}
		return nil
	}))
	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}

	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Millisecond)
	defer cancel()
	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded while job remains pending", err)
	}
	select {
	case <-reserveSent:
		t.Fatal("place() sent ReserveSeat to executor missing a required label")
	default:
	}
}

func TestPlaceKeepsAskingBelowTopKAfterIncompatibleRejects(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{TopK: 2, BidTimeout: time.Second})
	asked := make(chan string, 4)
	addReplyingCandidateSession(t, m, "wrong-label", []string{"linux"}, 0.0, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "wrong-label should not be asked",
	})
	addReplyingCandidateSession(t, m, "incompatible-a", []string{"linux", "arm64"}, 0.0, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "no runner",
	})
	addReplyingCandidateSession(t, m, "incompatible-b", []string{"linux", "arm64"}, 0.25, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "bad image",
	})
	addReplyingCandidateSession(t, m, "compatible-below-window", []string{"linux", "arm64"}, 0.5, asked, reserveReply{
		accepted: true,
	})

	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	slot, err := m.place(ctx, "dummy", wid, wf)
	if err != nil {
		t.Fatalf("place() error = %v, want compatible lower-ranked executor", err)
	}
	defer slot.Release()

	ms, ok := slot.(*millSlot)
	if !ok {
		t.Fatalf("place() slot type = %T, want *millSlot", slot)
	}
	if ms.lease.nodeID != "compatible-below-window" {
		t.Fatalf("place() chose node %q, want compatible-below-window", ms.lease.nodeID)
	}
	if got, want := drainAsked(asked), []string{"incompatible-a", "incompatible-b", "compatible-below-window"}; !sameStringMultiset(got, want) {
		t.Fatalf("ReserveSeat asked nodes = %v, want %v", got, want)
	}
}

func TestPlaceKeepsTransientOnlyRejectsPending(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{TopK: 1, BidTimeout: 10 * time.Millisecond})
	asked := make(chan string, 2)
	addReplyingCandidateSession(t, m, "busy-a", []string{"linux", "arm64"}, 0.5, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_TRANSIENT,
		reason:      "draining",
	})
	addReplyingCandidateSession(t, m, "busy-b", []string{"linux", "arm64"}, 0.75, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_TRANSIENT,
		reason:      "no slot",
	})

	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Millisecond)
	defer cancel()
	_, err := m.place(ctx, "dummy", wid, wf)
	if err != context.DeadlineExceeded {
		t.Fatalf("place() error = %v, want DeadlineExceeded while transient rejects leave job pending", err)
	}
	if got, want := drainAsked(asked), []string{"busy-a", "busy-b"}; !sameStringMultiset(got, want) {
		t.Fatalf("ReserveSeat asked nodes = %v, want %v", got, want)
	}
}

func TestPlaceReportsOnlyEligiblePermanentIncompatibleRejects(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{TopK: 1, BidTimeout: time.Second})
	asked := make(chan string, 3)
	addReplyingCandidateSession(t, m, "wrong-label", []string{"linux"}, 0.25, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "wrong-label should not be asked",
	})
	addReplyingCandidateSession(t, m, "incompatible-a", []string{"linux", "arm64"}, 0.5, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "no qemu-system-aarch64",
	})
	addReplyingCandidateSession(t, m, "incompatible-b", []string{"linux", "arm64"}, 0.75, asked, reserveReply{
		accepted:    false,
		rejectClass: millv1.RejectClass_REJECT_CLASS_INCOMPATIBLE,
		reason:      "image arch unsupported",
	})

	wf := testWorkflowWithRunsOn("build", []string{"linux", "arm64"})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	_, err := m.place(ctx, "dummy", wid, wf)
	if err == nil {
		t.Fatal("place() error = nil, want permanent incompatible failure")
	}
	errText := err.Error()
	for _, want := range []string{"no compatible executor for dummy", "no qemu-system-aarch64", "image arch unsupported"} {
		if !strings.Contains(errText, want) {
			t.Fatalf("place() error = %q, want it to contain %q", errText, want)
		}
	}
	if strings.Contains(errText, "wrong-label should not be asked") {
		t.Fatalf("place() error = %q, included a missing-label candidate as incompatible", errText)
	}
	if got, want := drainAsked(asked), []string{"incompatible-a", "incompatible-b"}; !sameStringMultiset(got, want) {
		t.Fatalf("ReserveSeat asked nodes = %v, want %v", got, want)
	}
}

func TestMaxPendingRejects(t *testing.T) {
	l := slog.New(slog.NewTextHandler(io.Discard, nil))
	m := New(l, Config{MaxPending: 1})

	m.mu.Lock()
	m.pending = 1
	m.mu.Unlock()

	wf2 := testWorkflow("b")
	_, err := m.place(context.Background(), "dummy", models.WorkflowId{Name: "b"}, wf2)
	if err == nil {
		t.Fatal("place() past maxPending should error")
	}
}

func TestCancelledRunningLeaseSurvivesReleaseForReconnectReplay(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{ReconnectGrace: time.Minute})
	wid := models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	lease := newLease("lease-1", "node-1", "dummy")
	lease.wid = wid
	lease.setState(leaseRunning)
	if err := m.persistLease(lease, leaseRowRunning); err != nil {
		t.Fatalf("persistLease: %v", err)
	}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()

	m.destroy(wid)
	slot := &millSlot{fleet: m, lease: lease}
	slot.Release()

	m.mu.Lock()
	_, retained := m.leases[lease.id]
	m.mu.Unlock()
	if !retained {
		t.Fatal("slot release removed a cancellation-requested running lease before its terminal result")
	}
	if rows, err := bdb.ListMillLeases(); err != nil || len(rows) != 1 {
		t.Fatalf("durable leases after slot release = %+v, err = %v; want retained lease", rows, err)
	}

	var sentMu sync.Mutex
	var sent []*millproto.Message
	sess := newSession("node-1", scriptedEncoder(func(msg *millproto.Message) error {
		sentMu.Lock()
		sent = append(sent, msg)
		sentMu.Unlock()
		return nil
	}), discardLogger())
	if _, ok := m.attachSession(sess); !ok {
		t.Fatal("attachSession rejected reconnect")
	}
	m.sessionReady(sess)
	sentMu.Lock()
	var replayed bool
	for _, msg := range sent {
		if cancel := msg.GetCancelAttempt(); cancel != nil && cancel.GetLeaseId() == lease.id {
			replayed = true
		}
	}
	sentMu.Unlock()
	if !replayed {
		t.Fatal("reconnect did not replay CancelAttempt for retained lease")
	}

	if err := m.onAttemptResult(sess, &millv1.AttemptResult{
		Offset:         1,
		LeaseId:        lease.id,
		TerminalStatus: string(models.StatusKindCancelled),
	}); err != nil {
		t.Fatalf("onAttemptResult: %v", err)
	}
	m.mu.Lock()
	_, retained = m.leases[lease.id]
	m.mu.Unlock()
	if retained {
		t.Fatal("terminal result did not clean retained cancelled lease")
	}
	if rows, err := bdb.ListMillLeases(); err != nil || len(rows) != 0 {
		t.Fatalf("durable leases after terminal = %+v, err = %v; want none", rows, err)
	}
	slot.Release()
}

func TestRelaySideEffectFailureStopsReadLoopAtFailedOffset(t *testing.T) {
	m := New(discardLogger(), Config{LogDir: filepath.Join(t.TempDir(), "missing")})
	lease := newLease("lease-1", "node-1", "dummy")
	lease.wid = models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"}
	m.mu.Lock()
	m.leases[lease.id] = lease
	m.mu.Unlock()
	sess := newSession("node-1", nopEncoder(), discardLogger())
	if _, ok := m.attachSession(sess); !ok {
		t.Fatal("attachSession rejected first session")
	}

	var frames bytes.Buffer
	enc := millproto.NewEncoder(&frames)
	if err := enc.Encode(&millproto.Message{LogLine: &millv1.LogLine{
		Offset: 1, LeaseId: lease.id, RawJson: []byte(`{"line":"first"}`),
	}}); err != nil {
		t.Fatalf("encode failing log relay: %v", err)
	}
	if err := enc.Encode(&millproto.Message{AttemptResult: &millv1.AttemptResult{
		Offset: 2, LeaseId: lease.id, TerminalStatus: string(models.StatusKindSuccess),
	}}); err != nil {
		t.Fatalf("encode later relay: %v", err)
	}

	err := sess.readLoop(m, millproto.NewDecoder(&frames))
	if err == nil {
		t.Fatal("readLoop continued after relayed log side effect failed")
	}
	m.mu.Lock()
	offset := m.nodeOffset[sess.nodeID]
	m.mu.Unlock()
	if offset != 0 {
		t.Fatalf("in-memory relay offset = %d, want 0 so reconnect retries offset 1", offset)
	}
	if lease.getState() == leaseDone {
		t.Fatal("readLoop applied offset 2 after offset 1 failed")
	}
}

func TestFallbackBidGetsFullTimeout(t *testing.T) {
	m := New(discardLogger(), Config{TopK: 1, BidTimeout: 30 * time.Millisecond})
	addCandidateSession(t, m, "silent", nil, 0, nil)
	var fallback *millSession
	fallback = addCandidateSession(t, m, "fallback", nil, 0.5, scriptedEncoder(func(msg *millproto.Message) error {
		if reserve := msg.GetReserveSeat(); reserve != nil {
			leaseID := reserve.GetLeaseId()
			time.AfterFunc(10*time.Millisecond, func() {
				fallback.deliver(leaseID, &millproto.Message{ReserveResult: &millv1.ReserveResult{
					LeaseId:  leaseID,
					Accepted: true,
				}})
			})
		}
		return nil
	}))

	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	lease, err := m.bid(ctx, "dummy", testWorkflow("build"))
	if err != nil {
		t.Fatalf("bid: %v", err)
	}
	if lease == nil || lease.nodeID != "fallback" {
		t.Fatalf("bid winner = %+v, want fallback after silent incumbent times out", lease)
	}
}

func TestSessionRequestCancelledBeforeRegistrationOrSend(t *testing.T) {
	sent := 0
	sess := newSession("node-1", scriptedEncoder(func(*millproto.Message) error {
		sent++
		return nil
	}), discardLogger())
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := sess.request(ctx, "lease-1", &millproto.Message{ReleaseLease: &millv1.ReleaseLease{LeaseId: "lease-1"}})
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("request error = %v, want context.Canceled", err)
	}
	if sent != 0 {
		t.Fatalf("request sent %d messages for an already-cancelled context, want 0", sent)
	}
	sess.mu.Lock()
	pending := len(sess.pending)
	sess.mu.Unlock()
	if pending != 0 {
		t.Fatalf("request left %d pending waiters, want 0", pending)
	}
}

func TestAttachSessionReplacesSilentIncumbentButRejectsActiveDuplicate(t *testing.T) {
	m := New(discardLogger(), Config{ReconnectGrace: time.Minute})
	old := newSession("node-1", nopEncoder(), discardLogger())
	transportClosed := make(chan struct{})
	old.closeTransport = func() error {
		close(transportClosed)
		return nil
	}
	if _, ok := m.attachSession(old); !ok {
		t.Fatal("first attach rejected")
	}
	m.mu.Lock()
	old.lastSeen = time.Now().Add(-2 * m.cfg.ReconnectGrace)
	m.mu.Unlock()

	replacement := newSession("node-1", nopEncoder(), discardLogger())
	if _, ok := m.attachSession(replacement); !ok {
		t.Fatal("silent incumbent blocked authenticated replacement")
	}
	select {
	case <-transportClosed:
	default:
		t.Fatal("replacing a silent incumbent did not close its transport")
	}
	m.mu.Lock()
	replacement.lastSeen = time.Now().Add(-2 * m.cfg.ReconnectGrace)
	m.mu.Unlock()
	if err := replacement.dispatch(m, &millproto.Message{NodeSnapshot: &millv1.NodeSnapshot{NodeId: "node-1"}}); err != nil {
		t.Fatalf("periodic snapshot dispatch: %v", err)
	}
	if _, ok := m.attachSession(newSession("node-1", nopEncoder(), discardLogger())); ok {
		t.Fatal("active replacement did not reject a duplicate session")
	}
}

func TestPlaceReleasesRemoteReservationWhenInitialPersistenceFails(t *testing.T) {
	m, bdb := restoreTestMill(t, Config{BidTimeout: time.Second})
	if err := bdb.Close(); err != nil {
		t.Fatalf("close db: %v", err)
	}
	released := make(chan string, 1)
	var sess *millSession
	sess = addCandidateSession(t, m, "node-1", nil, 0, scriptedEncoder(func(msg *millproto.Message) error {
		switch {
		case msg.GetReserveSeat() != nil:
			leaseID := msg.GetReserveSeat().GetLeaseId()
			sess.deliver(leaseID, &millproto.Message{ReserveResult: &millv1.ReserveResult{
				LeaseId:  leaseID,
				Accepted: true,
			}})
		case msg.GetReleaseLease() != nil:
			released <- msg.GetReleaseLease().GetLeaseId()
		}
		return nil
	}))

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	slot, err := m.place(
		ctx,
		"dummy",
		models.WorkflowId{PipelineId: models.PipelineId{Knot: "k", Rkey: "r"}, Name: "build"},
		testWorkflow("build"),
	)
	if err == nil {
		t.Fatal("place succeeded after reserved lease persistence failed")
	}
	if slot != nil {
		t.Fatalf("place returned slot %T after persistence failure", slot)
	}
	select {
	case leaseID := <-released:
		if leaseID == "" {
			t.Fatal("ReleaseLease had empty lease id")
		}
	case <-time.After(time.Second):
		t.Fatal("persistence failure did not compensate with ReleaseLease")
	}
	m.mu.Lock()
	leases := len(m.leases)
	m.mu.Unlock()
	if leases != 0 {
		t.Fatalf("mill published %d leases after initial persistence failure, want 0", leases)
	}
}
