148 lines
3.3 KiB
Go
148 lines
3.3 KiB
Go
package task
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/play-with-docker/play-with-docker/pwd/types"
|
|
"github.com/play-with-docker/play-with-docker/storage"
|
|
)
|
|
|
|
type Task interface {
|
|
Name() string
|
|
Run(ctx context.Context, instance *types.Instance) error
|
|
}
|
|
|
|
type SchedulerApi interface {
|
|
Schedule(session *types.Session) error
|
|
Unschedule(session *types.Session) error
|
|
Start()
|
|
Stop()
|
|
AddTask(task Task) error
|
|
RemoveTask(task Task) error
|
|
}
|
|
|
|
type scheduledSession struct {
|
|
session *types.Session
|
|
cancel context.CancelFunc
|
|
ctx context.Context
|
|
ticker *time.Ticker
|
|
}
|
|
|
|
type scheduler struct {
|
|
scheduledSessions map[string]*scheduledSession
|
|
storage storage.StorageApi
|
|
tasks map[string]Task
|
|
started bool
|
|
}
|
|
|
|
func NewScheduler(s storage.StorageApi) (*scheduler, error) {
|
|
sch := &scheduler{storage: s}
|
|
|
|
sch.tasks = make(map[string]Task)
|
|
|
|
err := sch.loadFromStorage()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return sch, nil
|
|
}
|
|
|
|
func (s *scheduler) loadFromStorage() error {
|
|
sessions, err := s.storage.SessionGetAll()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
s.scheduledSessions = make(map[string]*scheduledSession, len(sessions))
|
|
for _, session := range sessions {
|
|
s.register(session)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *scheduler) AddTask(task Task) error {
|
|
if _, found := s.tasks[task.Name()]; found {
|
|
return fmt.Errorf("Task [%s] was already added", task.Name())
|
|
}
|
|
s.tasks[task.Name()] = task
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *scheduler) RemoveTask(task Task) error {
|
|
if _, found := s.tasks[task.Name()]; !found {
|
|
return fmt.Errorf("Task [%s] doesn't exist", task.Name())
|
|
}
|
|
delete(s.tasks, task.Name())
|
|
|
|
return nil
|
|
}
|
|
|
|
func (s *scheduler) Start() {
|
|
for _, session := range s.scheduledSessions {
|
|
go s.cron(session)
|
|
}
|
|
s.started = true
|
|
}
|
|
|
|
func (s *scheduler) register(session *types.Session) *scheduledSession {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
s.scheduledSessions[session.Id] = &scheduledSession{session: session, cancel: cancel, ctx: ctx}
|
|
return s.scheduledSessions[session.Id]
|
|
}
|
|
|
|
func (s *scheduler) cron(session *scheduledSession) {
|
|
session.ticker = time.NewTicker(1 * time.Second)
|
|
|
|
for {
|
|
select {
|
|
case <-session.ticker.C:
|
|
s.processSession(session.ctx, session.session)
|
|
case <-session.ctx.Done():
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *scheduler) processSession(ctx context.Context, session *types.Session) {
|
|
for _, ins := range session.Instances {
|
|
go s.processInstance(ctx, ins)
|
|
}
|
|
}
|
|
|
|
func (s *scheduler) processInstance(ctx context.Context, instance *types.Instance) {
|
|
for _, task := range s.tasks {
|
|
task.Run(ctx, instance)
|
|
}
|
|
}
|
|
|
|
func (s *scheduler) Schedule(session *types.Session) error {
|
|
if !s.started {
|
|
return fmt.Errorf("Can only schedule sessions after the scheduler has been started.")
|
|
}
|
|
if _, found := s.scheduledSessions[session.Id]; found {
|
|
return fmt.Errorf("Session [%s] was already scheduled", session.Id)
|
|
}
|
|
go s.cron(s.register(session))
|
|
return nil
|
|
}
|
|
|
|
func (s *scheduler) Unschedule(session *types.Session) error {
|
|
if !s.started {
|
|
return fmt.Errorf("Can only schedule sessions after the scheduler has been started.")
|
|
}
|
|
if _, found := s.scheduledSessions[session.Id]; !found {
|
|
return fmt.Errorf("Session [%s] in not scheduled", session.Id)
|
|
}
|
|
|
|
scheduledSession := s.scheduledSessions[session.Id]
|
|
scheduledSession.cancel()
|
|
scheduledSession.ticker.Stop()
|
|
delete(s.scheduledSessions, session.Id)
|
|
|
|
return nil
|
|
}
|