filed

Job queue using FUSE

git clone git://mccd.space/filed

jobs.go (5542B)

 package store
 
 import (
 	"database/sql"
 	"log/slog"
 	"os"
 	"sync"
 	"time"
 
 	_ "modernc.org/sqlite"
 )
 
 const (
 	StatePending   = "pending"
 	StateRunning   = "active"
 	StateCompleted = "complete"
 	StateFailed    = "failed"
 )
 
 // A single unit of work
 type Job struct {
 	ID      string
 	State   string
 	INode   int
 	Command string
 	// Command output on completion
 	Output    []byte
 	Attempts  int
 	CreatedAt time.Time
 	UpdatedAt time.Time
 }
 
 type Store struct {
 	db     *sql.DB
 	mu     sync.RWMutex
 	config Config
 }
 
 func NewStore(filepath string) (*Store, error) {
 	slog.Info("STORE: Connecting", "filepath", filepath)
 	db, err := sql.Open("sqlite", filepath)
 	if err != nil {
 		return nil, err
 	}
 
 	// Otherwise there are strange race conditions
 	db.SetMaxOpenConns(1)
 
 	if _, err := db.Exec("PRAGMA journal_mode=WAL;PRAGMA busy_timeout=20000;"); err != nil {
 		return nil, err
 	}
 
 	s := &Store{db: db}
 	if err := s.initSchema(); err != nil {
 		db.Close()
 		return nil, err
 	}
 	s.PopulateConfig()
 
 	return s, nil
 }
 
 func (s *Store) initSchema() error {
 	query := `
 	CREATE TABLE IF NOT EXISTS jobs (
 		id TEXT PRIMARY KEY,
 		inode INTEGER,
 		state TEXT NOT NULL,
 		command TEXT NOT NULL,
 		output BLOB,
 		attempts INTEGER DEFAULT 0,
 		created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
 		updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
 	);
 	CREATE INDEX IF NOT EXISTS idx_state ON jobs(state);
 	CREATE TABLE IF NOT EXISTS config (
 		id INTEGER PRIMARY KEY CHECK (id = 0),
 		max_attempts INTEGER DEFAULT 3,
 		timeout_sec INTEGER DEFAULT 180,
 		max_job_count INTEGER DEFAULT 20,
 		backoff_mult INTEGER DEFAULT 2,
 		backoff_base INTEGER DEFAULT 2
 	);
 	INSERT INTO config (id) VALUES (0) ON CONFLICT DO NOTHING;
 	CREATE TRIGGER IF NOT EXISTS auto_increment_inode_trigger
 		AFTER INSERT ON jobs
 		WHEN new.inode IS NULL
 		BEGIN
 			UPDATE jobs
 			SET inode = (SELECT IFNULL(MAX(inode), 0) + 1 FROM jobs)
 			WHERE id = new.id;
 	END;
 	CREATE TABLE IF NOT EXISTS file_meta (
 		inode INTEGER PRIMARY KEY,
 		mode INTEGER,
 		uid INTEGER,
 		gid INTEGER
 	);
 	
 	`
 	_, err := s.db.Exec(query)
 	return err
 }
 
 func (s *Store) Close() error {
 	return s.db.Close()
 }
 
 // Adds a new job with the 'pending' state.
 func (s *Store) CreateJob(id string, command string) (*Job, error) {
 	query := `
 		INSERT INTO jobs (id, state, command, created_at, updated_at)
 		VALUES (?, ?, ?, ?, ?)
 	`
 	now := time.Now()
 	j := &Job{
 		ID:        id,
 		State:     StatePending,
 		Command:   command,
 		CreatedAt: now,
 		UpdatedAt: now,
 		Attempts:  0,
 	}
 
 	_, err := s.db.Exec(query, id, StatePending, command, now, now)
 	if err != nil {
 		return nil, err
 	}
 	return j, nil
 }
 
 func (s *Store) GetJob(id string) (*Job, error) {
 	query := `SELECT id, state, inode, command, output, attempts, created_at FROM jobs WHERE id = ?`
 	j := &Job{}
 	err := s.db.QueryRow(query, id).Scan(
 		&j.ID, &j.State, &j.INode, &j.Command, &j.Output, &j.Attempts, &j.CreatedAt,
 	)
 	if err == sql.ErrNoRows {
 		return nil, os.ErrNotExist
 	}
 	return j, err
 }
 
 func (s *Store) DeleteJob(id string) error {
 	res, err := s.db.Exec("DELETE FROM jobs WHERE id = ?", id)
 	if err != nil {
 		return err
 	}
 	rows, _ := res.RowsAffected()
 	if rows == 0 {
 		return os.ErrNotExist
 	}
 	return nil
 }
 
 func (s *Store) ListJobsByState(state string) ([]Job, error) {
 	// Since it's a queue, it should be first in first out.
 	rows, err := s.db.Query("SELECT id, command, attempts, created_at,updated_at FROM jobs WHERE state = ? ORDER BY created_at", state)
 	if err != nil {
 		return nil, err
 	}
 	defer rows.Close()
 
 	var jobs []Job
 	for rows.Next() {
 		var j Job
 		if err := rows.Scan(&j.ID, &j.Command, &j.Attempts, &j.CreatedAt, &j.UpdatedAt); err != nil {
 			return nil, err
 		}
 		j.State = state
 		jobs = append(jobs, j)
 	}
 	return jobs, nil
 }
 
 // Picks up a pending job and moves it to 'active'.
 func (s *Store) AttemptJob(id string) error {
 	tx, err := s.db.Begin()
 	if err != nil {
 		return err
 	}
 	defer tx.Rollback()
 
 	var currentState string
 	err = tx.QueryRow("SELECT state FROM jobs WHERE id = ?", id).Scan(&currentState)
 	if err != nil {
 		return err
 	}
 
 	query := `
 		UPDATE jobs 
 		SET state = ?, attempts = attempts + 1, updated_at = ? 
 		WHERE id = ?
 	`
 	_, err = tx.Exec(query, StateRunning, time.Now(), id)
 	if err != nil {
 		return err
 	}
 
 	return tx.Commit()
 }
 
 // Moves a job back to 'pending'.
 func (s *Store) RestartJob(id string, output []byte) error {
 	query := `
 		UPDATE jobs 
 		SET state = ?, output = ?, updated_at = ?
 		WHERE id = ?
 	`
 	res, err := s.db.Exec(query, StatePending, output, time.Now(), id)
 	if err != nil {
 		return err
 	}
 	rows, _ := res.RowsAffected()
 	if rows == 0 {
 		return os.ErrNotExist
 	}
 	return nil
 }
 
 func (s *Store) ResetJob(id string) error {
 	query := `
 		UPDATE jobs 
 		SET state = ?, attempts = 0, updated_at = ?
 		WHERE id = ?
 	`
 	res, err := s.db.Exec(query, StatePending, time.Now(), id)
 	if err != nil {
 		return err
 	}
 	rows, _ := res.RowsAffected()
 	if rows == 0 {
 		return os.ErrNotExist
 	}
 	return nil
 }
 
 func (s *Store) CompleteJob(id string, output []byte) error {
 	_, err := s.db.Exec("UPDATE jobs SET state = ?, output = ?, updated_at = ? WHERE id = ?", StateCompleted, output, time.Now(), id)
 	return err
 }
 
 func (s *Store) WriteOutput(id string, output []byte) error {
 	_, err := s.db.Exec("UPDATE jobs SET output = ?, updated_at = ? WHERE id = ?", output, time.Now(), id)
 	return err
 }
 
 func (s *Store) FailJob(id string) error {
 	_, err := s.db.Exec("UPDATE jobs SET state = ?, updated_at = ? WHERE id = ?", StateFailed, time.Now(), id)
 	return err
 }