flamenco/internal/manager/persistence/worker_sleep_schedule.go

package persistence

// SPDX-License-Identifier: GPL-3.0-or-later

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"

	"github.com/rs/zerolog/log"
	"projects.blender.org/studio/flamenco/internal/manager/persistence/sqlc"
)

// SleepSchedule belongs to a Worker, and determines when it's automatically
// sent to the 'asleep' and 'awake' states.
type SleepSchedule struct {
	Model

	WorkerID uint    `gorm:"default:0;unique;index"`
	Worker   *Worker `gorm:"foreignkey:WorkerID;references:ID;constraint:OnDelete:CASCADE"`

	IsActive bool `gorm:"default:false;index"`

	// Space-separated two-letter strings indicating days of week the schedule is
	// active ("mo", "tu", etc.). Empty means "every day".
	DaysOfWeek string    `gorm:"default:''"`
	StartTime  TimeOfDay `gorm:"default:''"`
	EndTime    TimeOfDay `gorm:"default:''"`

	NextCheck time.Time
}

// FetchWorkerSleepSchedule fetches the worker's sleep schedule.
// It does not fetch the worker itself. If you need that, call
// `FetchSleepScheduleWorker()` afterwards.
func (db *DB) FetchWorkerSleepSchedule(ctx context.Context, workerUUID string) (*SleepSchedule, error) {
	logger := log.With().Str("worker", workerUUID).Logger()
	logger.Trace().Msg("fetching worker sleep schedule")

	queries := db.queries()

	sqlcSched, err := queries.FetchWorkerSleepSchedule(ctx, workerUUID)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return nil, nil
	case err != nil:
		return nil, err
	}

	return convertSqlcSleepSchedule(sqlcSched)
}

func (db *DB) SetWorkerSleepSchedule(ctx context.Context, workerUUID string, schedule *SleepSchedule) error {
	logger := log.With().Str("worker", workerUUID).Logger()
	logger.Trace().Msg("setting worker sleep schedule")

	worker, err := db.FetchWorker(ctx, workerUUID)
	if err != nil {
		return fmt.Errorf("fetching worker %q: %w", workerUUID, err)
	}
	schedule.WorkerID = worker.ID
	schedule.Worker = worker

	// Only store timestamps in UTC.
	if schedule.NextCheck.Location() != time.UTC {
		schedule.NextCheck = schedule.NextCheck.UTC()
	}

	queries := db.queries()
	params := sqlc.SetWorkerSleepScheduleParams{
		CreatedAt:  db.gormDB.NowFunc(),
		UpdatedAt:  db.now(),
		WorkerID:   int64(schedule.WorkerID),
		IsActive:   schedule.IsActive,
		DaysOfWeek: schedule.DaysOfWeek,
		StartTime:  schedule.StartTime.String(),
		EndTime:    schedule.EndTime.String(),
		NextCheck:  sql.NullTime{Time: schedule.NextCheck, Valid: !schedule.NextCheck.IsZero()},
	}

	id, err := queries.SetWorkerSleepSchedule(ctx, params)
	if err != nil {
		return fmt.Errorf("storing worker %q sleep schedule: %w", workerUUID, err)
	}
	schedule.ID = uint(id)
	return nil
}

func (db *DB) SetWorkerSleepScheduleNextCheck(ctx context.Context, schedule *SleepSchedule) error {
	// Only store timestamps in UTC.
	if schedule.NextCheck.Location() != time.UTC {
		schedule.NextCheck = schedule.NextCheck.UTC()
	}

	queries := db.queries()
	numAffected, err := queries.SetWorkerSleepScheduleNextCheck(
		ctx,
		sqlc.SetWorkerSleepScheduleNextCheckParams{
			ScheduleID: int64(schedule.ID),
			NextCheck:  sql.NullTime{Time: schedule.NextCheck, Valid: !schedule.NextCheck.IsZero()},
		})
	if err != nil {
		return fmt.Errorf("updating worker sleep schedule: %w", err)
	}
	if numAffected < 1 {
		return fmt.Errorf("could not find worker sleep schedule ID %d", schedule.ID)
	}
	return nil
}

// FetchSleepScheduleWorker sets the given schedule's `Worker` pointer.
func (db *DB) FetchSleepScheduleWorker(ctx context.Context, schedule *SleepSchedule) error {
	queries := db.queries()

	worker, err := queries.FetchWorkerByID(ctx, int64(schedule.WorkerID))
	if err != nil {
		schedule.Worker = nil
		return workerError(err, "finding worker by their sleep schedule")
	}

	schedule.Worker = convertSqlcWorker(worker)
	return nil
}

// FetchSleepSchedulesToCheck returns the sleep schedules that are due for a check.
func (db *DB) FetchSleepSchedulesToCheck(ctx context.Context) ([]*SleepSchedule, error) {
	now := db.now()

	log.Debug().
		Str("timeout", now.Time.String()).
		Msg("fetching sleep schedules that need checking")

	queries := db.queries()
	schedules, err := queries.FetchSleepSchedulesToCheck(ctx, now)
	if err != nil {
		return nil, err
	}

	gormSchedules := make([]*SleepSchedule, len(schedules))
	for index := range schedules {
		gormSched, err := convertSqlcSleepSchedule(schedules[index])
		if err != nil {
			return nil, err
		}
		gormSchedules[index] = gormSched
	}

	return gormSchedules, nil
}

func convertSqlcSleepSchedule(sqlcSchedule sqlc.SleepSchedule) (*SleepSchedule, error) {
	schedule := SleepSchedule{
		Model: Model{
			ID:        uint(sqlcSchedule.ID),
			CreatedAt: sqlcSchedule.CreatedAt,
			UpdatedAt: sqlcSchedule.UpdatedAt.Time,
		},
		WorkerID:   uint(sqlcSchedule.WorkerID),
		IsActive:   sqlcSchedule.IsActive,
		DaysOfWeek: sqlcSchedule.DaysOfWeek,
	}

	err := schedule.StartTime.Scan(sqlcSchedule.StartTime)
	if err != nil {
		return nil, fmt.Errorf("parsing schedule start time %q: %w", sqlcSchedule.StartTime, err)
	}

	err = schedule.EndTime.Scan(sqlcSchedule.EndTime)
	if err != nil {
		return nil, fmt.Errorf("parsing schedule end time %q: %w", sqlcSchedule.EndTime, err)
	}

	return &schedule, nil
}