Harden lease lifecycle durability

This commit is contained in:
kami
2026-07-30 14:34:29 +04:00
parent 1ff0af2e69
commit f6ee0e3060
40 changed files with 2108 additions and 590 deletions
+95 -24
View File
@@ -36,6 +36,7 @@ type worker struct {
sessions map[string]herdr.Session
leases map[string]lease
releases map[string]releaseTransaction
quarantined map[string]bool
statePath string
hard float64
registration federation.Worker
@@ -79,6 +80,7 @@ func (w *worker) health(ctx context.Context) federation.WorkerHealth {
}
type lease struct {
Epoch string `json:"epoch"`
HandoffRef string `json:"handoff_ref,omitempty"`
TransactionID string `json:"transaction_id,omitempty"`
AnchorSHA string `json:"anchor_sha,omitempty"`
@@ -117,24 +119,29 @@ type completionEvidence struct {
CompletedAt time.Time `json:"completed_at"`
}
type workerState struct {
Cursor uint64 `json:"cursor"`
Sessions map[string]herdr.Session `json:"sessions"`
Tasks map[string]domain.Task `json:"tasks"`
Leases map[string]lease `json:"leases"`
Releases map[string]releaseTransaction `json:"releases"`
Cursor uint64 `json:"cursor"`
Sessions map[string]herdr.Session `json:"sessions"`
Tasks map[string]domain.Task `json:"tasks"`
Leases map[string]lease `json:"leases"`
Releases map[string]releaseTransaction `json:"releases"`
Quarantined map[string]bool `json:"quarantined,omitempty"`
}
func (w *worker) load() {
func (w *worker) load() error {
b, e := os.ReadFile(w.statePath)
if e == nil {
var s workerState
if json.Unmarshal(b, &s) == nil {
w.cursor = s.Cursor
w.sessions = s.Sessions
w.tasks = s.Tasks
w.leases = s.Leases
w.releases = s.Releases
if err := json.Unmarshal(b, &s); err != nil {
return fmt.Errorf("corrupt worker state %s: %w", w.statePath, err)
}
w.cursor = s.Cursor
w.sessions = s.Sessions
w.tasks = s.Tasks
w.leases = s.Leases
w.releases = s.Releases
w.quarantined = s.Quarantined
} else if !errors.Is(e, os.ErrNotExist) {
return fmt.Errorf("read worker state %s: %w", w.statePath, e)
}
if w.sessions == nil {
w.sessions = map[string]herdr.Session{}
@@ -148,16 +155,68 @@ func (w *worker) load() {
if w.releases == nil {
w.releases = map[string]releaseTransaction{}
}
if w.quarantined == nil {
w.quarantined = map[string]bool{}
}
return nil
}
func (w *worker) save() error {
b, e := json.Marshal(workerState{Cursor: w.cursor, Sessions: w.sessions, Tasks: w.tasks, Leases: w.leases, Releases: w.releases})
b, e := json.Marshal(workerState{Cursor: w.cursor, Sessions: w.sessions, Tasks: w.tasks, Leases: w.leases, Releases: w.releases, Quarantined: w.quarantined})
if e != nil {
return e
}
if e := os.MkdirAll(filepath.Dir(w.statePath), 0700); e != nil {
return e
}
return os.WriteFile(w.statePath, b, 0600)
tmp := w.statePath + ".tmp"
f, e := os.OpenFile(tmp, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
if e != nil {
return e
}
if _, e = f.Write(b); e == nil {
e = f.Sync()
}
if closeErr := f.Close(); e == nil {
e = closeErr
}
if e != nil {
_ = os.Remove(tmp)
return e
}
if e = os.Rename(tmp, w.statePath); e != nil {
return e
}
dir, e := os.Open(filepath.Dir(w.statePath))
if e != nil {
return e
}
defer dir.Close()
return dir.Sync()
}
// quarantine stops a pane before its lost lease mapping can be forgotten.
// A failed close remains durable and is retried; it is never treated as a
// harmless cleanup error while the old harness could still be working.
func (w *worker) quarantine(ctx context.Context, taskID string, s herdr.Session) {
if w.herdr == nil {
w.quarantined[taskID] = true
return
}
if err := (herdr.CLIAdapter{Client: w.herdr, Harness: w.harness}).Kill(ctx, s); err != nil {
w.quarantined[taskID] = true
w.recordError(fmt.Errorf("quarantine %s: %w", taskID, err))
return
}
delete(w.sessions, taskID)
delete(w.quarantined, taskID)
}
func (w *worker) retryQuarantines(ctx context.Context) {
for taskID := range w.quarantined {
if s, ok := w.sessions[taskID]; ok {
w.quarantine(ctx, taskID, s)
}
}
}
type artifactCAS struct{ api federation.Client }
@@ -286,6 +345,9 @@ func taskHash(t domain.Task) string { b := continuity.RenderTaskFile(t); return
func (w *worker) releaseReady(ctx context.Context) {
for id, s := range w.sessions {
if w.quarantined[id] {
continue
}
if l := w.leases[id]; l.HandoffRef != "" && !l.PickupAcknowledged {
if err := w.ackPickup(ctx, id, s); err != nil {
w.recordError(err)
@@ -306,7 +368,7 @@ func (w *worker) releaseReady(ctx context.Context) {
log.Printf("upload completion %s: %v", id, err)
continue
}
if err = w.api.Complete(ctx, id, ref, evidence.ResultSHA, evidence.Branch, evidence.Remote, w.leases[id].Version, w.usageReceipt(s), w.sessionEvidence(ctx, id, s)); err != nil {
if err = w.api.Complete(ctx, id, ref, evidence.ResultSHA, evidence.Branch, evidence.Remote, w.leases[id].Epoch, w.leases[id].Version, w.usageReceipt(s), w.sessionEvidence(ctx, id, s)); err != nil {
w.recordError(fmt.Errorf("complete %s: %w", id, err))
log.Printf("complete %s: %v", id, err)
continue
@@ -439,7 +501,8 @@ func (w *worker) advanceRelease(ctx context.Context, id string, s herdr.Session)
_ = w.save()
}
if tx.Phase == "anchor_pushed" {
if err := w.api.Release(ctx, id, tx.Ref, tx.AnchorSHA, tx.ID, tx.LeaseVersion, w.sessionEvidence(ctx, id, s)); err != nil {
l := w.leases[id]
if err := w.api.Release(ctx, id, tx.Ref, tx.AnchorSHA, tx.ID, l.Epoch, tx.LeaseVersion, w.sessionEvidence(ctx, id, s)); err != nil {
tx.LastError, tx.UpdatedAt = err.Error(), time.Now().UTC()
w.releases[id] = tx
_ = w.save()
@@ -491,7 +554,7 @@ func (w *worker) ackPickup(ctx context.Context, id string, s herdr.Session) erro
if l.PickupAcknowledged {
return nil
}
if err := w.api.Pickup(ctx, id, l.HandoffRef, l.AnchorSHA, l.TransactionID, l.Version, w.sessionEvidence(ctx, id, s)); err != nil {
if err := w.api.Pickup(ctx, id, l.HandoffRef, l.AnchorSHA, l.TransactionID, l.Epoch, l.Version, w.sessionEvidence(ctx, id, s)); err != nil {
return fmt.Errorf("pickup %s acknowledgement: %w", id, err)
}
l.PickupAcknowledged = true
@@ -617,7 +680,7 @@ func (w *worker) renewLeases(ctx context.Context) {
w.recordError(fmt.Errorf("validate lease %s: %w", taskID, err))
continue
}
if err := w.api.Renew(ctx, taskID, l.Version, int((30 * time.Minute).Seconds())); err != nil {
if err := w.api.Renew(ctx, taskID, l.Epoch, l.Version, int((30 * time.Minute).Seconds())); err != nil {
w.recordError(fmt.Errorf("renew lease %s: %w", taskID, err))
log.Printf("renew lease %s: %v", taskID, err)
} else {
@@ -726,6 +789,7 @@ func (w *worker) once(ctx context.Context) error {
if e.Type == "TaskLeased" {
var p struct {
HarnessID string `json:"harness_id"`
Epoch string `json:"lease_epoch"`
HandoffRef string `json:"handoff_ref"`
TransactionID string `json:"transaction_id"`
AnchorSHA string `json:"anchor_sha"`
@@ -735,17 +799,18 @@ func (w *worker) once(ctx context.Context) error {
UntilNS int64 `json:"until_ns"`
}
_ = json.Unmarshal(e.Payload, &until)
w.leases[e.TaskID] = lease{HandoffRef: p.HandoffRef, TransactionID: p.TransactionID, AnchorSHA: p.AnchorSHA, Version: e.Version, Until: time.Unix(0, until.UntilNS)}
w.leases[e.TaskID] = lease{Epoch: p.Epoch, HandoffRef: p.HandoffRef, TransactionID: p.TransactionID, AnchorSHA: p.AnchorSHA, Version: e.Version, Until: time.Unix(0, until.UntilNS)}
}
}
if e.Type == "TaskLeaseRenewed" {
var p struct {
HarnessID string `json:"harness_id"`
Epoch string `json:"lease_epoch"`
UntilNS int64 `json:"until_ns"`
}
if json.Unmarshal(e.Payload, &p) == nil && p.HarnessID == w.harnessID {
l := w.leases[e.TaskID]
l.Version, l.Until = e.Version, time.Unix(0, p.UntilNS)
l.Version, l.Epoch, l.Until = e.Version, p.Epoch, time.Unix(0, p.UntilNS)
w.leases[e.TaskID] = l
}
}
@@ -755,6 +820,7 @@ func (w *worker) once(ctx context.Context) error {
if w.herdr == nil {
delete(w.sessions, e.TaskID)
} else if err := (herdr.CLIAdapter{Client: w.herdr, Harness: w.harness}).Kill(ctx, session); err != nil {
w.quarantined[e.TaskID] = true
w.recordError(fmt.Errorf("close completed pane %s: %w", e.TaskID, err))
continue
} else {
@@ -796,7 +862,9 @@ func (w *worker) once(ctx context.Context) error {
// TaskPickupValidated for its transaction. Do not erase its pane
// mapping merely because our own release event was replayed.
if _, releasing := w.releases[e.TaskID]; !releasing {
delete(w.sessions, e.TaskID)
if session, active := w.sessions[e.TaskID]; active {
w.quarantine(ctx, e.TaskID, session)
}
}
}
if e.Seq > w.cursor {
@@ -848,6 +916,7 @@ func (w *worker) once(ctx context.Context) error {
w.runCommands(ctx)
w.renewLeases(ctx)
}
w.retryQuarantines(ctx)
w.releaseReady(ctx)
if err := w.save(); err != nil {
return err
@@ -864,7 +933,7 @@ func (w *worker) reconcileLeases(ctx context.Context) error {
for _, task := range tasks {
w.tasks[task.ID] = task
if task.State == domain.StateLeased && task.Lease != nil && task.Lease.HarnessID == w.harnessID {
active[task.ID] = lease{HandoffRef: task.HandoffRef, TransactionID: task.ReleaseTransaction, AnchorSHA: task.ReleaseAnchor, Version: task.Version, Until: task.Lease.Until}
active[task.ID] = lease{Epoch: task.Lease.Epoch, HandoffRef: task.HandoffRef, TransactionID: task.ReleaseTransaction, AnchorSHA: task.ReleaseAnchor, Version: task.Version, Until: task.Lease.Until}
}
}
for taskID := range w.leases {
@@ -951,11 +1020,13 @@ func main() {
if v, err := strconv.ParseInt(os.Getenv("ORCHESTRA_CONTEXT_WINDOW"), 10, 64); err == nil && v > 0 {
window = v
}
w := &worker{api: federation.Client{BaseURL: required("ORCHESTRA_URL"), WorkerID: id, Token: token, AdmitToken: os.Getenv("ORCHESTRA_FEDERATION_ADMIT_TOKEN")}, harnessID: required("ORCHESTRA_WORKER_HERDR_ID"), harness: required("ORCHESTRA_WORKER_HARNESS"), repo: repo, root: root, remote: remote, projects: projects, tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{}, leases: map[string]lease{}, releases: map[string]releaseTransaction{}, statePath: os.Getenv("ORCHESTRA_WORKER_STATE"), hard: hard, soft: soft, window: window, registration: federation.Worker{ID: id, Address: os.Getenv("ORCHESTRA_WORKER_ADDRESS"), Capacity: 1, SupportedProjects: supported, Build: buildinfo.Current()}}
w := &worker{api: federation.Client{BaseURL: required("ORCHESTRA_URL"), WorkerID: id, Token: token, AdmitToken: os.Getenv("ORCHESTRA_FEDERATION_ADMIT_TOKEN")}, harnessID: required("ORCHESTRA_WORKER_HERDR_ID"), harness: required("ORCHESTRA_WORKER_HARNESS"), repo: repo, root: root, remote: remote, projects: projects, tasks: map[string]domain.Task{}, sessions: map[string]herdr.Session{}, leases: map[string]lease{}, releases: map[string]releaseTransaction{}, quarantined: map[string]bool{}, statePath: os.Getenv("ORCHESTRA_WORKER_STATE"), hard: hard, soft: soft, window: window, registration: federation.Worker{ID: id, Address: os.Getenv("ORCHESTRA_WORKER_ADDRESS"), Capacity: 1, SupportedProjects: supported, Build: buildinfo.Current()}}
if w.statePath == "" {
w.statePath = filepath.Join(w.root, ".orchestra-worker-state.json")
}
w.load()
if err := w.load(); err != nil {
log.Fatal(err)
}
w.herdr = herdr.New(required("ORCHESTRA_WORKER_HERDR"))
if id != w.harnessID {
log.Fatal("ORCHESTRA_WORKER_ID must equal ORCHESTRA_WORKER_HERDR_ID so leases and offline recovery have one owner")