blender/flamenco
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Transition from ex-GORM structs to sqlc structs (4/5)

Browse Source 
Replace old used-to-be-GORM datastructures (#104305) with sqlc-generated
structs. This also makes it possible to use more specific structs that
are more taylored to the specific queries, increasing efficiency.

This commit deals with the worker sleep schedule.

Functional changes are kept to a minimum, as the API still serves the
same data.

Because this work covers so much of Flamenco's code, it's been split up
into different commits. Each commit brings Flamenco to a state where it
compiles and unit tests pass. Only the result of the final commit has
actually been tested properly.

Ref: #104343
This commit is contained in:
Sybren A. Stüvel 2024-11-11 23:18:29 +01:00
parent c04e4992e0
commit ddced5a823
19 changed files with 339 additions and 336 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internal/manager/api_impl/interfaces.go
View File

@ -230,7 +230,7 @@ var _ TimeService = (clock.Clock)(nil)
type WorkerSleepScheduler interface {
FetchSchedule(ctx context.Context, workerUUID string) (*persistence.SleepSchedule, error)
SetSchedule(ctx context.Context, workerUUID string, schedule *persistence.SleepSchedule) error
SetSchedule(ctx context.Context, workerUUID string, schedule persistence.SleepSchedule) error
WorkerStatus(ctx context.Context, workerUUID string) (api.WorkerStatus, error)
}

6
internal/manager/api_impl/mocks/api_impl_mock.gen.go generated
View File

@ -1323,10 +1323,10 @@ func (m *MockWorkerSleepScheduler) EXPECT() *MockWorkerSleepSchedulerMockRecorde
}
// FetchSchedule mocks base method.
func (m *MockWorkerSleepScheduler) FetchSchedule(arg0 context.Context, arg1 string) (*persistence.SleepSchedule, error) {
func (m *MockWorkerSleepScheduler) FetchSchedule(arg0 context.Context, arg1 string) (*sqlc.SleepSchedule, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "FetchSchedule", arg0, arg1)
ret0, _ := ret[0].(*persistence.SleepSchedule)
ret0, _ := ret[0].(*sqlc.SleepSchedule)
ret1, _ := ret[1].(error)
return ret0, ret1
}
@ -1338,7 +1338,7 @@ func (mr *MockWorkerSleepSchedulerMockRecorder) FetchSchedule(arg0, arg1 interfa
}
// SetSchedule mocks base method.
func (m *MockWorkerSleepScheduler) SetSchedule(arg0 context.Context, arg1 string, arg2 *persistence.SleepSchedule) error {
func (m *MockWorkerSleepScheduler) SetSchedule(arg0 context.Context, arg1 string, arg2 sqlc.SleepSchedule) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SetSchedule", arg0, arg1, arg2)
ret0, _ := ret[0].(error)

2
internal/manager/api_impl/worker_sleep_schedule.go
View File

@ -74,7 +74,7 @@ func (f *Flamenco) SetWorkerSleepSchedule(e echo.Context, workerUUID string) err
}
// Send the sleep schedule to the scheduler.
err = f.sleepScheduler.SetSchedule(ctx, workerUUID, &dbSchedule)
err = f.sleepScheduler.SetSchedule(ctx, workerUUID, dbSchedule)
switch {
case errors.Is(err, persistence.ErrWorkerNotFound):
logger.Warn().Msg("SetWorkerSleepSchedule: worker does not exist")

9
internal/manager/persistence/sqlc/methods.go
View File

@ -3,14 +3,23 @@
package sqlc
import (
"database/sql"
"fmt"
"strings"
"time"
"projects.blender.org/studio/flamenco/pkg/api"
)
// SPDX-License-Identifier: GPL-3.0-or-later
func (ss *SleepSchedule) SetNextCheck(nextCheck time.Time) {
ss.NextCheck = sql.NullTime{
Time: nextCheck,
Valid: true,
}
}
func (w *Worker) Identifier() string {
// Avoid a panic when worker.Identifier() is called on a nil pointer.
if w == nil {

5
internal/manager/persistence/sqlc/models.go
View File

@ -10,6 +10,7 @@ import (
"time"
"projects.blender.org/studio/flamenco/pkg/api"
"projects.blender.org/studio/flamenco/pkg/time_of_day"
)
type Job struct {
@ -51,8 +52,8 @@ type SleepSchedule struct {
WorkerID int64
IsActive bool
DaysOfWeek string
StartTime string
EndTime string
StartTime time_of_day.TimeOfDay
EndTime time_of_day.TimeOfDay
NextCheck sql.NullTime
}

4
internal/manager/persistence/sqlc/query_workers.sql
View File

@ -205,7 +205,9 @@ WHERE ID=@schedule_id;
-- name: FetchSleepSchedulesToCheck :many
SELECT * FROM sleep_schedules
SELECT sqlc.embed(sleep_schedules), workers.uuid as workeruuid, workers.name as worker_name
FROM sleep_schedules
LEFT JOIN workers ON workers.id = sleep_schedules.worker_id
WHERE is_active
AND (next_check <= @next_check OR next_check IS NULL OR next_check = '');

45
internal/manager/persistence/sqlc/query_workers.sql.go
View File

@ -12,6 +12,7 @@ import (
"time"
"projects.blender.org/studio/flamenco/pkg/api"
"projects.blender.org/studio/flamenco/pkg/time_of_day"
)
const countWorkerTags = `-- name: CountWorkerTags :one
@ -147,30 +148,40 @@ func (q *Queries) DeleteWorkerTag(ctx context.Context, uuid string) (int64, erro
}
const fetchSleepSchedulesToCheck = `-- name: FetchSleepSchedulesToCheck :many
SELECT id, created_at, updated_at, worker_id, is_active, days_of_week, start_time, end_time, next_check FROM sleep_schedules
SELECT sleep_schedules.id, sleep_schedules.created_at, sleep_schedules.updated_at, sleep_schedules.worker_id, sleep_schedules.is_active, sleep_schedules.days_of_week, sleep_schedules.start_time, sleep_schedules.end_time, sleep_schedules.next_check, workers.uuid as workeruuid, workers.name as worker_name
FROM sleep_schedules
LEFT JOIN workers ON workers.id = sleep_schedules.worker_id
WHERE is_active
AND (next_check <= ?1 OR next_check IS NULL OR next_check = '')
`
func (q *Queries) FetchSleepSchedulesToCheck(ctx context.Context, nextCheck sql.NullTime) ([]SleepSchedule, error) {
type FetchSleepSchedulesToCheckRow struct {
SleepSchedule SleepSchedule
WorkerUUID sql.NullString
WorkerName sql.NullString
}
func (q *Queries) FetchSleepSchedulesToCheck(ctx context.Context, nextCheck sql.NullTime) ([]FetchSleepSchedulesToCheckRow, error) {
rows, err := q.db.QueryContext(ctx, fetchSleepSchedulesToCheck, nextCheck)
if err != nil {
return nil, err
}
defer rows.Close()
var items []SleepSchedule
var items []FetchSleepSchedulesToCheckRow
for rows.Next() {
var i SleepSchedule
var i FetchSleepSchedulesToCheckRow
if err := rows.Scan(
&i.ID,
&i.CreatedAt,
&i.UpdatedAt,
&i.WorkerID,
&i.IsActive,
&i.DaysOfWeek,
&i.StartTime,
&i.EndTime,
&i.NextCheck,
&i.SleepSchedule.ID,
&i.SleepSchedule.CreatedAt,
&i.SleepSchedule.UpdatedAt,
&i.SleepSchedule.WorkerID,
&i.SleepSchedule.IsActive,
&i.SleepSchedule.DaysOfWeek,
&i.SleepSchedule.StartTime,
&i.SleepSchedule.EndTime,
&i.SleepSchedule.NextCheck,
&i.WorkerUUID,
&i.WorkerName,
); err != nil {
return nil, err
}
@ -737,8 +748,8 @@ type SetWorkerSleepScheduleParams struct {
WorkerID int64
IsActive bool
DaysOfWeek string
StartTime string
EndTime string
StartTime time_of_day.TimeOfDay
EndTime time_of_day.TimeOfDay
NextCheck sql.NullTime
}
@ -858,8 +869,8 @@ type Test_CreateWorkerSleepScheduleParams struct {
WorkerID int64
IsActive bool
DaysOfWeek string
StartTime string
EndTime string
StartTime time_of_day.TimeOfDay
EndTime time_of_day.TimeOfDay
NextCheck sql.NullTime
}

118
internal/manager/persistence/worker_sleep_schedule.go
View File

@ -7,7 +7,6 @@ import (
"database/sql"
"errors"
"fmt"
"time"
"github.com/rs/zerolog/log"
"projects.blender.org/studio/flamenco/internal/manager/persistence/sqlc"
@ -15,41 +14,29 @@ import (
// SleepSchedule belongs to a Worker, and determines when it's automatically
// sent to the 'asleep' and 'awake' states.
type SleepSchedule struct {
Model
type SleepSchedule = sqlc.SleepSchedule
WorkerID uint
Worker *Worker
IsActive bool
// Space-separated two-letter strings indicating days of week the schedule is
// active ("mo", "tu", etc.). Empty means "every day".
DaysOfWeek string
StartTime TimeOfDay
EndTime TimeOfDay
NextCheck time.Time
type SleepScheduleOwned struct {
SleepSchedule SleepSchedule
WorkerName string
WorkerUUID string
}
// 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)
schedule, err := queries.FetchWorkerSleepSchedule(ctx, workerUUID)
switch {
case errors.Is(err, sql.ErrNoRows):
return nil, nil
case err != nil:
return nil, err
}
return convertSqlcSleepSchedule(sqlcSched)
return &schedule, nil
}
func (db *DB) SetWorkerSleepSchedule(ctx context.Context, workerUUID string, schedule *SleepSchedule) error {
@ -60,46 +47,38 @@ func (db *DB) SetWorkerSleepSchedule(ctx context.Context, workerUUID string, sch
if err != nil {
return fmt.Errorf("fetching worker %q: %w", workerUUID, err)
}
schedule.WorkerID = uint(worker.ID)
schedule.Worker = worker
// Only store timestamps in UTC.
if schedule.NextCheck.Location() != time.UTC {
schedule.NextCheck = schedule.NextCheck.UTC()
}
schedule.WorkerID = worker.ID
queries := db.queries()
params := sqlc.SetWorkerSleepScheduleParams{
CreatedAt: db.now(),
UpdatedAt: db.nowNullable(),
WorkerID: int64(schedule.WorkerID),
WorkerID: 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()},
StartTime: schedule.StartTime,
EndTime: schedule.EndTime,
NextCheck: nullTimeToUTC(schedule.NextCheck),
}
id, err := queries.SetWorkerSleepSchedule(ctx, params)
if err != nil {
return fmt.Errorf("storing worker %q sleep schedule: %w", workerUUID, err)
}
schedule.ID = uint(id)
schedule.ID = id
schedule.NextCheck = params.NextCheck
schedule.CreatedAt = params.CreatedAt
schedule.UpdatedAt = params.UpdatedAt
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()
}
func (db *DB) SetWorkerSleepScheduleNextCheck(ctx context.Context, schedule SleepSchedule) error {
queries := db.queries()
numAffected, err := queries.SetWorkerSleepScheduleNextCheck(
ctx,
sqlc.SetWorkerSleepScheduleNextCheckParams{
ScheduleID: int64(schedule.ID),
NextCheck: sql.NullTime{Time: schedule.NextCheck, Valid: !schedule.NextCheck.IsZero()},
NextCheck: nullTimeToUTC(schedule.NextCheck),
})
if err != nil {
return fmt.Errorf("updating worker sleep schedule: %w", err)
@ -110,22 +89,19 @@ func (db *DB) SetWorkerSleepScheduleNextCheck(ctx context.Context, schedule *Sle
return nil
}
// FetchSleepScheduleWorker sets the given schedule's `Worker` pointer.
func (db *DB) FetchSleepScheduleWorker(ctx context.Context, schedule *SleepSchedule) error {
// FetchSleepScheduleWorker returns the given schedule's associated Worker.
func (db *DB) FetchSleepScheduleWorker(ctx context.Context, schedule SleepSchedule) (*Worker, error) {
queries := db.queries()
worker, err := queries.FetchWorkerByID(ctx, int64(schedule.WorkerID))
worker, err := queries.FetchWorkerByID(ctx, schedule.WorkerID)
if err != nil {
schedule.Worker = nil
return workerError(err, "finding worker by their sleep schedule")
return nil, workerError(err, "finding worker by their sleep schedule")
}
return &worker, nil
}
schedule.Worker = &worker
return nil
}
// FetchSleepSchedulesToCheck returns the sleep schedules that are due for a check.
func (db *DB) FetchSleepSchedulesToCheck(ctx context.Context) ([]*SleepSchedule, error) {
// FetchSleepSchedulesToCheck returns the sleep schedules that are due for a check, with their owning Worker.
func (db *DB) FetchSleepSchedulesToCheck(ctx context.Context) ([]SleepScheduleOwned, error) {
now := db.nowNullable()
log.Debug().
@ -133,44 +109,16 @@ func (db *DB) FetchSleepSchedulesToCheck(ctx context.Context) ([]*SleepSchedule,
Msg("fetching sleep schedules that need checking")
queries := db.queries()
schedules, err := queries.FetchSleepSchedulesToCheck(ctx, now)
rows, 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
schedules := make([]SleepScheduleOwned, len(rows))
for index, row := range rows {
schedules[index].SleepSchedule = row.SleepSchedule
schedules[index].WorkerName = row.WorkerName.String
schedules[index].WorkerUUID = row.WorkerUUID.String
}
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
return schedules, nil
}

175
internal/manager/persistence/worker_sleep_schedule_test.go
View File

@ -3,6 +3,7 @@ package persistence
// SPDX-License-Identifier: GPL-3.0-or-later
import (
"database/sql"
"testing"
"time"
@ -10,6 +11,7 @@ import (
"github.com/stretchr/testify/require"
"projects.blender.org/studio/flamenco/internal/uuid"
"projects.blender.org/studio/flamenco/pkg/api"
"projects.blender.org/studio/flamenco/pkg/time_of_day"
)
func TestFetchWorkerSleepSchedule(t *testing.T) {
@ -31,29 +33,27 @@ func TestFetchWorkerSleepSchedule(t *testing.T) {
// Not an existing Worker.
fetched, err := db.FetchWorkerSleepSchedule(ctx, "2cf6153a-3d4e-49f4-a5c0-1c9fc176e155")
require.NoError(t, err, "non-existent worker should not cause an error")
assert.Nil(t, fetched)
assert.Zero(t, fetched)
// No sleep schedule.
fetched, err = db.FetchWorkerSleepSchedule(ctx, linuxWorker.UUID)
require.NoError(t, err, "non-existent schedule should not cause an error")
assert.Nil(t, fetched)
assert.Zero(t, fetched)
// Create a sleep schedule.
created := SleepSchedule{
WorkerID: uint(linuxWorker.ID),
Worker: &linuxWorker,
WorkerID: linuxWorker.ID,
IsActive: true,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
}
err = db.SetWorkerSleepSchedule(ctx, linuxWorker.UUID, &created)
require.NoError(t, err)
fetched, err = db.FetchWorkerSleepSchedule(ctx, linuxWorker.UUID)
require.NoError(t, err)
assertEqualSleepSchedule(t, uint(linuxWorker.ID), created, *fetched)
assertEqualSleepSchedule(t, linuxWorker.ID, created, *fetched)
}
func TestFetchSleepScheduleWorker(t *testing.T) {
@ -74,33 +74,32 @@ func TestFetchSleepScheduleWorker(t *testing.T) {
// Create a sleep schedule.
created := SleepSchedule{
WorkerID: uint(linuxWorker.ID),
Worker: &linuxWorker,
WorkerID: linuxWorker.ID,
IsActive: true,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
}
err = db.SetWorkerSleepSchedule(ctx, linuxWorker.UUID, &created)
require.NoError(t, err)
dbSchedule, err := db.FetchWorkerSleepSchedule(ctx, linuxWorker.UUID)
require.NoError(t, err)
assert.Nil(t, dbSchedule.Worker, "worker should be nil when fetching schedule")
require.NotNil(t, dbSchedule)
err = db.FetchSleepScheduleWorker(ctx, dbSchedule)
worker, err := db.FetchSleepScheduleWorker(ctx, *dbSchedule)
require.NoError(t, err)
if assert.NotNil(t, dbSchedule.Worker) {
if assert.NotNil(t, worker) {
// Compare a few fields. If these are good, the correct worker has been fetched.
assert.Equal(t, linuxWorker.ID, dbSchedule.Worker.ID)
assert.Equal(t, linuxWorker.UUID, dbSchedule.Worker.UUID)
assert.Equal(t, linuxWorker.ID, worker.ID)
assert.Equal(t, linuxWorker.UUID, worker.UUID)
}
// Deleting the Worker should result in a specific error when fetching the schedule again.
require.NoError(t, db.DeleteWorker(ctx, linuxWorker.UUID))
assert.ErrorIs(t, db.FetchSleepScheduleWorker(ctx, dbSchedule), ErrWorkerNotFound)
assert.Nil(t, dbSchedule.Worker)
worker, err = db.FetchSleepScheduleWorker(ctx, *dbSchedule)
assert.ErrorIs(t, err, ErrWorkerNotFound)
assert.Nil(t, worker)
}
func TestSetWorkerSleepSchedule(t *testing.T) {
@ -120,13 +119,11 @@ func TestSetWorkerSleepSchedule(t *testing.T) {
require.NoError(t, err)
schedule := SleepSchedule{
WorkerID: uint(linuxWorker.ID),
Worker: &linuxWorker,
WorkerID: linuxWorker.ID,
IsActive: true,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
}
// Not an existing Worker.
@ -138,52 +135,47 @@ func TestSetWorkerSleepSchedule(t *testing.T) {
require.NoError(t, err)
fetched, err := db.FetchWorkerSleepSchedule(ctx, linuxWorker.UUID)
require.NoError(t, err)
assertEqualSleepSchedule(t, uint(linuxWorker.ID), schedule, *fetched)
assertEqualSleepSchedule(t, linuxWorker.ID, schedule, *fetched)
// Overwrite the schedule with one that already has a database ID.
newSchedule := schedule
newSchedule.IsActive = false
newSchedule.DaysOfWeek = "mo,tu,we,th,fr"
newSchedule.StartTime = TimeOfDay{2, 0}
newSchedule.EndTime = TimeOfDay{6, 0}
newSchedule.StartTime = time_of_day.New(2, 0)
newSchedule.EndTime = time_of_day.New(6, 0)
err = db.SetWorkerSleepSchedule(ctx, linuxWorker.UUID, &newSchedule)
require.NoError(t, err)
fetched, err = db.FetchWorkerSleepSchedule(ctx, linuxWorker.UUID)
require.NoError(t, err)
assertEqualSleepSchedule(t, uint(linuxWorker.ID), newSchedule, *fetched)
assertEqualSleepSchedule(t, linuxWorker.ID, newSchedule, *fetched)
// Overwrite the schedule with a freshly constructed one.
newerSchedule := SleepSchedule{
WorkerID: uint(linuxWorker.ID),
Worker: &linuxWorker,
WorkerID: linuxWorker.ID,
IsActive: true,
DaysOfWeek: "mo",
StartTime: TimeOfDay{3, 0},
EndTime: TimeOfDay{15, 0},
StartTime: time_of_day.New(3, 0),
EndTime: time_of_day.New(15, 0),
}
err = db.SetWorkerSleepSchedule(ctx, linuxWorker.UUID, &newerSchedule)
require.NoError(t, err)
fetched, err = db.FetchWorkerSleepSchedule(ctx, linuxWorker.UUID)
require.NoError(t, err)
assertEqualSleepSchedule(t, uint(linuxWorker.ID), newerSchedule, *fetched)
assertEqualSleepSchedule(t, linuxWorker.ID, newerSchedule, *fetched)
// Clear the sleep schedule.
emptySchedule := SleepSchedule{
WorkerID: uint(linuxWorker.ID),
Worker: &linuxWorker,
WorkerID: linuxWorker.ID,
IsActive: false,
DaysOfWeek: "",
StartTime: emptyToD,
EndTime: emptyToD,
StartTime: time_of_day.Empty(),
EndTime: time_of_day.Empty(),
}
err = db.SetWorkerSleepSchedule(ctx, linuxWorker.UUID, &emptySchedule)
require.NoError(t, err)
fetched, err = db.FetchWorkerSleepSchedule(ctx, linuxWorker.UUID)
require.NoError(t, err)
assertEqualSleepSchedule(t, uint(linuxWorker.ID), emptySchedule, *fetched)
assertEqualSleepSchedule(t, linuxWorker.ID, emptySchedule, *fetched)
}
func TestSetWorkerSleepScheduleNextCheck(t *testing.T) {
@ -195,24 +187,23 @@ func TestSetWorkerSleepScheduleNextCheck(t *testing.T) {
})
schedule := SleepSchedule{
Worker: &w,
IsActive: true,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
}
err := db.SetWorkerSleepSchedule(ctx, w.UUID, &schedule)
require.NoError(t, err)
future := db.now().Add(5 * time.Hour)
schedule.NextCheck = future
schedule.NextCheck = sql.NullTime{Time: future, Valid: true}
err = db.SetWorkerSleepScheduleNextCheck(ctx, &schedule)
err = db.SetWorkerSleepScheduleNextCheck(ctx, schedule)
require.NoError(t, err)
fetched, err := db.FetchWorkerSleepSchedule(ctx, schedule.Worker.UUID)
fetched, err := db.FetchWorkerSleepSchedule(ctx, w.UUID)
require.NoError(t, err)
assertEqualSleepSchedule(t, uint(schedule.Worker.ID), schedule, *fetched)
assertEqualSleepSchedule(t, w.ID, schedule, *fetched)
}
func TestFetchSleepSchedulesToCheck(t *testing.T) {
@ -225,67 +216,64 @@ func TestFetchSleepSchedulesToCheck(t *testing.T) {
db.nowfunc = func() time.Time { return mockedNow }
worker0 := createWorker(ctx, t, db, func(w *Worker) {
w.UUID = "2b1f857a-fd64-484b-9c17-cf89bbe47be7"
w.Name = "дрон 1"
w.Status = api.WorkerStatusAwake
})
worker1 := createWorker(ctx, t, db, func(w *Worker) {
w.UUID = "4475738e-41eb-47b2-8bca-2bbcabab69bb"
w.Name = "дрон 2"
w.Status = api.WorkerStatusAwake
})
worker2 := createWorker(ctx, t, db, func(w *Worker) {
w.UUID = "dc251817-6a11-4548-a36a-07b0d50b4c21"
w.Name = "дрон 3"
w.Status = api.WorkerStatusAwake
})
worker3 := createWorker(ctx, t, db, func(w *Worker) {
w.UUID = "874d5fc6-5784-4d43-8c20-6e7e73fc1b8d"
w.Name = "дрон 4"
w.Status = api.WorkerStatusAwake
})
schedule0 := SleepSchedule{ // Next check in the past -> should be checked.
Worker: &Worker{
UUID: "2b1f857a-fd64-484b-9c17-cf89bbe47be7",
Name: "дрон 1",
Status: api.WorkerStatusAwake,
},
IsActive: true,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
NextCheck: mockedPast,
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
NextCheck: sql.NullTime{Time: mockedPast, Valid: true},
}
schedule1 := SleepSchedule{ // Next check in future -> should not be checked.
Worker: &Worker{
UUID: "4475738e-41eb-47b2-8bca-2bbcabab69bb",
Name: "дрон 2",
Status: api.WorkerStatusAwake,
},
IsActive: true,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
NextCheck: mockedFuture,
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
NextCheck: sql.NullTime{Time: mockedFuture, Valid: true},
}
schedule2 := SleepSchedule{ // Next check is zero value -> should be checked.
Worker: &Worker{
UUID: "dc251817-6a11-4548-a36a-07b0d50b4c21",
Name: "дрон 3",
Status: api.WorkerStatusAwake,
},
IsActive: true,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
NextCheck: time.Time{}, // zero value for time.
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
NextCheck: sql.NullTime{}, // zero value for time.
}
schedule3 := SleepSchedule{ // Schedule inactive -> should not be checked.
Worker: &Worker{
UUID: "874d5fc6-5784-4d43-8c20-6e7e73fc1b8d",
Name: "дрон 4",
Status: api.WorkerStatusAwake,
},
IsActive: false,
DaysOfWeek: "mo,tu,th,fr",
StartTime: TimeOfDay{18, 0},
EndTime: TimeOfDay{9, 0},
NextCheck: mockedPast, // next check in the past, so if active it would be checked.
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
NextCheck: sql.NullTime{Time: mockedPast, Valid: true}, // next check in the past, so if active it would be checked.
}
// Create the workers and sleep schedules.
scheds := []*SleepSchedule{&schedule0, &schedule1, &schedule2, &schedule3}
workers := []*Worker{worker0, worker1, worker2, worker3}
for idx := range scheds {
saveTestWorker(t, db, scheds[idx].Worker)
err := db.SetWorkerSleepSchedule(ctx, scheds[idx].Worker.UUID, scheds[idx])
err := db.SetWorkerSleepSchedule(ctx, workers[idx].UUID, scheds[idx])
require.NoError(t, err)
}
@ -293,15 +281,16 @@ func TestFetchSleepSchedulesToCheck(t *testing.T) {
require.NoError(t, err)
require.Len(t, toCheck, 2)
assertEqualSleepSchedule(t, uint(schedule0.Worker.ID), schedule0, *toCheck[0])
assert.Nil(t, toCheck[0].Worker, "the Worker should NOT be fetched")
assertEqualSleepSchedule(t, uint(schedule2.Worker.ID), schedule1, *toCheck[1])
assert.Nil(t, toCheck[1].Worker, "the Worker should NOT be fetched")
assert.Equal(t, worker0.Name, toCheck[0].WorkerName)
assert.Equal(t, worker0.UUID, toCheck[0].WorkerUUID)
assert.Equal(t, worker2.Name, toCheck[1].WorkerName)
assert.Equal(t, worker2.UUID, toCheck[1].WorkerUUID)
assertEqualSleepSchedule(t, worker0.ID, schedule0, toCheck[0].SleepSchedule)
assertEqualSleepSchedule(t, worker2.ID, schedule1, toCheck[1].SleepSchedule)
}
func assertEqualSleepSchedule(t *testing.T, workerID uint, expect, actual SleepSchedule) {
func assertEqualSleepSchedule(t *testing.T, workerID int64, expect, actual SleepSchedule) {
assert.Equal(t, workerID, actual.WorkerID, "sleep schedule is assigned to different worker")
assert.Nil(t, actual.Worker, "the Worker itself should not be fetched")
assert.Equal(t, expect.IsActive, actual.IsActive, "IsActive does not match")
assert.Equal(t, expect.DaysOfWeek, actual.DaysOfWeek, "DaysOfWeek does not match")
assert.Equal(t, expect.StartTime, actual.StartTime, "StartTime does not match")

26
internal/manager/persistence/workers.go
View File

@ -7,7 +7,6 @@ import (
"database/sql"
"errors"
"fmt"
"time"
"github.com/rs/zerolog/log"
"projects.blender.org/studio/flamenco/internal/manager/persistence/sqlc"
@ -16,29 +15,6 @@ import (
type Worker = sqlc.Worker
type Worker__gorm struct {
Model
DeletedAt sql.NullTime
UUID string
Secret string
Name string
Address string // 39 = max length of IPv6 address.
Platform string
Software string
Status api.WorkerStatus
LastSeenAt time.Time // Should contain UTC timestamps.
CanRestart bool
StatusRequested api.WorkerStatus
LazyStatusRequest bool
SupportedTaskTypes string // comma-separated list of task types.
Tags []*WorkerTag
}
func (db *DB) CreateWorker(ctx context.Context, w *Worker) error {
queries := db.queries()
@ -56,7 +32,7 @@ func (db *DB) CreateWorker(ctx context.Context, w *Worker) error {
StatusRequested: w.StatusRequested,
LazyStatusRequest: w.LazyStatusRequest,
SupportedTaskTypes: w.SupportedTaskTypes,
DeletedAt: sql.NullTime(w.DeletedAt),
DeletedAt: w.DeletedAt,
CanRestart: w.CanRestart,
}

7
internal/manager/sleep_scheduler/calculations.go
View File

@ -8,6 +8,7 @@ import (
"projects.blender.org/studio/flamenco/internal/manager/persistence"
"projects.blender.org/studio/flamenco/pkg/api"
"projects.blender.org/studio/flamenco/pkg/time_of_day"
)
// scheduledWorkerStatus returns the expected worker status at the given date/time.
@ -17,7 +18,7 @@ func scheduledWorkerStatus(now time.Time, sched *persistence.SleepSchedule) api.
return api.WorkerStatusAwake
}
tod := persistence.MakeTimeOfDay(now)
tod := time_of_day.MakeTimeOfDay(now)
if !sched.IsActive {
return api.WorkerStatusAwake
@ -56,10 +57,10 @@ func cleanupDaysOfWeek(daysOfWeek string) string {
}
// Return a timestamp when the next scheck for this schedule is due.
func calculateNextCheck(now time.Time, schedule *persistence.SleepSchedule) time.Time {
func calculateNextCheck(now time.Time, schedule persistence.SleepSchedule) time.Time {
// calcNext returns the given time of day on "today" if that hasn't passed
// yet, otherwise on "tomorrow".
calcNext := func(tod persistence.TimeOfDay) time.Time {
calcNext := func(tod time_of_day.TimeOfDay) time.Time {
nextCheck := tod.OnDate(now).In(time.Local)
if nextCheck.Before(now) {
nextCheck = nextCheck.AddDate(0, 0, 1)

37
internal/manager/sleep_scheduler/calculations_test.go
View File

@ -9,48 +9,49 @@ import (
"projects.blender.org/studio/flamenco/internal/manager/persistence"
"projects.blender.org/studio/flamenco/pkg/api"
"projects.blender.org/studio/flamenco/pkg/time_of_day"
)
func TestCalculateNextCheck(t *testing.T) {
_, mocks, _ := testFixtures(t)
var sched persistence.SleepSchedule
empty := persistence.EmptyTimeOfDay()
empty := time_of_day.Empty()
// Below, S, N, and E respectively mean Start, Now, and End times.
// Their order shows their relation to "Now". Lower-case letters mean "no value".
// Note that N can never be before 's' or after 'e'.
// S N E -> E
sched = persistence.SleepSchedule{StartTime: mkToD(9, 0), EndTime: mkToD(18, 0)}
assert.Equal(t, mocks.todayAt(18, 0), calculateNextCheck(mocks.todayAt(11, 16), &sched))
sched = persistence.SleepSchedule{StartTime: time_of_day.New(9, 0), EndTime: time_of_day.New(18, 0)}
assert.Equal(t, mocks.todayAt(18, 0), calculateNextCheck(mocks.todayAt(11, 16), sched))
// S E N -> end of day
sched = persistence.SleepSchedule{StartTime: mkToD(9, 0), EndTime: mkToD(18, 0)}
assert.Equal(t, mocks.endOfDay(), calculateNextCheck(mocks.todayAt(19, 16), &sched))
sched = persistence.SleepSchedule{StartTime: time_of_day.New(9, 0), EndTime: time_of_day.New(18, 0)}
assert.Equal(t, mocks.endOfDay(), calculateNextCheck(mocks.todayAt(19, 16), sched))
// N S E -> S
sched = persistence.SleepSchedule{StartTime: mkToD(9, 0), EndTime: mkToD(18, 0)}
assert.Equal(t, mocks.todayAt(9, 0), calculateNextCheck(mocks.todayAt(8, 47), &sched))
sched = persistence.SleepSchedule{StartTime: time_of_day.New(9, 0), EndTime: time_of_day.New(18, 0)}
assert.Equal(t, mocks.todayAt(9, 0), calculateNextCheck(mocks.todayAt(8, 47), sched))
// s N e -> end of day
sched = persistence.SleepSchedule{StartTime: empty, EndTime: empty}
assert.Equal(t, mocks.endOfDay(), calculateNextCheck(mocks.todayAt(7, 47), &sched))
assert.Equal(t, mocks.endOfDay(), calculateNextCheck(mocks.todayAt(7, 47), sched))
// S N e -> end of day
sched = persistence.SleepSchedule{StartTime: mkToD(9, 0), EndTime: empty}
assert.Equal(t, mocks.endOfDay(), calculateNextCheck(mocks.todayAt(10, 47), &sched))
sched = persistence.SleepSchedule{StartTime: time_of_day.New(9, 0), EndTime: empty}
assert.Equal(t, mocks.endOfDay(), calculateNextCheck(mocks.todayAt(10, 47), sched))
// s N E -> E
sched = persistence.SleepSchedule{StartTime: empty, EndTime: mkToD(18, 0)}
assert.Equal(t, mocks.todayAt(18, 0), calculateNextCheck(mocks.todayAt(7, 47), &sched))
sched = persistence.SleepSchedule{StartTime: empty, EndTime: time_of_day.New(18, 0)}
assert.Equal(t, mocks.todayAt(18, 0), calculateNextCheck(mocks.todayAt(7, 47), sched))
}
func TestScheduledWorkerStatus(t *testing.T) {
_, mocks, _ := testFixtures(t)
var sched persistence.SleepSchedule
empty := persistence.EmptyTimeOfDay()
empty := time_of_day.Empty()
// No schedule means 'awake'.
assert.Equal(t, api.WorkerStatusAwake, scheduledWorkerStatus(mocks.todayAt(11, 16), nil))
@ -63,7 +64,7 @@ func TestScheduledWorkerStatus(t *testing.T) {
// to each day.
// S N E -> asleep
sched = persistence.SleepSchedule{StartTime: mkToD(9, 0), EndTime: mkToD(18, 0), IsActive: true}
sched = persistence.SleepSchedule{StartTime: time_of_day.New(9, 0), EndTime: time_of_day.New(18, 0), IsActive: true}
assert.Equal(t, api.WorkerStatusAsleep, scheduledWorkerStatus(mocks.todayAt(11, 16), &sched))
// S E N -> awake
@ -77,11 +78,11 @@ func TestScheduledWorkerStatus(t *testing.T) {
assert.Equal(t, api.WorkerStatusAsleep, scheduledWorkerStatus(mocks.todayAt(7, 47), &sched))
// S N e -> asleep
sched = persistence.SleepSchedule{StartTime: mkToD(9, 0), EndTime: empty, IsActive: true}
sched = persistence.SleepSchedule{StartTime: time_of_day.New(9, 0), EndTime: empty, IsActive: true}
assert.Equal(t, api.WorkerStatusAsleep, scheduledWorkerStatus(mocks.todayAt(10, 47), &sched))
// s N E -> asleep
sched = persistence.SleepSchedule{StartTime: empty, EndTime: mkToD(18, 0), IsActive: true}
sched = persistence.SleepSchedule{StartTime: empty, EndTime: time_of_day.New(18, 0), IsActive: true}
assert.Equal(t, api.WorkerStatusAsleep, scheduledWorkerStatus(mocks.todayAt(7, 47), &sched))
// Test DaysOfWeek logic, but only with explicit start & end times. The logic
@ -89,7 +90,7 @@ func TestScheduledWorkerStatus(t *testing.T) {
// The mocked "today" is a Tuesday.
// S N E unmentioned day -> awake
sched = persistence.SleepSchedule{DaysOfWeek: "mo we", StartTime: mkToD(9, 0), EndTime: mkToD(18, 0), IsActive: true}
sched = persistence.SleepSchedule{DaysOfWeek: "mo we", StartTime: time_of_day.New(9, 0), EndTime: time_of_day.New(18, 0), IsActive: true}
assert.Equal(t, api.WorkerStatusAwake, scheduledWorkerStatus(mocks.todayAt(11, 16), &sched))
// S E N unmentioned day -> awake
@ -99,7 +100,7 @@ func TestScheduledWorkerStatus(t *testing.T) {
assert.Equal(t, api.WorkerStatusAwake, scheduledWorkerStatus(mocks.todayAt(8, 47), &sched))
// S N E mentioned day -> asleep
sched = persistence.SleepSchedule{DaysOfWeek: "tu th fr", StartTime: mkToD(9, 0), EndTime: mkToD(18, 0), IsActive: true}
sched = persistence.SleepSchedule{DaysOfWeek: "tu th fr", StartTime: time_of_day.New(9, 0), EndTime: time_of_day.New(18, 0), IsActive: true}
assert.Equal(t, api.WorkerStatusAsleep, scheduledWorkerStatus(mocks.todayAt(11, 16), &sched))
// S E N mentioned day -> awake

7
internal/manager/sleep_scheduler/interfaces.go
View File

@ -16,11 +16,10 @@ import (
type PersistenceService interface {
FetchWorkerSleepSchedule(ctx context.Context, workerUUID string) (*persistence.SleepSchedule, error)
SetWorkerSleepSchedule(ctx context.Context, workerUUID string, schedule *persistence.SleepSchedule) error
// FetchSleepScheduleWorker sets the given schedule's `Worker` pointer.
FetchSleepScheduleWorker(ctx context.Context, schedule *persistence.SleepSchedule) error
FetchSleepSchedulesToCheck(ctx context.Context) ([]*persistence.SleepSchedule, error)
FetchSleepScheduleWorker(ctx context.Context, schedule persistence.SleepSchedule) (*persistence.Worker, error)
FetchSleepSchedulesToCheck(ctx context.Context) ([]persistence.SleepScheduleOwned, error)
SetWorkerSleepScheduleNextCheck(ctx context.Context, schedule *persistence.SleepSchedule) error
SetWorkerSleepScheduleNextCheck(ctx context.Context, schedule persistence.SleepSchedule) error
SaveWorkerStatus(ctx context.Context, w *persistence.Worker) error
}

33
internal/manager/sleep_scheduler/mocks/interfaces_mock.gen.go generated
View File

@ -37,14 +37,29 @@ func (m *MockPersistenceService) EXPECT() *MockPersistenceServiceMockRecorder {
return m.recorder
}
// FetchSleepScheduleWorker mocks base method.
func (m *MockPersistenceService) FetchSleepScheduleWorker(arg0 context.Context, arg1 *persistence.SleepSchedule) error {
// CreateWorker mocks base method.
func (m *MockPersistenceService) CreateWorker(arg0 context.Context, arg1 *sqlc.Worker) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "FetchSleepScheduleWorker", arg0, arg1)
ret := m.ctrl.Call(m, "CreateWorker", arg0, arg1)
ret0, _ := ret[0].(error)
return ret0
}
// CreateWorker indicates an expected call of CreateWorker.
func (mr *MockPersistenceServiceMockRecorder) CreateWorker(arg0, arg1 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CreateWorker", reflect.TypeOf((*MockPersistenceService)(nil).CreateWorker), arg0, arg1)
}
// FetchSleepScheduleWorker mocks base method.
func (m *MockPersistenceService) FetchSleepScheduleWorker(arg0 context.Context, arg1 sqlc.SleepSchedule) (*sqlc.Worker, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "FetchSleepScheduleWorker", arg0, arg1)
ret0, _ := ret[0].(*sqlc.Worker)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// FetchSleepScheduleWorker indicates an expected call of FetchSleepScheduleWorker.
func (mr *MockPersistenceServiceMockRecorder) FetchSleepScheduleWorker(arg0, arg1 interface{}) *gomock.Call {
mr.mock.ctrl.T.Helper()
@ -52,10 +67,10 @@ func (mr *MockPersistenceServiceMockRecorder) FetchSleepScheduleWorker(arg0, arg
}
// FetchSleepSchedulesToCheck mocks base method.
func (m *MockPersistenceService) FetchSleepSchedulesToCheck(arg0 context.Context) ([]*persistence.SleepSchedule, error) {
func (m *MockPersistenceService) FetchSleepSchedulesToCheck(arg0 context.Context) ([]persistence.SleepScheduleOwned, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "FetchSleepSchedulesToCheck", arg0)
ret0, _ := ret[0].([]*persistence.SleepSchedule)
ret0, _ := ret[0].([]persistence.SleepScheduleOwned)
ret1, _ := ret[1].(error)
return ret0, ret1
}
@ -67,10 +82,10 @@ func (mr *MockPersistenceServiceMockRecorder) FetchSleepSchedulesToCheck(arg0 in
}
// FetchWorkerSleepSchedule mocks base method.
func (m *MockPersistenceService) FetchWorkerSleepSchedule(arg0 context.Context, arg1 string) (*persistence.SleepSchedule, error) {
func (m *MockPersistenceService) FetchWorkerSleepSchedule(arg0 context.Context, arg1 string) (*sqlc.SleepSchedule, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "FetchWorkerSleepSchedule", arg0, arg1)
ret0, _ := ret[0].(*persistence.SleepSchedule)
ret0, _ := ret[0].(*sqlc.SleepSchedule)
ret1, _ := ret[1].(error)
return ret0, ret1
}
@ -96,7 +111,7 @@ func (mr *MockPersistenceServiceMockRecorder) SaveWorkerStatus(arg0, arg1 interf
}
// SetWorkerSleepSchedule mocks base method.
func (m *MockPersistenceService) SetWorkerSleepSchedule(arg0 context.Context, arg1 string, arg2 *persistence.SleepSchedule) error {
func (m *MockPersistenceService) SetWorkerSleepSchedule(arg0 context.Context, arg1 string, arg2 *sqlc.SleepSchedule) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SetWorkerSleepSchedule", arg0, arg1, arg2)
ret0, _ := ret[0].(error)
@ -110,7 +125,7 @@ func (mr *MockPersistenceServiceMockRecorder) SetWorkerSleepSchedule(arg0, arg1,
}
// SetWorkerSleepScheduleNextCheck mocks base method.
func (m *MockPersistenceService) SetWorkerSleepScheduleNextCheck(arg0 context.Context, arg1 *persistence.SleepSchedule) error {
func (m *MockPersistenceService) SetWorkerSleepScheduleNextCheck(arg0 context.Context, arg1 sqlc.SleepSchedule) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "SetWorkerSleepScheduleNextCheck", arg0, arg1)
ret0, _ := ret[0].(error)

45
internal/manager/sleep_scheduler/sleep_scheduler.go
View File

@ -66,7 +66,7 @@ func (ss *SleepScheduler) FetchSchedule(ctx context.Context, workerUUID string)
// SetSleepSchedule stores the given schedule as the worker's new sleep schedule.
// The new schedule is immediately applied to the Worker.
func (ss *SleepScheduler) SetSchedule(ctx context.Context, workerUUID string, schedule *persistence.SleepSchedule) error {
func (ss *SleepScheduler) SetSchedule(ctx context.Context, workerUUID string, schedule persistence.SleepSchedule) error {
// Ensure 'start' actually preceeds 'end'.
if schedule.StartTime.HasValue() &&
schedule.EndTime.HasValue() &&
@ -75,15 +75,15 @@ func (ss *SleepScheduler) SetSchedule(ctx context.Context, workerUUID string, sc
}
schedule.DaysOfWeek = cleanupDaysOfWeek(schedule.DaysOfWeek)
schedule.NextCheck = ss.calculateNextCheck(schedule)
schedule.SetNextCheck(ss.calculateNextCheck(schedule))
if err := ss.persist.SetWorkerSleepSchedule(ctx, workerUUID, schedule); err != nil {
if err := ss.persist.SetWorkerSleepSchedule(ctx, workerUUID, &schedule); err != nil {
return fmt.Errorf("persisting sleep schedule of worker %s: %w", workerUUID, err)
}
logger := addLoggerFields(zerolog.Ctx(ctx), schedule)
logger := addLoggerFields(zerolog.Ctx(ctx), schedule, workerUUID, "")
logger.Info().
Str("worker", schedule.Worker.Identifier()).
Str("worker", workerUUID).
Msg("sleep scheduler: new schedule for worker")
return ss.ApplySleepSchedule(ctx, schedule)
@ -106,28 +106,24 @@ func (ss *SleepScheduler) scheduledWorkerStatus(sched *persistence.SleepSchedule
}
// Return a timestamp when the next scheck for this schedule is due.
func (ss *SleepScheduler) calculateNextCheck(schedule *persistence.SleepSchedule) time.Time {
func (ss *SleepScheduler) calculateNextCheck(schedule persistence.SleepSchedule) time.Time {
now := ss.clock.Now()
return calculateNextCheck(now, schedule)
}
// ApplySleepSchedule sets worker.StatusRequested if the scheduler demands a status change.
func (ss *SleepScheduler) ApplySleepSchedule(ctx context.Context, schedule *persistence.SleepSchedule) error {
func (ss *SleepScheduler) ApplySleepSchedule(ctx context.Context, schedule persistence.SleepSchedule) error {
// Find the Worker managed by this schedule.
worker := schedule.Worker
if worker == nil {
err := ss.persist.FetchSleepScheduleWorker(ctx, schedule)
worker, err := ss.persist.FetchSleepScheduleWorker(ctx, schedule)
if err != nil {
return err
}
worker = schedule.Worker
}
if !ss.mayUpdateWorker(worker) {
return nil
}
scheduled := ss.scheduledWorkerStatus(schedule)
scheduled := ss.scheduledWorkerStatus(&schedule)
if scheduled == "" ||
(worker.StatusRequested == scheduled && !worker.LazyStatusRequest) ||
(worker.Status == scheduled && worker.StatusRequested == "") {
@ -204,9 +200,11 @@ func (ss *SleepScheduler) CheckSchedules(ctx context.Context) {
}
}
func (ss *SleepScheduler) checkSchedule(ctx context.Context, schedule *persistence.SleepSchedule) {
func (ss *SleepScheduler) checkSchedule(ctx context.Context, ownedSchedule persistence.SleepScheduleOwned) {
// Compute the next time to check.
schedule.NextCheck = ss.calculateNextCheck(schedule)
schedule := ownedSchedule.SleepSchedule
schedule.SetNextCheck(ss.calculateNextCheck(schedule))
err := ss.persist.SetWorkerSleepScheduleNextCheck(ctx, schedule)
switch {
case errors.Is(ctx.Err(), context.Canceled):
@ -215,7 +213,8 @@ func (ss *SleepScheduler) checkSchedule(ctx context.Context, schedule *persisten
case err != nil:
log.Error().
Err(err).
Str("worker", schedule.Worker.Identifier()).
Str("workerName", ownedSchedule.WorkerName).
Str("workerUUID", ownedSchedule.WorkerUUID).
Msg("sleep scheduler: error refreshing worker's sleep schedule")
return
}
@ -230,12 +229,13 @@ func (ss *SleepScheduler) checkSchedule(ctx context.Context, schedule *persisten
// soft-deleted (and thus foreign key constraints don't trigger deletion of
// the sleep schedule).
log.Debug().
Uint("worker", schedule.WorkerID).
Int64("worker", schedule.WorkerID).
Msg("sleep scheduler: sleep schedule's owning worker cannot be found; not applying the schedule")
case err != nil:
log.Error().
Err(err).
Str("worker", schedule.Worker.Identifier()).
Str("workerName", ownedSchedule.WorkerName).
Str("workerUUID", ownedSchedule.WorkerUUID).
Msg("sleep scheduler: error applying worker's sleep schedule")
}
}
@ -246,11 +246,14 @@ func (ss *SleepScheduler) mayUpdateWorker(worker *persistence.Worker) bool {
return !shouldSkip
}
func addLoggerFields(logger *zerolog.Logger, schedule *persistence.SleepSchedule) zerolog.Logger {
func addLoggerFields(logger *zerolog.Logger, schedule persistence.SleepSchedule, workerUUID, workerName string) zerolog.Logger {
logCtx := logger.With()
if schedule.Worker != nil {
logCtx = logCtx.Str("worker", schedule.Worker.Identifier())
if workerUUID != "" {
logCtx = logCtx.Str("workerUUID", workerUUID)
}
if workerName != "" {
logCtx = logCtx.Str("workerName", workerName)
}
logCtx = logCtx.

120
internal/manager/sleep_scheduler/sleep_scheduler_test.go
View File

@ -4,6 +4,7 @@ package sleep_scheduler
import (
"context"
"database/sql"
"testing"
"time"
@ -15,6 +16,7 @@ import (
"projects.blender.org/studio/flamenco/internal/manager/persistence"
"projects.blender.org/studio/flamenco/internal/manager/sleep_scheduler/mocks"
"projects.blender.org/studio/flamenco/pkg/api"
"projects.blender.org/studio/flamenco/pkg/time_of_day"
)
func TestFetchSchedule(t *testing.T) {
@ -33,22 +35,24 @@ func TestSetSchedule(t *testing.T) {
ss, mocks, ctx := testFixtures(t)
workerUUID := "aeb49d8a-6903-41b3-b545-77b7a1c0ca19"
worker := persistence.Worker{
UUID: workerUUID,
Status: api.WorkerStatusAwake,
}
sched := persistence.SleepSchedule{
IsActive: true,
DaysOfWeek: " mo tu we",
StartTime: mkToD(9, 0),
EndTime: mkToD(18, 0),
Worker: &persistence.Worker{
UUID: workerUUID,
Status: api.WorkerStatusAwake,
},
StartTime: time_of_day.New(9, 0),
EndTime: time_of_day.New(18, 0),
WorkerID: worker.ID,
}
expectSavedSchedule := sched
expectSavedSchedule.DaysOfWeek = "mo tu we" // Expect a cleanup
expectNextCheck := mocks.todayAt(18, 0) // "now" is at 11:14:47, expect a check at the end time.
expectSavedSchedule.NextCheck = expectNextCheck
expectSavedSchedule.SetNextCheck(expectNextCheck)
mocks.persist.EXPECT().FetchSleepScheduleWorker(ctx, expectSavedSchedule).Return(&worker, nil)
// Expect the new schedule to be saved.
mocks.persist.EXPECT().SetWorkerSleepSchedule(ctx, workerUUID, &expectSavedSchedule)
@ -58,40 +62,41 @@ func TestSetSchedule(t *testing.T) {
mocks.persist.EXPECT().SaveWorkerStatus(ctx, gomock.Any())
mocks.broadcaster.EXPECT().BroadcastWorkerUpdate(gomock.Any())
err := ss.SetSchedule(ctx, workerUUID, &sched)
err := ss.SetSchedule(ctx, workerUUID, sched)
require.NoError(t, err)
}
func TestSetScheduleSwappedStartEnd(t *testing.T) {
ss, mocks, ctx := testFixtures(t)
workerUUID := "aeb49d8a-6903-41b3-b545-77b7a1c0ca19"
// Worker already in the right state, so no saving/broadcasting expected.
worker := persistence.Worker{
ID: 47,
UUID: "aeb49d8a-6903-41b3-b545-77b7a1c0ca19",
Name: "test worker",
Status: api.WorkerStatusAsleep,
}
sched := persistence.SleepSchedule{
IsActive: true,
DaysOfWeek: "mo tu we",
StartTime: mkToD(18, 0),
EndTime: mkToD(9, 0),
// Worker already in the right state, so no saving/broadcasting expected.
Worker: &persistence.Worker{
UUID: workerUUID,
Status: api.WorkerStatusAsleep,
},
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
WorkerID: worker.ID,
}
expectSavedSchedule := persistence.SleepSchedule{
IsActive: true,
DaysOfWeek: "mo tu we",
StartTime: mkToD(9, 0), // Expect start and end time to be corrected.
EndTime: mkToD(18, 0),
NextCheck: mocks.todayAt(18, 0), // "now" is at 11:14:47, expect a check at the end time.
Worker: sched.Worker,
StartTime: time_of_day.New(9, 0), // Expect start and end time to be corrected.
EndTime: time_of_day.New(18, 0),
NextCheck: sql.NullTime{Time: mocks.todayAt(18, 0), Valid: true}, // "now" is at 11:14:47, expect a check at the end time.
WorkerID: worker.ID,
}
mocks.persist.EXPECT().SetWorkerSleepSchedule(ctx, workerUUID, &expectSavedSchedule)
mocks.persist.EXPECT().FetchSleepScheduleWorker(ctx, expectSavedSchedule).Return(&worker, nil)
mocks.persist.EXPECT().SetWorkerSleepSchedule(ctx, worker.UUID, &expectSavedSchedule)
err := ss.SetSchedule(ctx, workerUUID, &sched)
err := ss.SetSchedule(ctx, worker.UUID, sched)
require.NoError(t, err)
}
@ -99,19 +104,36 @@ func TestSetScheduleSwappedStartEnd(t *testing.T) {
func TestCheckSleepScheduleAtShutdown(t *testing.T) {
ss, mocks, _ := testFixtures(t)
worker := persistence.Worker{
ID: 47,
UUID: "aeb49d8a-6903-41b3-b545-77b7a1c0ca19",
Name: "test worker",
Status: api.WorkerStatusAsleep,
}
sched := persistence.SleepSchedule{
IsActive: true,
DaysOfWeek: "mo tu we",
StartTime: mkToD(18, 0),
EndTime: mkToD(9, 0),
Worker: nil,
StartTime: time_of_day.New(18, 0),
EndTime: time_of_day.New(9, 0),
}
// Construct the updated-and-about-to-be-saved schedule.
updatedSched := sched
updatedSched.NextCheck = sql.NullTime{
Time: sched.StartTime.OnDate(mocks.clock.Now()),
Valid: true,
}
// Cancel the context to mimick the Manager shutting down.
ctx, cancel := context.WithCancel(context.Background())
cancel()
mocks.persist.EXPECT().SetWorkerSleepScheduleNextCheck(ctx, &sched).Return(context.Canceled)
ss.checkSchedule(ctx, &sched)
mocks.persist.EXPECT().SetWorkerSleepScheduleNextCheck(ctx, updatedSched).Return(context.Canceled)
ss.checkSchedule(ctx, persistence.SleepScheduleOwned{
SleepSchedule: sched,
WorkerName: worker.Name,
WorkerUUID: worker.UUID,
})
}
func TestApplySleepSchedule(t *testing.T) {
@ -126,8 +148,8 @@ func TestApplySleepSchedule(t *testing.T) {
sched := persistence.SleepSchedule{
IsActive: true,
DaysOfWeek: "mo tu we",
StartTime: mkToD(9, 0),
EndTime: mkToD(18, 0),
StartTime: time_of_day.New(9, 0),
EndTime: time_of_day.New(18, 0),
}
testForExpectedStatus := func(expectedNewStatus api.WorkerStatus) {
@ -136,11 +158,7 @@ func TestApplySleepSchedule(t *testing.T) {
testWorker := worker
// Expect the Worker to be fetched.
mocks.persist.EXPECT().FetchSleepScheduleWorker(ctx, &testSchedule).DoAndReturn(
func(ctx context.Context, schedule *persistence.SleepSchedule) error {
schedule.Worker = &testWorker
return nil
})
mocks.persist.EXPECT().FetchSleepScheduleWorker(ctx, testSchedule).Return(&testWorker, nil)
// Construct the worker as we expect it to be saved to the database.
savedWorker := testWorker
@ -156,7 +174,7 @@ func TestApplySleepSchedule(t *testing.T) {
})
// Actually apply the sleep schedule.
err := ss.ApplySleepSchedule(ctx, &testSchedule)
err := ss.ApplySleepSchedule(ctx, testSchedule)
require.NoError(t, err)
// Check the SocketIO broadcast.
@ -200,8 +218,8 @@ func TestApplySleepScheduleNoStatusChange(t *testing.T) {
sched := persistence.SleepSchedule{
IsActive: true,
DaysOfWeek: "mo tu we",
StartTime: mkToD(9, 0),
EndTime: mkToD(18, 0),
StartTime: time_of_day.New(9, 0),
EndTime: time_of_day.New(18, 0),
}
runTest := func() {
@ -210,14 +228,10 @@ func TestApplySleepScheduleNoStatusChange(t *testing.T) {
testWorker := worker
// Expect the Worker to be fetched.
mocks.persist.EXPECT().FetchSleepScheduleWorker(ctx, &testSchedule).DoAndReturn(
func(ctx context.Context, schedule *persistence.SleepSchedule) error {
schedule.Worker = &testWorker
return nil
})
mocks.persist.EXPECT().FetchSleepScheduleWorker(ctx, testSchedule).Return(&testWorker, nil)
// Apply the sleep schedule. This should not trigger any persistence or broadcasts.
err := ss.ApplySleepSchedule(ctx, &testSchedule)
err := ss.ApplySleepSchedule(ctx, testSchedule)
require.NoError(t, err)
}
@ -284,6 +298,16 @@ func testFixtures(t *testing.T) (*SleepScheduler, TestMocks, context.Context) {
return ss, mocks, ctx
}
func mkToD(hour, minute int) persistence.TimeOfDay {
return persistence.TimeOfDay{Hour: hour, Minute: minute}
func createTestWorker(updaters ...func(*persistence.Worker)) persistence.Worker {
w := persistence.Worker{
ID: 47,
UUID: "4f2d3755-c365-429f-8017-44356427c069",
Name: "schedule test worker",
}
for _, updater := range updaters {
updater(&w)
}
return w
}

14
internal/manager/persistence/time_of_day.go → pkg/time_of_day/time_of_day.go
View File

@ -1,4 +1,4 @@
package persistence
package time_of_day
// SPDX-License-Identifier: GPL-3.0-or-later
@ -27,14 +27,22 @@ type TimeOfDay struct {
Minute int
}
// New returns a new TimeOfDay.
func New(Hour, Minute int) TimeOfDay {
return TimeOfDay{
Hour: Hour,
Minute: Minute,
}
}
// MakeTimeOfDay converts a time.Time into a TimeOfDay.
func MakeTimeOfDay(someTime time.Time) TimeOfDay {
return TimeOfDay{someTime.Hour(), someTime.Minute()}
}
// EmptyTimeOfDay returns a TimeOfDay struct with no value.
// Empty returns a TimeOfDay struct with no value.
// See `TimeOfDay.HasValue()`.
func EmptyTimeOfDay() TimeOfDay {
func Empty() TimeOfDay {
return TimeOfDay{Hour: timeOfDayNoValue, Minute: timeOfDayNoValue}
}

2
internal/manager/persistence/time_of_day_test.go → pkg/time_of_day/time_of_day_test.go
View File

@ -1,4 +1,4 @@
package persistence
package time_of_day
// SPDX-License-Identifier: GPL-3.0-or-later

16
sqlc.yaml
View File

@ -19,6 +19,10 @@ sql:
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: workers.status_requested
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: sleep_schedules.start_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
- column: sleep_schedules.end_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
rename:
uuid: "UUID"
uuids: "UUIDs"
@ -44,6 +48,10 @@ sql:
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: workers.status_requested
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: sleep_schedules.start_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
- column: sleep_schedules.end_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
rename:
uuid: "UUID"
uuids: "UUIDs"
@ -69,6 +77,10 @@ sql:
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: workers.status_requested
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: sleep_schedules.start_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
- column: sleep_schedules.end_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
rename:
uuid: "UUID"
uuids: "UUIDs"
@ -94,6 +106,10 @@ sql:
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: workers.status_requested
go_type: { type: "WorkerStatus", import: "projects.blender.org/studio/flamenco/pkg/api" }
- column: sleep_schedules.start_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
- column: sleep_schedules.end_time
go_type: { type: "TimeOfDay", import: "projects.blender.org/studio/flamenco/pkg/time_of_day" }
rename:
uuid: "UUID"
uuids: "UUIDs"