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 }