manager.go

package taskmanager

import (
	"errors"
	"github.com/go-co-op/gocron/v2"
	"github.com/google/uuid"
	"sync"
)

// Manager task manager.
type Manager struct {
	scheduler gocron.Scheduler
	tasks     map[string]*Task
	lock      sync.RWMutex
}

// NewManager returns a new instance of the Manager structure.
func NewManager() *Manager {
	return &Manager{
		tasks: make(map[string]*Task),
	}
}

// Start starts task manager.
func (m *Manager) Start(opts ...gocron.SchedulerOption) (err error) {
	m.lock.Lock()

	if m.scheduler != nil {
		m.lock.Unlock()
		return errors.New("already running")
	}

	// Exit early.
	if len(m.tasks) == 0 {
		m.lock.Unlock()
		return nil
	}

	m.scheduler, err = initializeScheduler(opts...)

	if err != nil {
		m.lock.Unlock()
		return err
	}

	for _, task := range m.tasks {
		// Check that the task is correct.
		err = task.Validate()

		if err != nil {
			m.lock.Unlock()
			return err
		}

		// Skip tasks that don't need to be added to the planner right away.
		if !task.withJob {
			continue
		}

		jobID, err := updateSchedulerJob(task.name, task.definition, uuid.Nil, task.taskFunc, m.scheduler)

		if err != nil {
			m.lock.Unlock()
			return err
		}

		task.jobID = jobID
	}

	// Start the scheduler.
	m.scheduler.Start()
	m.lock.Unlock()
	return
}

// Stop stops the scheduler and deletes all tasks.
func (m *Manager) Stop() error {
	m.lock.Lock()

	// Already stopped.
	if m.scheduler == nil {
		m.lock.Unlock()
		return nil
	}

	err := m.scheduler.Shutdown()

	if err != nil {
		m.lock.Unlock()
		return err
	}

	if m.tasks != nil {
		for _, task := range m.tasks {
			task.jobID = uuid.Nil
		}
	}

	m.scheduler = nil
	m.lock.Unlock()
	return nil
}

// GetTask returns task by name.
func (m *Manager) GetTask(name string) (*Task, bool) {
	m.lock.RLock()
	if m.tasks == nil {
		m.lock.RUnlock()
		return nil, false
	}
	value, ok := m.tasks[name]
	m.lock.RUnlock()
	return value, ok
}

// GetTasks returns a list of tasks.
func (m *Manager) GetTasks() map[string]*Task {
	m.lock.RLock()
	values := make(map[string]*Task)
	for _, task := range m.tasks {
		values[task.name] = task
	}
	m.lock.RUnlock()
	return values
}

// AddTask adds a new task to the manager.
func (m *Manager) AddTask(name, definition string, withJob bool, function interface{}, parameters ...interface{}) *Manager {
	m.lock.Lock()

	if m.scheduler != nil {
		m.lock.Unlock()
		return m
	}

	if m.tasks == nil {
		m.tasks = make(map[string]*Task)
	}

	m.tasks[name] = &Task{
		name:       name,
		definition: definition,
		withJob:    withJob,
		jobID:      uuid.Nil,
		taskFunc:   gocron.NewTask(function, parameters...),
	}

	m.lock.Unlock()
	return m
}

// StartTask starts the job in the scheduler.
func (m *Manager) StartTask(name string) error {
	m.lock.Lock()

	if m.tasks == nil {
		m.lock.Unlock()
		return ErrTaskNotFound
	}

	task, ok := m.tasks[name]

	if !ok {
		m.lock.Unlock()
		return ErrTaskNotFound
	}

	if task.jobID != uuid.Nil {
		m.lock.Unlock()
		return nil
	}

	if m.scheduler == nil {
		m.lock.Unlock()
		return ErrNilScheduler
	}

	jobID, err := updateSchedulerJob(task.name, task.definition, task.jobID, task.taskFunc, m.scheduler)

	if err != nil {
		m.lock.Unlock()
		return err
	}

	task.jobID = jobID
	m.lock.Unlock()
	return nil
}

// StopTask deletes the job from the scheduler.
func (m *Manager) StopTask(name string) error {
	m.lock.Lock()

	if m.tasks == nil {
		m.lock.Unlock()
		return ErrTaskNotFound
	}

	task, ok := m.tasks[name]

	if !ok {
		m.lock.Unlock()
		return ErrTaskNotFound
	}

	if task.jobID != uuid.Nil {
		if m.scheduler == nil {
			m.lock.Unlock()
			return ErrNilScheduler
		}

		err := m.scheduler.RemoveJob(task.jobID)

		if err != nil {
			m.lock.Unlock()
			return err
		}
	}

	task.jobID = uuid.Nil
	m.lock.Unlock()
	return nil
}

// UpdateTaskDefinition updates the task definition.
func (m *Manager) UpdateTaskDefinition(name, definition string) error {
	m.lock.Lock()

	if m.tasks == nil {
		m.lock.Unlock()
		return ErrTaskNotFound
	}

	task, ok := m.tasks[name]

	if !ok {
		m.lock.Unlock()
		return ErrTaskNotFound
	}

	if task.jobID != uuid.Nil {
		if m.scheduler == nil {
			m.lock.Unlock()
			return ErrNilScheduler
		}

		jobID, err := updateSchedulerJob(task.name, task.definition, task.jobID, task.taskFunc, m.scheduler)

		if err != nil {
			m.lock.Unlock()
			return err
		}

		task.jobID = jobID
	}

	task.definition = definition
	m.lock.Unlock()
	return nil
}