package jobs import ( "context" "encoding/json" "log" "sync" "github.com/jackc/pgx/v5" "github.com/riverqueue/river" "github.com/riverqueue/river/rivertype" ) type Output struct { Processed int `json:"processed"` Succeeded int `json:"succeeded"` Warnings []string `json:"warnings,omitempty"` } func OutputOf(job *rivertype.JobRow) Output { var out Output if raw := job.Output(); raw != nil { _ = json.Unmarshal(raw, &out) } return out } // Tracker counts a job's work items and persists each move as the job's output. type Tracker struct { jobId int64 mu sync.Mutex out Output } func NewTracker(jobId int64) *Tracker { return &Tracker{jobId: jobId} } func (t *Tracker) Succeed(ctx context.Context) { t.mu.Lock() t.out.Processed++ t.out.Succeeded++ current := t.out t.mu.Unlock() t.persist(ctx, current) } func (t *Tracker) Skip(ctx context.Context, reason string) { t.mu.Lock() t.out.Processed++ t.out.Warnings = append(t.out.Warnings, reason) current := t.out t.mu.Unlock() t.persist(ctx, current) } func (t *Tracker) Output() Output { t.mu.Lock() defer t.mu.Unlock() return t.out } func (t *Tracker) persist(ctx context.Context, current Output) { // No client in the context means the body runs outside River, as in unit tests. client, err := river.ClientFromContextSafely[pgx.Tx](ctx) if err != nil { return } if err := river.RecordOutput(ctx, current); err != nil { log.Printf("[jobs] job %d could not record progress: %v", t.jobId, err) return } if _, err := client.JobUpdate(ctx, t.jobId, &river.JobUpdateParams{}); err != nil { log.Printf("[jobs] job %d could not persist progress: %v", t.jobId, err) } }