sift/internal/dcache/batch.go

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")
	}
	// One transaction: a job never has only some of its targets.
	tx, err := d.sql.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	for _, t := range targets {
		if _, err := tx.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 tx.Commit()
}

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())
	// A queued or running job is never swept: it is still in a worker's rotation. The restart
	// reaper marks such jobs failed, and they are swept after that. Both deletes are one
	// transaction, so a job never loses its targets but keeps its row.
	tx, err := d.sql.Begin()
	if err != nil {
		return 0
	}
	defer func() { _ = tx.Rollback() }()
	const expired = `SELECT id FROM batch_job WHERE lease + ? <= ? AND status NOT IN ('queued','running')`
	if _, err := tx.Exec(`DELETE FROM batch_target WHERE job_id IN (`+expired+`)`, ttlSec, now); err != nil {
		return 0
	}
	res, err := tx.Exec(`DELETE FROM batch_job WHERE id IN (`+expired+`)`, ttlSec, now)
	if err != nil {
		return 0
	}
	if err := tx.Commit(); 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.deleteBatchJob(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
}

// deleteBatchJob removes a job and its targets in one transaction.
func (d *DB) deleteBatchJob(id string) error {
	tx, err := d.sql.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	if _, err := tx.Exec(`DELETE FROM batch_target WHERE job_id = ?`, id); err != nil {
		return err
	}
	if _, err := tx.Exec(`DELETE FROM batch_job WHERE id = ?`, id); err != nil {
		return err
	}
	return tx.Commit()
}