package db import ( "context" "database/sql" "fmt" "time" "tangled.org/core/eventstream" "tangled.org/core/notifier" "tangled.org/core/spindle/models" ) // recovery uses persisted identity, fencing and quota state type MillLease struct { LeaseID string NodeID string Epoch string Engine string Knot string Rkey string Workflow string State string QuotaReservationID string OwnerDID string RepoDID string MillRecordsTerminalMetrics bool } type ExecutorCursor struct { NodeID string Epoch string AckedSeqno uint64 } type OutboxRow struct { Epoch string Seqno uint64 Payload []byte ByteSize int64 Control bool } type OutboxDeletion struct { Rows int64 Bytes int64 } func (d *DB) SaveMillLease(l MillLease) error { _, err := d.Exec( `insert into mill_leases ( lease_id, node_id, epoch, engine, knot, rkey, workflow, state, quota_reservation_id, owner_did, repo_did, mill_records_terminal_metrics ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) on conflict(lease_id) do update set state = excluded.state, quota_reservation_id = excluded.quota_reservation_id, owner_did = excluded.owner_did, repo_did = excluded.repo_did, mill_records_terminal_metrics = excluded.mill_records_terminal_metrics`, l.LeaseID, l.NodeID, l.Epoch, l.Engine, l.Knot, l.Rkey, l.Workflow, l.State, l.QuotaReservationID, l.OwnerDID, l.RepoDID, l.MillRecordsTerminalMetrics, ) return err } func (d *DB) DeleteMillLease(leaseID string) error { _, err := d.Exec(`delete from mill_leases where lease_id = ?`, leaseID) return err } func (d *DB) ListMillLeases() ([]MillLease, error) { rows, err := d.Query(` select lease_id, node_id, epoch, engine, knot, rkey, workflow, state, coalesce(quota_reservation_id, ''), coalesce(owner_did, ''), coalesce(repo_did, ''), mill_records_terminal_metrics from mill_leases `) if err != nil { return nil, err } defer rows.Close() var leases []MillLease for rows.Next() { var l MillLease if err := rows.Scan( &l.LeaseID, &l.NodeID, &l.Epoch, &l.Engine, &l.Knot, &l.Rkey, &l.Workflow, &l.State, &l.QuotaReservationID, &l.OwnerDID, &l.RepoDID, &l.MillRecordsTerminalMetrics, ); err != nil { return nil, err } leases = append(leases, l) } return leases, rows.Err() } func (d *DB) ListExecutorCursors() ([]ExecutorCursor, error) { rows, err := d.Query(`select node_id, epoch, acked_seqno from mill_executor_cursors`) if err != nil { return nil, err } defer rows.Close() var cursors []ExecutorCursor for rows.Next() { var c ExecutorCursor if err := rows.Scan(&c.NodeID, &c.Epoch, &c.AckedSeqno); err != nil { return nil, err } cursors = append(cursors, c) } return cursors, rows.Err() } func (d *DB) SetOutboxEpoch(epoch string) error { return retrySQLite(func() error { tx, err := d.Begin() if err != nil { return err } defer tx.Rollback() if _, err := tx.Exec(`delete from mill_outbox_rows`); err != nil { return err } if _, err := tx.Exec(`delete from mill_outbox_state`); err != nil { return err } if _, err := tx.Exec(`insert into mill_outbox_state (epoch, next_seqno) values (?, 1)`, epoch); err != nil { return err } return tx.Commit() }) } func (d *DB) AppendOutboxRow(payload []byte, control bool) (uint64, error) { var nextSeqno uint64 err := retrySQLite(func() error { tx, err := d.Begin() if err != nil { return err } defer tx.Rollback() var epoch string err = tx.QueryRow(` update mill_outbox_state set next_seqno = next_seqno + 1 where epoch = (select epoch from mill_outbox_state limit 1) returning epoch, next_seqno - 1 `).Scan(&epoch, &nextSeqno) if err == sql.ErrNoRows { return fmt.Errorf("no outbox epoch set") } else if err != nil { return err } byteSize := int64(len(payload)) controlVal := 0 if control { controlVal = 1 } if _, err := tx.Exec( `insert into mill_outbox_rows (epoch, seqno, payload, byte_size, control) values (?, ?, ?, ?, ?)`, epoch, nextSeqno, payload, byteSize, controlVal, ); err != nil { return err } return tx.Commit() }) if err != nil { return 0, err } return nextSeqno, nil } func (d *DB) DeleteOutboxPrefix(ackedSeqno uint64) (OutboxDeletion, error) { var deleted OutboxDeletion err := retrySQLite(func() error { tx, err := d.Begin() if err != nil { return err } defer tx.Rollback() var epoch string err = tx.QueryRow(`select epoch from mill_outbox_state limit 1`).Scan(&epoch) if err == sql.ErrNoRows { return nil } if err != nil { return err } deleted = OutboxDeletion{} if err := tx.QueryRow(` select count(*), coalesce(sum(byte_size), 0) from mill_outbox_rows where epoch = ? and seqno <= ? `, epoch, ackedSeqno).Scan(&deleted.Rows, &deleted.Bytes); err != nil { return err } if _, err := tx.Exec( `delete from mill_outbox_rows where epoch = ? and seqno <= ?`, epoch, ackedSeqno, ); err != nil { return err } return tx.Commit() }) if err != nil { return OutboxDeletion{}, err } return deleted, nil } func (d *DB) ListOutboxRows() ([]OutboxRow, error) { rows, err := d.Query(`select epoch, seqno, payload, byte_size, control from mill_outbox_rows order by seqno`) if err != nil { return nil, err } defer rows.Close() var out []OutboxRow for rows.Next() { var r OutboxRow var controlVal int if err := rows.Scan(&r.Epoch, &r.Seqno, &r.Payload, &r.ByteSize, &controlVal); err != nil { return nil, err } r.Control = (controlVal != 0) out = append(out, r) } return out, rows.Err() } func (d *DB) ListOutboxRowsAfter(seqno uint64, limit int) ([]OutboxRow, error) { rows, err := d.Query(` select epoch, seqno, payload, byte_size, control from mill_outbox_rows where epoch = (select epoch from mill_outbox_state limit 1) and seqno > ? order by seqno limit ? `, seqno, limit) if err != nil { return nil, err } defer rows.Close() var out []OutboxRow for rows.Next() { var row OutboxRow var control int if err := rows.Scan(&row.Epoch, &row.Seqno, &row.Payload, &row.ByteSize, &control); err != nil { return nil, err } row.Control = control != 0 out = append(out, row) } return out, rows.Err() } func (d *DB) GetOutboxState() (string, uint64, error) { var epoch string var nextSeqno uint64 err := d.QueryRow(`select epoch, next_seqno from mill_outbox_state limit 1`).Scan(&epoch, &nextSeqno) if err == sql.ErrNoRows { return "", 0, nil } return epoch, nextSeqno, err } type EventBatchTx struct { tx *sql.Tx db *DB insert eventstream.Inserter } func (tx *EventBatchTx) InsertStatusEvent(pipelineAtUri, workflow, status string, workflowError *string, exitCode *int64) error { event, err := statusEvent(pipelineAtUri, workflow, status, workflowError, exitCode) if err != nil { return err } return tx.insert(tx.tx, event) } func (tx *EventBatchTx) DeleteLease(leaseID string) error { _, err := tx.tx.Exec(`delete from mill_leases where lease_id = ?`, leaseID) return err } func (tx *EventBatchTx) AdvanceCursor(nodeID, epoch string, seqno uint64) error { _, err := tx.tx.Exec( `insert into mill_executor_cursors (node_id, epoch, acked_seqno) values (?, ?, ?) on conflict(node_id, epoch) do update set acked_seqno = excluded.acked_seqno where excluded.acked_seqno > mill_executor_cursors.acked_seqno`, nodeID, epoch, seqno, ) return err } // start from zero if a row is missing func (d *DB) GetExecutorCursor(nodeID, epoch string) (uint64, error) { var seqno uint64 err := d.QueryRow( `select acked_seqno from mill_executor_cursors where node_id = ? and epoch = ?`, nodeID, epoch, ).Scan(&seqno) if err == sql.ErrNoRows { return 0, nil } return seqno, err } // forgets everything the mill applied from this node, picked up on the // executor's next reconnect since cursors are re-read at session start func (d *DB) DeleteExecutorCursors(nodeID string) (int64, error) { res, err := d.Exec(`delete from mill_executor_cursors where node_id = ?`, nodeID) if err != nil { return 0, err } return res.RowsAffected() } // marks everything at or below seqno as applied, for skipping a backlog // the mill can never receive func (d *DB) SetExecutorCursors(nodeID string, seqno uint64) (int64, error) { res, err := d.Exec(`update mill_executor_cursors set acked_seqno = ? where node_id = ?`, seqno, nodeID) if err != nil { return 0, err } return res.RowsAffected() } // clears artifacts waiting on leases from the old stream, replaying them // under a new epoch would trip the mill's lease-epoch check func (d *DB) ClearPendingArtifacts() error { _, err := d.Exec(`delete from executor_pending_artifacts`) return err } func (tx *EventBatchTx) InsertArtifactRef(leaseID, repoDid string, wid models.WorkflowId, ref, hash string) error { if err := validateWorkflowIdentity(wid); err != nil { return err } _, err := tx.tx.Exec( `insert into mill_artifacts (lease_id, repo_did, knot, rkey, workflow, ref, hash) values (?, ?, ?, ?, ?, ?, ?)`, leaseID, repoDid, wid.Knot, wid.Rkey, wid.Name, ref, hash, ) return err } type PendingArtifact struct { LeaseID string Knot string Rkey string Workflow string Status string Error string ExitCode int64 Ref string Hash string FailureClass string FailureReason string MillRecordsTerminalMetrics bool } func (d *DB) SavePendingArtifact( leaseID string, wid models.WorkflowId, status, errStr string, exitCode int64, ref, hash, failureClass, failureReason string, millRecordsTerminalMetrics bool, ) error { _, err := d.Exec( `insert into executor_pending_artifacts ( lease_id, knot, rkey, workflow, status, error, exit_code, ref, hash, failure_class, failure_reason, mill_records_terminal_metrics ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) on conflict(lease_id) do update set knot = excluded.knot, rkey = excluded.rkey, workflow = excluded.workflow, status = excluded.status, error = excluded.error, exit_code = excluded.exit_code, ref = excluded.ref, hash = excluded.hash, failure_class = excluded.failure_class, failure_reason = excluded.failure_reason, mill_records_terminal_metrics = excluded.mill_records_terminal_metrics`, leaseID, wid.Knot, wid.Rkey, wid.Name, status, errStr, exitCode, ref, hash, failureClass, failureReason, millRecordsTerminalMetrics, ) return err } func (d *DB) RemovePendingArtifact(leaseID string) error { _, err := d.Exec(`delete from executor_pending_artifacts where lease_id = ?`, leaseID) return err } func (d *DB) ListPendingArtifacts() ([]PendingArtifact, error) { rows, err := d.Query(` select lease_id, knot, rkey, workflow, status, error, exit_code, ref, hash, failure_class, failure_reason, mill_records_terminal_metrics from executor_pending_artifacts `) if err != nil { return nil, err } defer rows.Close() var res []PendingArtifact for rows.Next() { var p PendingArtifact if err := rows.Scan( &p.LeaseID, &p.Knot, &p.Rkey, &p.Workflow, &p.Status, &p.Error, &p.ExitCode, &p.Ref, &p.Hash, &p.FailureClass, &p.FailureReason, &p.MillRecordsTerminalMetrics, ); err != nil { return nil, err } res = append(res, p) } return res, rows.Err() } // ApplyEventBatch runs fn in a transaction and retries it while sqlite reports // transient lock errors. fn may run more than once, so non-transactional side // effects must be returned through its result: the result is published only // from the attempt that commits, and a failed attempt's result is discarded. func ApplyEventBatch[T any](d *DB, n *notifier.Notifier, fn func(tx *EventBatchTx) (T, error)) (T, error) { return ApplyEventBatchCtx(context.Background(), d, n, fn) } func ApplyEventBatchCtx[T any](ctx context.Context, d *DB, n *notifier.Notifier, fn func(tx *EventBatchTx) (T, error)) (T, error) { var result T if err := ctx.Err(); err != nil { var zero T return zero, err } err := retrySQLiteCtx(ctx, 20*time.Second, func() error { if err := ctx.Err(); err != nil { return err } return eventstream.WithClock(func(insert eventstream.Inserter) error { tx, err := d.Begin() if err != nil { return err } defer tx.Rollback() batchTx := &EventBatchTx{ tx: tx, db: d, insert: insert, } v, err := fn(batchTx) if err != nil { return err } if err := tx.Commit(); err != nil { return err } result = v return nil }) }) if err != nil { var zero T return zero, err } if n != nil { n.NotifyAll() } return result, nil } // ApplyEventBatch is the package-level ApplyEventBatch for callers without a // result. func (d *DB) ApplyEventBatch(n *notifier.Notifier, fn func(tx *EventBatchTx) error) error { return d.ApplyEventBatchCtx(context.Background(), n, fn) } func (d *DB) ApplyEventBatchCtx(ctx context.Context, n *notifier.Notifier, fn func(tx *EventBatchTx) error) error { _, err := ApplyEventBatchCtx(ctx, d, n, func(tx *EventBatchTx) (struct{}, error) { return struct{}{}, fn(tx) }) return err } func (d *DB) DeleteMillLeasesByRepo(repoDid string) ([]string, error) { tx, err := d.Begin() if err != nil { return nil, err } defer tx.Rollback() rows, err := tx.Query(`select lease_id from mill_leases where repo_did = ?`, repoDid) if err != nil { return nil, err } var leaseIDs []string for rows.Next() { var id string if err := rows.Scan(&id); err != nil { rows.Close() return nil, err } leaseIDs = append(leaseIDs, id) } if err := rows.Err(); err != nil { rows.Close() return nil, err } rows.Close() for _, id := range leaseIDs { if _, err := tx.Exec(`delete from executor_pending_artifacts where lease_id = ?`, id); err != nil { return nil, err } } if _, err := tx.Exec(`delete from mill_leases where repo_did = ?`, repoDid); err != nil { return nil, err } // return ids only after the deletion commits if err := tx.Commit(); err != nil { return nil, err } return leaseIDs, nil } func (d *DB) ListArtifactRefsByRepo(repoDid string) ([]string, error) { rows, err := d.Query(` select ref from mill_artifacts where repo_did = ? union select pa.ref from executor_pending_artifacts pa join mill_leases ml on ml.lease_id = pa.lease_id where ml.repo_did = ?`, repoDid, repoDid, ) if err != nil { return nil, err } defer rows.Close() var refs []string for rows.Next() { var ref string if err := rows.Scan(&ref); err != nil { return nil, err } refs = append(refs, ref) } return refs, rows.Err() } func (d *DB) DeleteArtifactRefsByRepo(repoDid string) error { _, err := d.Exec(`delete from mill_artifacts where repo_did = ?`, repoDid) return err }