Compare commits
1 Commits
b8a55cca82
...
847ef923bf
Author | SHA1 | Date | |
---|---|---|---|
|
847ef923bf |
@ -6,10 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||||||
|
|
||||||
## [Unreleased]
|
## [Unreleased]
|
||||||
|
|
||||||
## [0.1.0] - 2025-01-18
|
## [0.1.0] - 2025-01-17
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
- add worker distributor and model registry
|
|
||||||
- enable worker process
|
- enable worker process
|
||||||
- add migration
|
- add migration
|
||||||
- add executor (#3)
|
- add executor (#3)
|
||||||
|
1
go.mod
1
go.mod
@ -3,7 +3,6 @@ module git.front.kjuulh.io/kjuulh/orbis
|
|||||||
go 1.23.4
|
go 1.23.4
|
||||||
|
|
||||||
require (
|
require (
|
||||||
github.com/adhocore/gronx v1.19.5
|
|
||||||
github.com/golang-migrate/migrate/v4 v4.18.1
|
github.com/golang-migrate/migrate/v4 v4.18.1
|
||||||
github.com/google/uuid v1.6.0
|
github.com/google/uuid v1.6.0
|
||||||
github.com/jackc/pgx/v5 v5.7.2
|
github.com/jackc/pgx/v5 v5.7.2
|
||||||
|
2
go.sum
2
go.sum
@ -2,8 +2,6 @@ github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161 h1:L/gRVlceqvL25
|
|||||||
github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
|
github.com/Azure/go-ansiterm v0.0.0-20230124172434-306776ec8161/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
|
||||||
github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
|
github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
|
||||||
github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
|
github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
|
||||||
github.com/adhocore/gronx v1.19.5 h1:cwIG4nT1v9DvadxtHBe6MzE+FZ1JDvAUC45U2fl4eSQ=
|
|
||||||
github.com/adhocore/gronx v1.19.5/go.mod h1:7oUY1WAU8rEJWmAxXR2DN0JaO4gi9khSgKjiRypqteg=
|
|
||||||
github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
|
github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
|
||||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||||
|
@ -4,11 +4,8 @@ import (
|
|||||||
"log/slog"
|
"log/slog"
|
||||||
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/executor"
|
"git.front.kjuulh.io/kjuulh/orbis/internal/executor"
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/modelschedule"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/scheduler"
|
"git.front.kjuulh.io/kjuulh/orbis/internal/scheduler"
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/worker"
|
"git.front.kjuulh.io/kjuulh/orbis/internal/worker"
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/workprocessor"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/workscheduler"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type App struct {
|
type App struct {
|
||||||
@ -30,27 +27,9 @@ func (a *App) Scheduler() *scheduler.Scheduler {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (a *App) Executor() *executor.Executor {
|
func (a *App) Executor() *executor.Executor {
|
||||||
return executor.NewExecutor(
|
return executor.NewExecutor(a.logger.With("component", "executor"))
|
||||||
a.logger.With("component", "executor"),
|
|
||||||
ModelRegistry(),
|
|
||||||
a.ModelSchedule(),
|
|
||||||
a.Worker(),
|
|
||||||
a.WorkScheduler(),
|
|
||||||
)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (a *App) Worker() *worker.Worker {
|
func (a *App) Worker() *worker.Worker {
|
||||||
return worker.NewWorker(Postgres(), a.logger, a.WorkProcessor())
|
return worker.NewWorker(Postgres(), a.logger)
|
||||||
}
|
|
||||||
|
|
||||||
func (a *App) WorkScheduler() *workscheduler.WorkScheduler {
|
|
||||||
return workscheduler.NewWorkScheduler(Postgres(), a.logger)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (a *App) WorkProcessor() *workprocessor.WorkProcessor {
|
|
||||||
return workprocessor.NewWorkProcessor(a.WorkScheduler(), a.logger)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (a *App) ModelSchedule() *modelschedule.ModelSchedule {
|
|
||||||
return modelschedule.NewModelSchedule(a.logger, Postgres())
|
|
||||||
}
|
}
|
||||||
|
@ -1,10 +0,0 @@
|
|||||||
package app
|
|
||||||
|
|
||||||
import (
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/modelregistry"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/utilities"
|
|
||||||
)
|
|
||||||
|
|
||||||
var ModelRegistry = utilities.Singleton(func() (*modelregistry.ModelRegistry, error) {
|
|
||||||
return modelregistry.NewModelRegistry(), nil
|
|
||||||
})
|
|
@ -2,84 +2,26 @@ package executor
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"fmt"
|
|
||||||
"log/slog"
|
"log/slog"
|
||||||
"time"
|
|
||||||
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/modelregistry"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/modelschedule"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/worker"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/workscheduler"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
type Executor struct {
|
type Executor struct {
|
||||||
logger *slog.Logger
|
logger *slog.Logger
|
||||||
|
|
||||||
modelRegistry *modelregistry.ModelRegistry
|
|
||||||
modelSchedule *modelschedule.ModelSchedule
|
|
||||||
worker *worker.Worker
|
|
||||||
workerscheduler *workscheduler.WorkScheduler
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewExecutor(
|
func NewExecutor(logger *slog.Logger) *Executor {
|
||||||
logger *slog.Logger,
|
|
||||||
modelRegistry *modelregistry.ModelRegistry,
|
|
||||||
modelSchedule *modelschedule.ModelSchedule,
|
|
||||||
worker *worker.Worker,
|
|
||||||
workerscheduler *workscheduler.WorkScheduler,
|
|
||||||
) *Executor {
|
|
||||||
return &Executor{
|
return &Executor{
|
||||||
logger: logger,
|
logger: logger,
|
||||||
|
|
||||||
modelRegistry: modelRegistry,
|
|
||||||
modelSchedule: modelSchedule,
|
|
||||||
worker: worker,
|
|
||||||
workerscheduler: workerscheduler,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (e *Executor) DispatchEvents(ctx context.Context) error {
|
func (e *Executor) DispatchEvents(ctx context.Context) error {
|
||||||
e.logger.InfoContext(ctx, "dispatching events")
|
e.logger.InfoContext(ctx, "dispatching events")
|
||||||
|
|
||||||
start := time.Now().Add(-time.Second * 30)
|
// TODO: Process updates to models
|
||||||
end := time.Now()
|
// TODO: Insert new cron for runtime
|
||||||
|
// TODO: Calculate time since last run
|
||||||
models, err := e.modelRegistry.GetModels()
|
// TODO: Send events for workers to pick up
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("failed to get models from registry: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
registeredWorkers, err := e.worker.GetWorkers(ctx)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("failed to find workers: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
workers, err := e.workerscheduler.GetWorkers(ctx, registeredWorkers)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("failed to find workers: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
for workers := range workers.IterateSlice(2000) {
|
|
||||||
for _, model := range models {
|
|
||||||
modelRuns, lastRun, err := e.modelSchedule.GetNext(ctx, model, start, end, uint(len(workers)))
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
for i, modelRun := range modelRuns {
|
|
||||||
worker := workers[i]
|
|
||||||
e.logger.DebugContext(ctx, "dispatching model run", "modelRun", modelRun.Model.Name, "start", modelRun.Start, "end", modelRun.End)
|
|
||||||
|
|
||||||
if err := e.workerscheduler.InsertModelRun(ctx, worker, &modelRun); err != nil {
|
|
||||||
return fmt.Errorf("failed to register model run: %w", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := e.modelSchedule.UpdateModelRun(ctx, model, lastRun); err != nil {
|
|
||||||
return fmt.Errorf("failed to update checkpoint for model: %w", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
@ -1,105 +0,0 @@
|
|||||||
package modelregistry
|
|
||||||
|
|
||||||
import "sync"
|
|
||||||
|
|
||||||
type Model struct {
|
|
||||||
Name string `json:"name"`
|
|
||||||
Schedule string `json:"schedule"`
|
|
||||||
Lookup string `json:"lookup"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type ModelRegistry struct {
|
|
||||||
models []Model
|
|
||||||
modelsLock sync.RWMutex
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewModelRegistry() *ModelRegistry {
|
|
||||||
return &ModelRegistry{
|
|
||||||
models: make([]Model, 0, 128), // We start off with a capacity of 128 models. Makes memory more deterministic
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *ModelRegistry) GetModels() ([]Model, error) {
|
|
||||||
m.modelsLock.RLock()
|
|
||||||
defer m.modelsLock.RUnlock()
|
|
||||||
|
|
||||||
return []Model{
|
|
||||||
{
|
|
||||||
Name: "69C42481-650D-46E0-9C96-3D61B96565FB",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "5B0F96E5-BC37-427E-B615-E635156386F0",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "5A511180-6613-4F8E-9125-2E8FE272E03C",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "4EFE77E8-082B-4828-8527-635E5B6253D9",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "E53DCA1E-641B-421A-9DB6-1A6F09F3D96D",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "63BDC98C-ECBA-44FD-BFAE-056B1C004077",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "C18A1948-0045-4099-AC58-5B7C587AC0F0",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "8B87FDB5-A119-43C0-8D15-9B517577A8AE",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "4E121E78-1CBD-4BC1-8A10-14354B76E553",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "14AF7CDF-783F-4DFE-8B3D-E4C23C12AEDC",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "60FE99E9-4EF7-40A5-A19D-9A439BA12B24",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "622C03C2-CAF7-4708-B26D-D536E3C3F4DD",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "E7BC4A8D-FFF6-4A8D-A48B-6569340746E4",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "7CB66BA1-CF1E-4FA6-8C32-F048D88FCE54",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
{
|
|
||||||
Name: "82518E22-EFED-4AA8-AC19-FF3D81ECE609",
|
|
||||||
Schedule: "* * * * * *",
|
|
||||||
Lookup: "something",
|
|
||||||
},
|
|
||||||
}, nil
|
|
||||||
|
|
||||||
//return m.models, nil
|
|
||||||
}
|
|
@ -1,18 +0,0 @@
|
|||||||
-- name: Ping :one
|
|
||||||
SELECT 1;
|
|
||||||
|
|
||||||
-- name: GetLast :one
|
|
||||||
SELECT last_run
|
|
||||||
FROM
|
|
||||||
model_schedules
|
|
||||||
WHERE
|
|
||||||
model_name = $1
|
|
||||||
LIMIT 1;
|
|
||||||
|
|
||||||
-- name: UpsertModel :exec
|
|
||||||
INSERT INTO model_schedules (model_name, last_run)
|
|
||||||
VALUES ($1, $2)
|
|
||||||
ON CONFLICT (model_name)
|
|
||||||
DO UPDATE SET
|
|
||||||
last_run = excluded.last_run;
|
|
||||||
|
|
@ -1,32 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
|
|
||||||
"github.com/jackc/pgx/v5"
|
|
||||||
"github.com/jackc/pgx/v5/pgconn"
|
|
||||||
)
|
|
||||||
|
|
||||||
type DBTX interface {
|
|
||||||
Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
|
|
||||||
Query(context.Context, string, ...interface{}) (pgx.Rows, error)
|
|
||||||
QueryRow(context.Context, string, ...interface{}) pgx.Row
|
|
||||||
}
|
|
||||||
|
|
||||||
func New(db DBTX) *Queries {
|
|
||||||
return &Queries{db: db}
|
|
||||||
}
|
|
||||||
|
|
||||||
type Queries struct {
|
|
||||||
db DBTX
|
|
||||||
}
|
|
||||||
|
|
||||||
func (q *Queries) WithTx(tx pgx.Tx) *Queries {
|
|
||||||
return &Queries{
|
|
||||||
db: tx,
|
|
||||||
}
|
|
||||||
}
|
|
@ -1,30 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"github.com/google/uuid"
|
|
||||||
"github.com/jackc/pgx/v5/pgtype"
|
|
||||||
)
|
|
||||||
|
|
||||||
type ModelSchedule struct {
|
|
||||||
ModelName string `json:"model_name"`
|
|
||||||
LastRun pgtype.Timestamptz `json:"last_run"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type WorkSchedule struct {
|
|
||||||
ScheduleID uuid.UUID `json:"schedule_id"`
|
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
|
||||||
StartRun pgtype.Timestamptz `json:"start_run"`
|
|
||||||
EndRun pgtype.Timestamptz `json:"end_run"`
|
|
||||||
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
|
|
||||||
State string `json:"state"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type WorkerRegister struct {
|
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
|
||||||
Capacity int32 `json:"capacity"`
|
|
||||||
HeartBeat pgtype.Timestamptz `json:"heart_beat"`
|
|
||||||
}
|
|
@ -1,19 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
|
|
||||||
"github.com/jackc/pgx/v5/pgtype"
|
|
||||||
)
|
|
||||||
|
|
||||||
type Querier interface {
|
|
||||||
GetLast(ctx context.Context, modelName string) (pgtype.Timestamptz, error)
|
|
||||||
Ping(ctx context.Context) (int32, error)
|
|
||||||
UpsertModel(ctx context.Context, arg *UpsertModelParams) error
|
|
||||||
}
|
|
||||||
|
|
||||||
var _ Querier = (*Queries)(nil)
|
|
@ -1,57 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
// source: queries.sql
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
|
|
||||||
"github.com/jackc/pgx/v5/pgtype"
|
|
||||||
)
|
|
||||||
|
|
||||||
const getLast = `-- name: GetLast :one
|
|
||||||
SELECT last_run
|
|
||||||
FROM
|
|
||||||
model_schedules
|
|
||||||
WHERE
|
|
||||||
model_name = $1
|
|
||||||
LIMIT 1
|
|
||||||
`
|
|
||||||
|
|
||||||
func (q *Queries) GetLast(ctx context.Context, modelName string) (pgtype.Timestamptz, error) {
|
|
||||||
row := q.db.QueryRow(ctx, getLast, modelName)
|
|
||||||
var last_run pgtype.Timestamptz
|
|
||||||
err := row.Scan(&last_run)
|
|
||||||
return last_run, err
|
|
||||||
}
|
|
||||||
|
|
||||||
const ping = `-- name: Ping :one
|
|
||||||
SELECT 1
|
|
||||||
`
|
|
||||||
|
|
||||||
func (q *Queries) Ping(ctx context.Context) (int32, error) {
|
|
||||||
row := q.db.QueryRow(ctx, ping)
|
|
||||||
var column_1 int32
|
|
||||||
err := row.Scan(&column_1)
|
|
||||||
return column_1, err
|
|
||||||
}
|
|
||||||
|
|
||||||
const upsertModel = `-- name: UpsertModel :exec
|
|
||||||
INSERT INTO model_schedules (model_name, last_run)
|
|
||||||
VALUES ($1, $2)
|
|
||||||
ON CONFLICT (model_name)
|
|
||||||
DO UPDATE SET
|
|
||||||
last_run = excluded.last_run
|
|
||||||
`
|
|
||||||
|
|
||||||
type UpsertModelParams struct {
|
|
||||||
ModelName string `json:"model_name"`
|
|
||||||
LastRun pgtype.Timestamptz `json:"last_run"`
|
|
||||||
}
|
|
||||||
|
|
||||||
func (q *Queries) UpsertModel(ctx context.Context, arg *UpsertModelParams) error {
|
|
||||||
_, err := q.db.Exec(ctx, upsertModel, arg.ModelName, arg.LastRun)
|
|
||||||
return err
|
|
||||||
}
|
|
@ -1,105 +0,0 @@
|
|||||||
package modelschedule
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"errors"
|
|
||||||
"fmt"
|
|
||||||
"log/slog"
|
|
||||||
"time"
|
|
||||||
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/modelregistry"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/modelschedule/repositories"
|
|
||||||
"github.com/adhocore/gronx"
|
|
||||||
"github.com/jackc/pgx/v5"
|
|
||||||
"github.com/jackc/pgx/v5/pgtype"
|
|
||||||
"github.com/jackc/pgx/v5/pgxpool"
|
|
||||||
)
|
|
||||||
|
|
||||||
//go:generate sqlc generate
|
|
||||||
|
|
||||||
type ModelRunSchedule struct {
|
|
||||||
Model *modelregistry.Model
|
|
||||||
Start time.Time
|
|
||||||
End time.Time
|
|
||||||
}
|
|
||||||
|
|
||||||
type ModelSchedule struct {
|
|
||||||
logger *slog.Logger
|
|
||||||
|
|
||||||
db *pgxpool.Pool
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewModelSchedule(logger *slog.Logger, db *pgxpool.Pool) *ModelSchedule {
|
|
||||||
return &ModelSchedule{
|
|
||||||
logger: logger,
|
|
||||||
|
|
||||||
db: db,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *ModelSchedule) GetNext(
|
|
||||||
ctx context.Context,
|
|
||||||
model modelregistry.Model,
|
|
||||||
start time.Time,
|
|
||||||
end time.Time,
|
|
||||||
amount uint,
|
|
||||||
) (models []ModelRunSchedule, lastExecuted *time.Time, err error) {
|
|
||||||
repo := repositories.New(m.db)
|
|
||||||
|
|
||||||
var startRun time.Time
|
|
||||||
lastRun, err := repo.GetLast(ctx, model.Name)
|
|
||||||
if err != nil {
|
|
||||||
if !errors.Is(err, pgx.ErrNoRows) {
|
|
||||||
return nil, nil, fmt.Errorf("failed to get last run for mode: %s: %w", model.Name, err)
|
|
||||||
}
|
|
||||||
|
|
||||||
startRun = start
|
|
||||||
} else {
|
|
||||||
startRun = lastRun.Time
|
|
||||||
}
|
|
||||||
|
|
||||||
times := make([]ModelRunSchedule, 0, amount)
|
|
||||||
for {
|
|
||||||
next, err := gronx.NextTickAfter(model.Schedule, startRun, false)
|
|
||||||
if err != nil {
|
|
||||||
return nil, nil, fmt.Errorf("failed to find next model schedule: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if next.Equal(time.Time{}) {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
|
|
||||||
if next.After(end) {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
|
|
||||||
times = append(times, ModelRunSchedule{
|
|
||||||
Model: &model,
|
|
||||||
Start: startRun,
|
|
||||||
End: next,
|
|
||||||
})
|
|
||||||
startRun = next
|
|
||||||
|
|
||||||
if len(times) >= int(amount) {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if len(times) == 0 {
|
|
||||||
return nil, nil, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
return times, &startRun, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (m *ModelSchedule) UpdateModelRun(ctx context.Context, model modelregistry.Model, lastRun *time.Time) error {
|
|
||||||
repo := repositories.New(m.db)
|
|
||||||
|
|
||||||
return repo.UpsertModel(ctx, &repositories.UpsertModelParams{
|
|
||||||
ModelName: model.Name,
|
|
||||||
LastRun: pgtype.Timestamptz{
|
|
||||||
Time: *lastRun,
|
|
||||||
Valid: true,
|
|
||||||
},
|
|
||||||
})
|
|
||||||
}
|
|
@ -1,21 +0,0 @@
|
|||||||
version: "2"
|
|
||||||
sql:
|
|
||||||
- queries: queries.sql
|
|
||||||
schema: ../persistence/migrations/
|
|
||||||
engine: "postgresql"
|
|
||||||
gen:
|
|
||||||
go:
|
|
||||||
out: "repositories"
|
|
||||||
package: "repositories"
|
|
||||||
sql_package: "pgx/v5"
|
|
||||||
emit_json_tags: true
|
|
||||||
emit_prepared_queries: true
|
|
||||||
emit_interface: true
|
|
||||||
emit_empty_slices: true
|
|
||||||
emit_result_struct_pointers: true
|
|
||||||
emit_params_struct_pointers: true
|
|
||||||
overrides:
|
|
||||||
- db_type: "uuid"
|
|
||||||
go_type:
|
|
||||||
import: "github.com/google/uuid"
|
|
||||||
type: "UUID"
|
|
@ -1,22 +1,4 @@
|
|||||||
CREATE TABLE worker_register (
|
CREATE TABLE worker_register (
|
||||||
worker_id UUID PRIMARY KEY NOT NULL
|
worker_id UUID PRIMARY KEY NOT NULL
|
||||||
, capacity INTEGER NOT NULL
|
|
||||||
, heart_beat TIMESTAMPTZ NOT NULL DEFAULT now()
|
, heart_beat TIMESTAMPTZ NOT NULL DEFAULT now()
|
||||||
);
|
);
|
||||||
|
|
||||||
CREATE TABLE model_schedules (
|
|
||||||
model_name TEXT PRIMARY KEY NOT NULL
|
|
||||||
, last_run TIMESTAMPTZ
|
|
||||||
);
|
|
||||||
|
|
||||||
CREATE TABLE work_schedule (
|
|
||||||
schedule_id UUID PRIMARY KEY NOT NULL
|
|
||||||
, worker_id UUID NOT NULL
|
|
||||||
, start_run TIMESTAMPTZ NOT NULL
|
|
||||||
, end_run TIMESTAMPTZ NOT NULL
|
|
||||||
, updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
|
||||||
, state TEXT NOT NULL
|
|
||||||
);
|
|
||||||
|
|
||||||
CREATE INDEX idx_work_schedule_worker ON work_schedule (worker_id);
|
|
||||||
CREATE INDEX idx_work_schedule_worker_updated ON work_schedule (worker_id, updated_at DESC);
|
|
||||||
|
@ -62,13 +62,12 @@ func (s *Scheduler) Execute(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
|
|
||||||
func (s *Scheduler) acquireLeader(ctx context.Context) (bool, error) {
|
func (s *Scheduler) acquireLeader(ctx context.Context) (bool, error) {
|
||||||
db, err := s.db.Acquire(ctx)
|
|
||||||
|
|
||||||
for {
|
for {
|
||||||
select {
|
select {
|
||||||
case <-ctx.Done():
|
case <-ctx.Done():
|
||||||
return false, nil
|
return false, nil
|
||||||
default:
|
default:
|
||||||
|
db, err := s.db.Acquire(ctx)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false, fmt.Errorf("failed to acquire db connection: %w", err)
|
return false, fmt.Errorf("failed to acquire db connection: %w", err)
|
||||||
}
|
}
|
||||||
|
@ -2,23 +2,14 @@
|
|||||||
SELECT 1;
|
SELECT 1;
|
||||||
|
|
||||||
-- name: RegisterWorker :exec
|
-- name: RegisterWorker :exec
|
||||||
INSERT INTO worker_register (worker_id, capacity)
|
INSERT INTO worker_register (worker_id)
|
||||||
VALUES (
|
VALUES (
|
||||||
$1
|
$1
|
||||||
, $2
|
|
||||||
);
|
);
|
||||||
|
|
||||||
-- name: GetWorkers :many
|
|
||||||
SELECT
|
|
||||||
worker_id
|
|
||||||
, capacity
|
|
||||||
FROM
|
|
||||||
worker_register;
|
|
||||||
|
|
||||||
-- name: UpdateWorkerHeartbeat :exec
|
-- name: UpdateWorkerHeartbeat :exec
|
||||||
UPDATE worker_register
|
UPDATE worker_register
|
||||||
SET
|
SET
|
||||||
heart_beat = now()
|
heart_beat = now()
|
||||||
WHERE
|
WHERE
|
||||||
worker_id = $1;
|
worker_id = $1;
|
||||||
|
|
||||||
|
@ -9,22 +9,7 @@ import (
|
|||||||
"github.com/jackc/pgx/v5/pgtype"
|
"github.com/jackc/pgx/v5/pgtype"
|
||||||
)
|
)
|
||||||
|
|
||||||
type ModelSchedule struct {
|
|
||||||
ModelName string `json:"model_name"`
|
|
||||||
LastRun pgtype.Timestamptz `json:"last_run"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type WorkSchedule struct {
|
|
||||||
ScheduleID uuid.UUID `json:"schedule_id"`
|
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
|
||||||
StartRun pgtype.Timestamptz `json:"start_run"`
|
|
||||||
EndRun pgtype.Timestamptz `json:"end_run"`
|
|
||||||
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
|
|
||||||
State string `json:"state"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type WorkerRegister struct {
|
type WorkerRegister struct {
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
WorkerID uuid.UUID `json:"worker_id"`
|
||||||
Capacity int32 `json:"capacity"`
|
|
||||||
HeartBeat pgtype.Timestamptz `json:"heart_beat"`
|
HeartBeat pgtype.Timestamptz `json:"heart_beat"`
|
||||||
}
|
}
|
||||||
|
@ -11,9 +11,8 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
type Querier interface {
|
type Querier interface {
|
||||||
GetWorkers(ctx context.Context) ([]*GetWorkersRow, error)
|
|
||||||
Ping(ctx context.Context) (int32, error)
|
Ping(ctx context.Context) (int32, error)
|
||||||
RegisterWorker(ctx context.Context, arg *RegisterWorkerParams) error
|
RegisterWorker(ctx context.Context, workerID uuid.UUID) error
|
||||||
UpdateWorkerHeartbeat(ctx context.Context, workerID uuid.UUID) error
|
UpdateWorkerHeartbeat(ctx context.Context, workerID uuid.UUID) error
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -11,39 +11,6 @@ import (
|
|||||||
"github.com/google/uuid"
|
"github.com/google/uuid"
|
||||||
)
|
)
|
||||||
|
|
||||||
const getWorkers = `-- name: GetWorkers :many
|
|
||||||
SELECT
|
|
||||||
worker_id
|
|
||||||
, capacity
|
|
||||||
FROM
|
|
||||||
worker_register
|
|
||||||
`
|
|
||||||
|
|
||||||
type GetWorkersRow struct {
|
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
|
||||||
Capacity int32 `json:"capacity"`
|
|
||||||
}
|
|
||||||
|
|
||||||
func (q *Queries) GetWorkers(ctx context.Context) ([]*GetWorkersRow, error) {
|
|
||||||
rows, err := q.db.Query(ctx, getWorkers)
|
|
||||||
if err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
defer rows.Close()
|
|
||||||
items := []*GetWorkersRow{}
|
|
||||||
for rows.Next() {
|
|
||||||
var i GetWorkersRow
|
|
||||||
if err := rows.Scan(&i.WorkerID, &i.Capacity); err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
items = append(items, &i)
|
|
||||||
}
|
|
||||||
if err := rows.Err(); err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
return items, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
const ping = `-- name: Ping :one
|
const ping = `-- name: Ping :one
|
||||||
SELECT 1
|
SELECT 1
|
||||||
`
|
`
|
||||||
@ -56,20 +23,14 @@ func (q *Queries) Ping(ctx context.Context) (int32, error) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
const registerWorker = `-- name: RegisterWorker :exec
|
const registerWorker = `-- name: RegisterWorker :exec
|
||||||
INSERT INTO worker_register (worker_id, capacity)
|
INSERT INTO worker_register (worker_id)
|
||||||
VALUES (
|
VALUES (
|
||||||
$1
|
$1
|
||||||
, $2
|
|
||||||
)
|
)
|
||||||
`
|
`
|
||||||
|
|
||||||
type RegisterWorkerParams struct {
|
func (q *Queries) RegisterWorker(ctx context.Context, workerID uuid.UUID) error {
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
_, err := q.db.Exec(ctx, registerWorker, workerID)
|
||||||
Capacity int32 `json:"capacity"`
|
|
||||||
}
|
|
||||||
|
|
||||||
func (q *Queries) RegisterWorker(ctx context.Context, arg *RegisterWorkerParams) error {
|
|
||||||
_, err := q.db.Exec(ctx, registerWorker, arg.WorkerID, arg.Capacity)
|
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -11,94 +11,36 @@ import (
|
|||||||
"github.com/jackc/pgx/v5/pgxpool"
|
"github.com/jackc/pgx/v5/pgxpool"
|
||||||
)
|
)
|
||||||
|
|
||||||
type workProcessor interface {
|
|
||||||
ProcessNext(ctx context.Context, worker_id uuid.UUID) error
|
|
||||||
}
|
|
||||||
|
|
||||||
//go:generate sqlc generate
|
//go:generate sqlc generate
|
||||||
|
|
||||||
type Worker struct {
|
type Worker struct {
|
||||||
workerID uuid.UUID
|
workerID uuid.UUID
|
||||||
|
|
||||||
db *pgxpool.Pool
|
db *pgxpool.Pool
|
||||||
workProcessor workProcessor
|
|
||||||
logger *slog.Logger
|
logger *slog.Logger
|
||||||
|
|
||||||
capacity uint
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewWorker(
|
func NewWorker(
|
||||||
db *pgxpool.Pool,
|
db *pgxpool.Pool,
|
||||||
logger *slog.Logger,
|
logger *slog.Logger,
|
||||||
workProcessor workProcessor,
|
|
||||||
) *Worker {
|
) *Worker {
|
||||||
return &Worker{
|
return &Worker{
|
||||||
workerID: uuid.New(),
|
workerID: uuid.New(),
|
||||||
db: db,
|
db: db,
|
||||||
workProcessor: workProcessor,
|
|
||||||
logger: logger,
|
logger: logger,
|
||||||
|
|
||||||
capacity: 50,
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *Worker) Setup(ctx context.Context) error {
|
func (w *Worker) Setup(ctx context.Context) error {
|
||||||
repo := repositories.New(w.db)
|
repo := repositories.New(w.db)
|
||||||
|
|
||||||
w.logger.InfoContext(ctx, "setting up worker", "worker_id", w.workerID)
|
if err := repo.RegisterWorker(ctx, w.workerID); err != nil {
|
||||||
if err := repo.RegisterWorker(
|
|
||||||
ctx,
|
|
||||||
&repositories.RegisterWorkerParams{
|
|
||||||
WorkerID: w.workerID,
|
|
||||||
Capacity: int32(w.capacity),
|
|
||||||
},
|
|
||||||
); err != nil {
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
type Workers struct {
|
|
||||||
Instances []WorkerInstance
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *Workers) Capacity() uint {
|
|
||||||
capacity := uint(0)
|
|
||||||
|
|
||||||
for _, worker := range w.Instances {
|
|
||||||
capacity += worker.Capacity
|
|
||||||
}
|
|
||||||
|
|
||||||
return capacity
|
|
||||||
}
|
|
||||||
|
|
||||||
type WorkerInstance struct {
|
|
||||||
WorkerID uuid.UUID
|
|
||||||
Capacity uint
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *Worker) GetWorkers(ctx context.Context) (*Workers, error) {
|
|
||||||
repo := repositories.New(w.db)
|
|
||||||
|
|
||||||
dbInstances, err := repo.GetWorkers(ctx)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to find workers: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
instances := make([]WorkerInstance, 0, len(dbInstances))
|
|
||||||
for _, dbInstance := range dbInstances {
|
|
||||||
instances = append(instances, WorkerInstance{
|
|
||||||
WorkerID: dbInstance.WorkerID,
|
|
||||||
Capacity: uint(dbInstance.Capacity),
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
return &Workers{
|
|
||||||
Instances: instances,
|
|
||||||
}, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *Worker) Start(ctx context.Context) error {
|
func (w *Worker) Start(ctx context.Context) error {
|
||||||
heartBeatCtx, heartBeatCancel := context.WithCancel(context.Background())
|
heartBeatCtx, heartBeatCancel := context.WithCancel(context.Background())
|
||||||
go func() {
|
go func() {
|
||||||
@ -127,10 +69,6 @@ func (w *Worker) Start(ctx context.Context) error {
|
|||||||
}()
|
}()
|
||||||
|
|
||||||
for {
|
for {
|
||||||
select {
|
|
||||||
case <-ctx.Done():
|
|
||||||
return nil
|
|
||||||
default:
|
|
||||||
if err := w.processWorkQueue(ctx); err != nil {
|
if err := w.processWorkQueue(ctx); err != nil {
|
||||||
// FIXME: dead letter item, right now we just log and continue
|
// FIXME: dead letter item, right now we just log and continue
|
||||||
|
|
||||||
@ -138,7 +76,6 @@ func (w *Worker) Start(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
|
||||||
|
|
||||||
func (w *Worker) updateHeartBeat(ctx context.Context) error {
|
func (w *Worker) updateHeartBeat(ctx context.Context) error {
|
||||||
repo := repositories.New(w.db)
|
repo := repositories.New(w.db)
|
||||||
@ -147,6 +84,8 @@ func (w *Worker) updateHeartBeat(ctx context.Context) error {
|
|||||||
return repo.UpdateWorkerHeartbeat(ctx, w.workerID)
|
return repo.UpdateWorkerHeartbeat(ctx, w.workerID)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (w *Worker) processWorkQueue(ctx context.Context) error {
|
func (w *Worker) processWorkQueue(_ context.Context) error {
|
||||||
return w.workProcessor.ProcessNext(ctx, w.workerID)
|
time.Sleep(time.Second)
|
||||||
|
|
||||||
|
return nil
|
||||||
}
|
}
|
||||||
|
@ -1,48 +0,0 @@
|
|||||||
package workprocessor
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"fmt"
|
|
||||||
"log/slog"
|
|
||||||
"time"
|
|
||||||
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/workscheduler"
|
|
||||||
"github.com/google/uuid"
|
|
||||||
)
|
|
||||||
|
|
||||||
type WorkProcessor struct {
|
|
||||||
workscheduler *workscheduler.WorkScheduler
|
|
||||||
logger *slog.Logger
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewWorkProcessor(workscheduler *workscheduler.WorkScheduler, logger *slog.Logger) *WorkProcessor {
|
|
||||||
return &WorkProcessor{
|
|
||||||
workscheduler: workscheduler,
|
|
||||||
logger: logger,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *WorkProcessor) ProcessNext(ctx context.Context, workerID uuid.UUID) error {
|
|
||||||
schedule, err := w.workscheduler.GetNext(ctx, workerID)
|
|
||||||
if err != nil {
|
|
||||||
return fmt.Errorf("failed to get next work item: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if schedule == nil {
|
|
||||||
// TODO: defer somewhere else
|
|
||||||
time.Sleep(time.Second)
|
|
||||||
return nil
|
|
||||||
}
|
|
||||||
|
|
||||||
w.logger.DebugContext(ctx, "handling item", "schedule", schedule)
|
|
||||||
|
|
||||||
if err := w.workscheduler.StartProcessing(ctx, *schedule); err != nil {
|
|
||||||
return fmt.Errorf("failed to start processing items: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if err := w.workscheduler.Archive(ctx, *schedule); err != nil {
|
|
||||||
return fmt.Errorf("failed to archive item: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
return nil
|
|
||||||
}
|
|
@ -1,56 +0,0 @@
|
|||||||
-- name: Ping :one
|
|
||||||
SELECT 1;
|
|
||||||
|
|
||||||
-- name: GetCurrentQueueSize :one
|
|
||||||
SELECT
|
|
||||||
COUNT(*) current_queue_size
|
|
||||||
FROM
|
|
||||||
work_schedule
|
|
||||||
WHERE
|
|
||||||
worker_id = $1
|
|
||||||
AND state <> 'archived';
|
|
||||||
|
|
||||||
-- name: InsertQueueItem :exec
|
|
||||||
INSERT INTO work_schedule
|
|
||||||
(
|
|
||||||
schedule_id
|
|
||||||
, worker_id
|
|
||||||
, start_run
|
|
||||||
, end_run
|
|
||||||
, state
|
|
||||||
)
|
|
||||||
VALUES
|
|
||||||
(
|
|
||||||
$1
|
|
||||||
, $2
|
|
||||||
, $3
|
|
||||||
, $4
|
|
||||||
, 'pending'
|
|
||||||
);
|
|
||||||
|
|
||||||
|
|
||||||
-- name: GetNext :one
|
|
||||||
SELECT
|
|
||||||
*
|
|
||||||
FROM
|
|
||||||
work_schedule
|
|
||||||
WHERE
|
|
||||||
worker_id = $1
|
|
||||||
AND state = 'pending'
|
|
||||||
ORDER BY updated_at DESC
|
|
||||||
LIMIT 1;
|
|
||||||
|
|
||||||
-- name: StartProcessing :exec
|
|
||||||
UPDATE work_schedule
|
|
||||||
SET
|
|
||||||
state = 'processing'
|
|
||||||
WHERE
|
|
||||||
schedule_id = $1;
|
|
||||||
|
|
||||||
-- name: Archive :exec
|
|
||||||
UPDATE work_schedule
|
|
||||||
SET
|
|
||||||
state = 'archived'
|
|
||||||
WHERE
|
|
||||||
schedule_id = $1;
|
|
||||||
|
|
@ -1,32 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
|
|
||||||
"github.com/jackc/pgx/v5"
|
|
||||||
"github.com/jackc/pgx/v5/pgconn"
|
|
||||||
)
|
|
||||||
|
|
||||||
type DBTX interface {
|
|
||||||
Exec(context.Context, string, ...interface{}) (pgconn.CommandTag, error)
|
|
||||||
Query(context.Context, string, ...interface{}) (pgx.Rows, error)
|
|
||||||
QueryRow(context.Context, string, ...interface{}) pgx.Row
|
|
||||||
}
|
|
||||||
|
|
||||||
func New(db DBTX) *Queries {
|
|
||||||
return &Queries{db: db}
|
|
||||||
}
|
|
||||||
|
|
||||||
type Queries struct {
|
|
||||||
db DBTX
|
|
||||||
}
|
|
||||||
|
|
||||||
func (q *Queries) WithTx(tx pgx.Tx) *Queries {
|
|
||||||
return &Queries{
|
|
||||||
db: tx,
|
|
||||||
}
|
|
||||||
}
|
|
@ -1,30 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"github.com/google/uuid"
|
|
||||||
"github.com/jackc/pgx/v5/pgtype"
|
|
||||||
)
|
|
||||||
|
|
||||||
type ModelSchedule struct {
|
|
||||||
ModelName string `json:"model_name"`
|
|
||||||
LastRun pgtype.Timestamptz `json:"last_run"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type WorkSchedule struct {
|
|
||||||
ScheduleID uuid.UUID `json:"schedule_id"`
|
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
|
||||||
StartRun pgtype.Timestamptz `json:"start_run"`
|
|
||||||
EndRun pgtype.Timestamptz `json:"end_run"`
|
|
||||||
UpdatedAt pgtype.Timestamptz `json:"updated_at"`
|
|
||||||
State string `json:"state"`
|
|
||||||
}
|
|
||||||
|
|
||||||
type WorkerRegister struct {
|
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
|
||||||
Capacity int32 `json:"capacity"`
|
|
||||||
HeartBeat pgtype.Timestamptz `json:"heart_beat"`
|
|
||||||
}
|
|
@ -1,22 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
|
|
||||||
"github.com/google/uuid"
|
|
||||||
)
|
|
||||||
|
|
||||||
type Querier interface {
|
|
||||||
Archive(ctx context.Context, scheduleID uuid.UUID) error
|
|
||||||
GetCurrentQueueSize(ctx context.Context, workerID uuid.UUID) (int64, error)
|
|
||||||
GetNext(ctx context.Context, workerID uuid.UUID) (*WorkSchedule, error)
|
|
||||||
InsertQueueItem(ctx context.Context, arg *InsertQueueItemParams) error
|
|
||||||
Ping(ctx context.Context) (int32, error)
|
|
||||||
StartProcessing(ctx context.Context, scheduleID uuid.UUID) error
|
|
||||||
}
|
|
||||||
|
|
||||||
var _ Querier = (*Queries)(nil)
|
|
@ -1,129 +0,0 @@
|
|||||||
// Code generated by sqlc. DO NOT EDIT.
|
|
||||||
// versions:
|
|
||||||
// sqlc v1.23.0
|
|
||||||
// source: queries.sql
|
|
||||||
|
|
||||||
package repositories
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
|
|
||||||
"github.com/google/uuid"
|
|
||||||
"github.com/jackc/pgx/v5/pgtype"
|
|
||||||
)
|
|
||||||
|
|
||||||
const archive = `-- name: Archive :exec
|
|
||||||
UPDATE work_schedule
|
|
||||||
SET
|
|
||||||
state = 'archived'
|
|
||||||
WHERE
|
|
||||||
schedule_id = $1
|
|
||||||
`
|
|
||||||
|
|
||||||
func (q *Queries) Archive(ctx context.Context, scheduleID uuid.UUID) error {
|
|
||||||
_, err := q.db.Exec(ctx, archive, scheduleID)
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
const getCurrentQueueSize = `-- name: GetCurrentQueueSize :one
|
|
||||||
SELECT
|
|
||||||
COUNT(*) current_queue_size
|
|
||||||
FROM
|
|
||||||
work_schedule
|
|
||||||
WHERE
|
|
||||||
worker_id = $1
|
|
||||||
AND state <> 'archived'
|
|
||||||
`
|
|
||||||
|
|
||||||
func (q *Queries) GetCurrentQueueSize(ctx context.Context, workerID uuid.UUID) (int64, error) {
|
|
||||||
row := q.db.QueryRow(ctx, getCurrentQueueSize, workerID)
|
|
||||||
var current_queue_size int64
|
|
||||||
err := row.Scan(¤t_queue_size)
|
|
||||||
return current_queue_size, err
|
|
||||||
}
|
|
||||||
|
|
||||||
const getNext = `-- name: GetNext :one
|
|
||||||
SELECT
|
|
||||||
schedule_id, worker_id, start_run, end_run, updated_at, state
|
|
||||||
FROM
|
|
||||||
work_schedule
|
|
||||||
WHERE
|
|
||||||
worker_id = $1
|
|
||||||
AND state = 'pending'
|
|
||||||
ORDER BY updated_at DESC
|
|
||||||
LIMIT 1
|
|
||||||
`
|
|
||||||
|
|
||||||
func (q *Queries) GetNext(ctx context.Context, workerID uuid.UUID) (*WorkSchedule, error) {
|
|
||||||
row := q.db.QueryRow(ctx, getNext, workerID)
|
|
||||||
var i WorkSchedule
|
|
||||||
err := row.Scan(
|
|
||||||
&i.ScheduleID,
|
|
||||||
&i.WorkerID,
|
|
||||||
&i.StartRun,
|
|
||||||
&i.EndRun,
|
|
||||||
&i.UpdatedAt,
|
|
||||||
&i.State,
|
|
||||||
)
|
|
||||||
return &i, err
|
|
||||||
}
|
|
||||||
|
|
||||||
const insertQueueItem = `-- name: InsertQueueItem :exec
|
|
||||||
INSERT INTO work_schedule
|
|
||||||
(
|
|
||||||
schedule_id
|
|
||||||
, worker_id
|
|
||||||
, start_run
|
|
||||||
, end_run
|
|
||||||
, state
|
|
||||||
)
|
|
||||||
VALUES
|
|
||||||
(
|
|
||||||
$1
|
|
||||||
, $2
|
|
||||||
, $3
|
|
||||||
, $4
|
|
||||||
, 'pending'
|
|
||||||
)
|
|
||||||
`
|
|
||||||
|
|
||||||
type InsertQueueItemParams struct {
|
|
||||||
ScheduleID uuid.UUID `json:"schedule_id"`
|
|
||||||
WorkerID uuid.UUID `json:"worker_id"`
|
|
||||||
StartRun pgtype.Timestamptz `json:"start_run"`
|
|
||||||
EndRun pgtype.Timestamptz `json:"end_run"`
|
|
||||||
}
|
|
||||||
|
|
||||||
func (q *Queries) InsertQueueItem(ctx context.Context, arg *InsertQueueItemParams) error {
|
|
||||||
_, err := q.db.Exec(ctx, insertQueueItem,
|
|
||||||
arg.ScheduleID,
|
|
||||||
arg.WorkerID,
|
|
||||||
arg.StartRun,
|
|
||||||
arg.EndRun,
|
|
||||||
)
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
const ping = `-- name: Ping :one
|
|
||||||
SELECT 1
|
|
||||||
`
|
|
||||||
|
|
||||||
func (q *Queries) Ping(ctx context.Context) (int32, error) {
|
|
||||||
row := q.db.QueryRow(ctx, ping)
|
|
||||||
var column_1 int32
|
|
||||||
err := row.Scan(&column_1)
|
|
||||||
return column_1, err
|
|
||||||
}
|
|
||||||
|
|
||||||
const startProcessing = `-- name: StartProcessing :exec
|
|
||||||
UPDATE work_schedule
|
|
||||||
SET
|
|
||||||
state = 'processing'
|
|
||||||
WHERE
|
|
||||||
schedule_id = $1
|
|
||||||
`
|
|
||||||
|
|
||||||
func (q *Queries) StartProcessing(ctx context.Context, scheduleID uuid.UUID) error {
|
|
||||||
_, err := q.db.Exec(ctx, startProcessing, scheduleID)
|
|
||||||
return err
|
|
||||||
}
|
|
@ -1,21 +0,0 @@
|
|||||||
version: "2"
|
|
||||||
sql:
|
|
||||||
- queries: queries.sql
|
|
||||||
schema: ../persistence/migrations/
|
|
||||||
engine: "postgresql"
|
|
||||||
gen:
|
|
||||||
go:
|
|
||||||
out: "repositories"
|
|
||||||
package: "repositories"
|
|
||||||
sql_package: "pgx/v5"
|
|
||||||
emit_json_tags: true
|
|
||||||
emit_prepared_queries: true
|
|
||||||
emit_interface: true
|
|
||||||
emit_empty_slices: true
|
|
||||||
emit_result_struct_pointers: true
|
|
||||||
emit_params_struct_pointers: true
|
|
||||||
overrides:
|
|
||||||
- db_type: "uuid"
|
|
||||||
go_type:
|
|
||||||
import: "github.com/google/uuid"
|
|
||||||
type: "UUID"
|
|
@ -1,179 +0,0 @@
|
|||||||
package workscheduler
|
|
||||||
|
|
||||||
import (
|
|
||||||
"context"
|
|
||||||
"errors"
|
|
||||||
"fmt"
|
|
||||||
"log/slog"
|
|
||||||
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/modelschedule"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/worker"
|
|
||||||
"git.front.kjuulh.io/kjuulh/orbis/internal/workscheduler/repositories"
|
|
||||||
"github.com/google/uuid"
|
|
||||||
"github.com/jackc/pgx/v5"
|
|
||||||
"github.com/jackc/pgx/v5/pgtype"
|
|
||||||
"github.com/jackc/pgx/v5/pgxpool"
|
|
||||||
)
|
|
||||||
|
|
||||||
//go:generate sqlc generate
|
|
||||||
|
|
||||||
type WorkScheduler struct {
|
|
||||||
db *pgxpool.Pool
|
|
||||||
logger *slog.Logger
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewWorkScheduler(
|
|
||||||
db *pgxpool.Pool,
|
|
||||||
logger *slog.Logger,
|
|
||||||
) *WorkScheduler {
|
|
||||||
return &WorkScheduler{
|
|
||||||
db: db,
|
|
||||||
logger: logger,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
type Worker struct {
|
|
||||||
Instance worker.WorkerInstance
|
|
||||||
RemainingCapacity uint
|
|
||||||
}
|
|
||||||
|
|
||||||
type Workers struct {
|
|
||||||
Workers []*Worker
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *Workers) IterateSlice(size uint) func(yield func([]Worker) bool) {
|
|
||||||
return func(yield func([]Worker) bool) {
|
|
||||||
if len(w.Workers) == 0 {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
|
|
||||||
workers := make([]Worker, 0)
|
|
||||||
acc := uint(0)
|
|
||||||
|
|
||||||
for {
|
|
||||||
exit := true
|
|
||||||
|
|
||||||
for _, worker := range w.Workers {
|
|
||||||
if acc == size {
|
|
||||||
if !yield(workers) {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
workers = make([]Worker, 0)
|
|
||||||
acc = uint(0)
|
|
||||||
|
|
||||||
}
|
|
||||||
|
|
||||||
if worker.RemainingCapacity <= 0 {
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
worker.RemainingCapacity--
|
|
||||||
workers = append(workers, *worker)
|
|
||||||
acc++
|
|
||||||
|
|
||||||
exit = false
|
|
||||||
}
|
|
||||||
|
|
||||||
if exit {
|
|
||||||
if len(workers) > 0 {
|
|
||||||
if !yield(workers) {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
return
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *WorkScheduler) GetWorkers(ctx context.Context, registeredWorkers *worker.Workers) (*Workers, error) {
|
|
||||||
|
|
||||||
w.logger.DebugContext(ctx, "found workers", "workers", len(registeredWorkers.Instances))
|
|
||||||
|
|
||||||
workers := make([]*Worker, 0, len(registeredWorkers.Instances))
|
|
||||||
for _, registeredWorker := range registeredWorkers.Instances {
|
|
||||||
remainingCapacity, err := w.GetWorker(ctx, ®isteredWorker)
|
|
||||||
if err != nil {
|
|
||||||
return nil, fmt.Errorf("failed to find capacity for worker: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if remainingCapacity == 0 {
|
|
||||||
w.logger.DebugContext(ctx, "skipping worker as no remaining capacity")
|
|
||||||
continue
|
|
||||||
}
|
|
||||||
|
|
||||||
workers = append(workers, &Worker{
|
|
||||||
Instance: registeredWorker,
|
|
||||||
RemainingCapacity: remainingCapacity,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
return &Workers{Workers: workers}, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *WorkScheduler) GetWorker(
|
|
||||||
ctx context.Context,
|
|
||||||
worker *worker.WorkerInstance,
|
|
||||||
) (uint, error) {
|
|
||||||
repo := repositories.New(w.db)
|
|
||||||
|
|
||||||
current_size, err := repo.GetCurrentQueueSize(ctx, worker.WorkerID)
|
|
||||||
if err != nil {
|
|
||||||
return 0, fmt.Errorf("failed to get current queue size: %s: %w", worker.WorkerID, err)
|
|
||||||
}
|
|
||||||
|
|
||||||
if int64(worker.Capacity)-current_size <= 0 {
|
|
||||||
return 0, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
return worker.Capacity - uint(current_size), nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *WorkScheduler) InsertModelRun(
|
|
||||||
ctx context.Context,
|
|
||||||
worker Worker,
|
|
||||||
modelRun *modelschedule.ModelRunSchedule,
|
|
||||||
) error {
|
|
||||||
repo := repositories.New(w.db)
|
|
||||||
|
|
||||||
return repo.InsertQueueItem(ctx, &repositories.InsertQueueItemParams{
|
|
||||||
ScheduleID: uuid.New(),
|
|
||||||
WorkerID: worker.Instance.WorkerID,
|
|
||||||
StartRun: pgtype.Timestamptz{
|
|
||||||
Time: modelRun.Start,
|
|
||||||
Valid: true,
|
|
||||||
},
|
|
||||||
EndRun: pgtype.Timestamptz{
|
|
||||||
Time: modelRun.End,
|
|
||||||
Valid: true,
|
|
||||||
},
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *WorkScheduler) GetNext(ctx context.Context, workerID uuid.UUID) (*uuid.UUID, error) {
|
|
||||||
repo := repositories.New(w.db)
|
|
||||||
|
|
||||||
schedule, err := repo.GetNext(ctx, workerID)
|
|
||||||
if err != nil {
|
|
||||||
if !errors.Is(err, pgx.ErrNoRows) {
|
|
||||||
return nil, fmt.Errorf("failed to get next worker item: %w", err)
|
|
||||||
}
|
|
||||||
|
|
||||||
return nil, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
return &schedule.ScheduleID, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *WorkScheduler) StartProcessing(ctx context.Context, scheduleID uuid.UUID) error {
|
|
||||||
repo := repositories.New(w.db)
|
|
||||||
|
|
||||||
return repo.StartProcessing(ctx, scheduleID)
|
|
||||||
}
|
|
||||||
|
|
||||||
func (w *WorkScheduler) Archive(ctx context.Context, scheduleID uuid.UUID) error {
|
|
||||||
repo := repositories.New(w.db)
|
|
||||||
|
|
||||||
return repo.Archive(ctx, scheduleID)
|
|
||||||
}
|
|
Loading…
x
Reference in New Issue
Block a user