dcache: batch job and target tables with a sliding lease, sweep and byte budget

Add the batch job store to internal/dcache (storage only; no HTTP, no worker). New internal/dcache/batch.go holds the BatchJob/BatchTarget types and all accessors; Open gains the batch_job/batch_target tables and indices, and DB gains BatchH/BatchM counters. Sliding lease: writes (UpdateBatchTarget, SetBatchJobStatus, TouchBatchJob, MarkInFlightFailed) bump the job lease to now; reads never write. BatchJobByIDem returns the newest (caller, idem_key) job with lease+ttl>now and status queued/running/done/partial. UpdateBatchTarget counts targets_done exactly once per target entering done/failed and adds len(result) to nbytes. SweepBatch removes by lease; EvictBatchOldest evicts whole non-live jobs oldest-lease-first; MarkInFlightFailed reconciles in-flight jobs on restart. Batch nbytes are kept out of Size() and the content evict union. Co-Authored-By: ornith-1.5-35b-a3b <ornith-1.5-35b-a3b@llama.cpp.invalid>
2026-09-26 20:38:42 -07:00
parent bdc093ee02
commit 2529cb8e84
4 changed files with 597 additions and 0 deletions
@@ -0,0 +1,314 @@
+package dcache
+
+import (
+	"database/sql"
+	"fmt"
+	"time"
+)
+
+// terminal statuses mark a target as finished; a target already in one is not
+// counted again on a re-write.
+func isTerminal(status string) bool {
+	return status == "done" || status == "failed"
+}
+
+type BatchJob struct {
+	ID           string
+	Caller       string
+	CallerClass  string
+	IDemKey      string
+	Status       string
+	Error        string
+	TargetsTotal int
+	TargetsDone  int
+	CreatedAt    int64
+	Lease        int64
+	NBytes       int64
+}
+
+type BatchTarget struct {
+	JobID   string
+	Idx     int
+	URL     string
+	Options string
+	Status  string
+	Result  *string
+	Error   *string
+	Hint    *string
+}
+
+func nullString(s *string) sql.NullString {
+	if s == nil {
+		return sql.NullString{Valid: false}
+	}
+	return sql.NullString{String: *s, Valid: true}
+}
+
+func nullPtr(n sql.NullString) *string {
+	if !n.Valid {
+		return nil
+	}
+	s := n.String
+	return &s
+}
+
+func (d *DB) CreateBatchJob(j BatchJob) error {
+	if d == nil {
+		return fmt.Errorf("dcache: nil db")
+	}
+	_, err := d.sql.Exec(`
+		INSERT INTO batch_job
+			(id, caller, caller_class, idem_key, status, error, targets_total, targets_done, created_at, lease, nbytes)
+		VALUES(?,?,?,?,?,?,?,?,?,?,?)
+	`, j.ID, j.Caller, j.CallerClass, j.IDemKey, j.Status, j.Error, j.TargetsTotal, j.TargetsDone, j.CreatedAt, j.Lease, j.NBytes)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (d *DB) BatchJob(id string) (BatchJob, bool) {
+	if d == nil {
+		return BatchJob{}, false
+	}
+	var j BatchJob
+	err := d.sql.QueryRow(`
+		SELECT id, caller, caller_class, idem_key, status, error,
+		       targets_total, targets_done, created_at, lease, nbytes
+		FROM batch_job WHERE id = ?
+	`, id).Scan(&j.ID, &j.Caller, &j.CallerClass, &j.IDemKey, &j.Status, &j.Error,
+		&j.TargetsTotal, &j.TargetsDone, &j.CreatedAt, &j.Lease, &j.NBytes)
+	if err != nil {
+		d.mu.Lock()
+		d.BatchM++
+		d.mu.Unlock()
+		return BatchJob{}, false
+	}
+	d.mu.Lock()
+	d.BatchH++
+	d.mu.Unlock()
+	return j, true
+}
+
+// BatchJobByIDem returns the id of the newest job for (caller, idemKey) that is
+// still within its lease and not failed/cancelled. Reads never write.
+func (d *DB) BatchJobByIDem(caller, idemKey string, now int64, ttl time.Duration) (string, bool) {
+	if d == nil {
+		return "", false
+	}
+	ttlSec := int64(ttl.Seconds())
+	var id string
+	err := d.sql.QueryRow(`
+		SELECT id FROM batch_job
+		WHERE caller = ? AND idem_key = ? AND lease + ? > ?
+			AND status IN ('queued','running','done','partial')
+		ORDER BY created_at DESC LIMIT 1
+	`, caller, idemKey, ttlSec, now).Scan(&id)
+	if err != nil {
+		return "", false
+	}
+	return id, true
+}
+
+func (d *DB) LiveBatchJobs(caller string) int {
+	if d == nil {
+		return 0
+	}
+	var n int
+	_ = d.sql.QueryRow(`
+		SELECT COUNT(*) FROM batch_job WHERE caller = ? AND status IN ('queued','running')
+	`, caller).Scan(&n)
+	return n
+}
+
+func (d *DB) LiveBatchJobsTotal() int {
+	if d == nil {
+		return 0
+	}
+	var n int
+	_ = d.sql.QueryRow(`SELECT COUNT(*) FROM batch_job WHERE status IN ('queued','running')`).Scan(&n)
+	return n
+}
+
+func (d *DB) PutBatchTargets(jobID string, targets []BatchTarget) error {
+	if d == nil {
+		return fmt.Errorf("dcache: nil db")
+	}
+	for _, t := range targets {
+		if _, err := d.sql.Exec(`
+			INSERT INTO batch_target(job_id, idx, url, options, status, result, error, hint)
+			VALUES(?,?,?,?,?,?,?,?)
+		`, t.JobID, t.Idx, t.URL, t.Options, t.Status,
+			nullString(t.Result), nullString(t.Error), nullString(t.Hint)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (d *DB) BatchTargets(jobID string) ([]BatchTarget, bool) {
+	if d == nil {
+		return nil, false
+	}
+	rows, err := d.sql.Query(`
+		SELECT job_id, idx, url, options, status, result, error, hint
+		FROM batch_target WHERE job_id = ? ORDER BY idx ASC
+	`, jobID)
+	if err != nil {
+		d.mu.Lock()
+		d.BatchM++
+		d.mu.Unlock()
+		return nil, false
+	}
+	defer rows.Close()
+	var out []BatchTarget
+	for rows.Next() {
+		var (
+			r, e, h sql.NullString
+			t       BatchTarget
+		)
+		if err := rows.Scan(&t.JobID, &t.Idx, &t.URL, &t.Options, &t.Status, &r, &e, &h); err != nil {
+			d.mu.Lock()
+			d.BatchM++
+			d.mu.Unlock()
+			return nil, false
+		}
+		t.Result = nullPtr(r)
+		t.Error = nullPtr(e)
+		t.Hint = nullPtr(h)
+		out = append(out, t)
+	}
+	if err := rows.Err(); err != nil {
+		d.mu.Lock()
+		d.BatchM++
+		d.mu.Unlock()
+		return nil, false
+	}
+	if len(out) == 0 {
+		return nil, false
+	}
+	return out, true
+}
+
+func (d *DB) UpdateBatchTarget(jobID string, idx int, now int64, status string, result, errMsg, hint *string) error {
+	if d == nil {
+		return fmt.Errorf("dcache: nil db")
+	}
+	tx, err := d.sql.Begin()
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback()
+
+	var prev string
+	err = tx.QueryRow(`SELECT status FROM batch_target WHERE job_id = ? AND idx = ?`, jobID, idx).Scan(&prev)
+	if err != nil {
+		return err
+	}
+
+	var resBytes int64
+	if result != nil {
+		resBytes = int64(len(*result))
+	}
+	deltaDone := 0
+	if isTerminal(status) && !isTerminal(prev) {
+		deltaDone = 1
+	}
+
+	if _, err := tx.Exec(`
+		UPDATE batch_target SET status = ?, result = ?, error = ?, hint = ? WHERE job_id = ? AND idx = ?
+	`, status, nullString(result), nullString(errMsg), nullString(hint), jobID, idx); err != nil {
+		return err
+	}
+	if _, err := tx.Exec(`
+		UPDATE batch_job SET lease = ?, targets_done = targets_done + ?, nbytes = nbytes + ? WHERE id = ?
+	`, now, deltaDone, resBytes, jobID); err != nil {
+		return err
+	}
+	return tx.Commit()
+}
+
+func (d *DB) SetBatchJobStatus(id, status, errMsg string, now int64) error {
+	if d == nil {
+		return fmt.Errorf("dcache: nil db")
+	}
+	_, err := d.sql.Exec(`UPDATE batch_job SET status = ?, error = ?, lease = ? WHERE id = ?`, status, errMsg, now, id)
+	return err
+}
+
+func (d *DB) TouchBatchJob(id string, now int64) error {
+	if d == nil {
+		return fmt.Errorf("dcache: nil db")
+	}
+	_, err := d.sql.Exec(`UPDATE batch_job SET lease = ? WHERE id = ?`, now, id)
+	return err
+}
+
+func (d *DB) SweepBatch(now int64, ttl time.Duration) int {
+	if d == nil {
+		return 0
+	}
+	ttlSec := int64(ttl.Seconds())
+	if _, err := d.sql.Exec(`
+		DELETE FROM batch_target WHERE job_id IN
+			(SELECT id FROM batch_job WHERE lease + ? <= ?)
+	`, ttlSec, now); err != nil {
+		return 0
+	}
+	res, err := d.sql.Exec(`DELETE FROM batch_job WHERE lease + ? <= ?`, ttlSec, now)
+	if err != nil {
+		return 0
+	}
+	n, _ := res.RowsAffected()
+	return int(n)
+}
+
+func (d *DB) EvictBatchOldest(maxBytes int64) int {
+	if d == nil {
+		return 0
+	}
+	evicted := 0
+	for d.BatchNbytes() > maxBytes {
+		var id string
+		err := d.sql.QueryRow(`
+			SELECT id FROM batch_job
+			WHERE status NOT IN ('queued','running')
+			ORDER BY lease ASC, id ASC
+			LIMIT 1
+		`).Scan(&id)
+		if err != nil {
+			break
+		}
+		if _, err := d.sql.Exec(`DELETE FROM batch_target WHERE job_id = ?`, id); err != nil {
+			break
+		}
+		if _, err := d.sql.Exec(`DELETE FROM batch_job WHERE id = ?`, id); err != nil {
+			break
+		}
+		evicted++
+	}
+	return evicted
+}
+
+func (d *DB) MarkInFlightFailed(reason string, now int64) int {
+	if d == nil {
+		return 0
+	}
+	res, err := d.sql.Exec(`
+		UPDATE batch_job SET status = 'failed', error = ?, lease = ? WHERE status IN ('queued','running')
+	`, reason, now)
+	if err != nil {
+		return 0
+	}
+	n, _ := res.RowsAffected()
+	return int(n)
+}
+
+func (d *DB) BatchNbytes() int64 {
+	if d == nil {
+		return 0
+	}
+	var n sql.NullInt64
+	_ = d.sql.QueryRow(`SELECT IFNULL(SUM(nbytes),0) FROM batch_job`).Scan(&n)
+	return n.Int64
+}