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") } for _, t := range targets { if _, err := d.sql.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 nil } 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()) if _, err := d.sql.Exec(` DELETE FROM batch_target WHERE job_id IN (SELECT id FROM batch_job WHERE lease + ? <= ?) `, ttlSec, now); err != nil { return 0 } res, err := d.sql.Exec(`DELETE FROM batch_job WHERE lease + ? <= ?`, ttlSec, now) if 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.sql.Exec(`DELETE FROM batch_target WHERE job_id = ?`, id); err != nil { break } if _, err := d.sql.Exec(`DELETE FROM batch_job WHERE id = ?`, 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 }