refactor: implement worker pool pattern, add persistent task queue, and decouple target loading logic

2026-08-12 23:28:31 -07:00
parent 51b9468737
commit 890469fed8
18 changed files with 694 additions and 442 deletions
@@ -0,0 +1,159 @@
+package db
+
+import (
+	"database/sql"
+	"encoding/json"
+	"fmt"
+	"log"
+)
+
+// EnqueueTask adds a new task to the queue
+func (db *DB) EnqueueTask(taskType string, payload interface{}) error {
+	payloadBytes, err := json.Marshal(payload)
+	if err != nil {
+		return fmt.Errorf("failed to marshal task payload: %w", err)
+	}
+
+	query := `
+	INSERT INTO tasks (task_type, payload, status)
+	VALUES (?, ?, 'PENDING')
+	`
+	_, err = db.Exec(query, taskType, string(payloadBytes))
+	return err
+}
+
+// ClaimNextTask atomically finds and claims the next pending task
+func (db *DB) ClaimNextTask() (*Task, error) {
+	tx, err := db.Begin()
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = tx.Rollback() }()
+
+	// Find the oldest PENDING task that is ready to be retried
+	query := `
+	SELECT id, task_type, payload, status, error_message, retry_count, next_retry_at, created_at, updated_at
+	FROM tasks
+	WHERE status = 'PENDING' AND next_retry_at <= CURRENT_TIMESTAMP
+	ORDER BY created_at ASC
+	LIMIT 1
+	`
+	
+	var task Task
+	var errMsg sql.NullString
+	err = tx.QueryRow(query).Scan(
+		&task.ID, &task.TaskType, &task.Payload, &task.Status, &errMsg,
+		&task.RetryCount, &task.NextRetryAt, &task.CreatedAt, &task.UpdatedAt,
+	)
+	
+	if err == sql.ErrNoRows {
+		return nil, nil // No tasks available
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	if errMsg.Valid {
+		task.ErrorMessage = errMsg.String
+	}
+
+	// Claim the task
+	updateQuery := `UPDATE tasks SET status = 'RUNNING', updated_at = CURRENT_TIMESTAMP WHERE id = ?`
+	_, err = tx.Exec(updateQuery, task.ID)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := tx.Commit(); err != nil {
+		return nil, err
+	}
+
+	return &task, nil
+}
+
+// CompleteTask marks a task as COMPLETED
+func (db *DB) CompleteTask(id int64) error {
+	query := `UPDATE tasks SET status = 'COMPLETED', updated_at = CURRENT_TIMESTAMP WHERE id = ?`
+	_, err := db.Exec(query, id)
+	return err
+}
+
+// FailTask marks a task as FAILED or PENDING (for retry) using exponential backoff
+func (db *DB) FailTask(id int64, errStr string) error {
+	// First get current retry count
+	var retryCount int
+	err := db.QueryRow(`SELECT retry_count FROM tasks WHERE id = ?`, id).Scan(&retryCount)
+	if err != nil {
+		log.Printf("[DB] FailTask: failed to get retry_count for task %d: %v", id, err)
+		return err
+	}
+
+	retryCount++
+	
+	// Max retries = 5 (1s, 2s, 4s, 8s, 16s)
+	status := "PENDING"
+	if retryCount > 5 {
+		status = "FAILED"
+	}
+
+	// Exponential backoff
+	backoffSeconds := 1 << retryCount // 2^retryCount
+	
+	query := `
+	UPDATE tasks 
+	SET status = ?, 
+	    error_message = ?, 
+	    retry_count = ?,
+	    next_retry_at = datetime('now', '+' || ? || ' seconds'),
+	    updated_at = CURRENT_TIMESTAMP
+	WHERE id = ?
+	`
+	_, err = db.Exec(query, status, errStr, retryCount, fmt.Sprintf("%d", backoffSeconds), id)
+	return err
+}
+
+// HasPendingFetchTask checks if there is already a PENDING or RUNNING fetch task for a given canonical ID
+func (db *DB) HasPendingFetchTask(canonicalID string) (bool, error) {
+	query := `
+	SELECT COUNT(*) FROM tasks 
+	WHERE task_type = 'FETCH' 
+	  AND status IN ('PENDING', 'RUNNING') 
+	  AND payload LIKE ?
+	`
+	likePattern := fmt.Sprintf(`%%"canonical_id":"%s"%%`, canonicalID)
+	
+	var count int
+	err := db.QueryRow(query, likePattern).Scan(&count)
+	if err != nil {
+		return false, err
+	}
+	return count > 0, nil
+}
+
+// ClearTaskQueue deletes all tasks or only pending/failed ones
+func (db *DB) ClearTaskQueue() error {
+	query := `DELETE FROM tasks WHERE status IN ('PENDING', 'FAILED', 'COMPLETED')`
+	_, err := db.Exec(query)
+	return err
+}
+
+// GetQueueStatus returns counts of tasks by status
+func (db *DB) GetQueueStatus() (map[string]int, error) {
+	query := `SELECT status, COUNT(*) FROM tasks GROUP BY status`
+	rows, err := db.Query(query)
+	if err != nil {
+		return nil, err
+	}
+	defer func() { _ = rows.Close() }()
+
+	counts := make(map[string]int)
+	for rows.Next() {
+		var status string
+		var count int
+		if err := rows.Scan(&status, &count); err != nil {
+			return nil, err
+		}
+		counts[status] = count
+	}
+	return counts, nil
+}