plane/monitor/cli/pkg/worker/worker.go

package worker

import (
	"context"
	"fmt"
	"sync"

	log "github.com/makeplane/plane-ee/monitor/lib/logger"
	"github.com/makeplane/plane-ee/monitor/pkg/constants/descriptors"
	error_msgs "github.com/makeplane/plane-ee/monitor/pkg/constants/errors"
)

type Job func(context.Context)

type JobDefination struct {
	Id  string
	Job Job
}

// Prime Worker is responsible to run jobs in the background and keeping in
// track of each of them
type PrimeWorker struct {
	sync.Mutex
	ctx        context.Context
	cancel     context.CancelFunc
	wg         *sync.WaitGroup
	jobs       []JobDefination
	Logger     *log.Handler
	ExitedJobs []string
	running    bool
}

func NewPrimeWorker(logger *log.Handler) *PrimeWorker {
	ctx, cancel := context.WithCancel(context.Background())

	return &PrimeWorker{
		wg:      &sync.WaitGroup{},
		ctx:     ctx,
		cancel:  cancel,
		jobs:    make([]JobDefination, 0),
		Logger:  logger,
		running: false,
	}
}

// ---------------- Controller Methods -----------------------
// Adds a new job the job queue with an identifier
func (w *PrimeWorker) RegisterJob(id string, job Job) {
	w.wg.Add(1)
	w.jobs = append(w.jobs, JobDefination{
		Id:  id,
		Job: job,
	})
}

// Returns the currently available jobs
func (w *PrimeWorker) GetJobCount() int {
	return len(w.jobs)
}

func (w *PrimeWorker) GetCompletedJobs() []string {
	return w.ExitedJobs
}

// Starts the jobs in background and updates the completed and errored jobs
// queue for the updates
func (w *PrimeWorker) StartJobsInBackground() {
	if w.running {
		w.Logger.Error(context.Background(), fmt.Sprintf(error_msgs.F_WORKER_ALREADY_RUNNING, w.GetJobCount(), len(w.GetCompletedJobs())))
		return
	}
	w.running = true
	for _, job := range w.jobs { // Running the job in background
		go func(w *PrimeWorker, job JobDefination) {
			// Modify the job queue and append the completed job to the completion
			// array when the job is done
			defer func(w *PrimeWorker, job JobDefination) {
				w.DequeueJob(job)
			}(w, job)

			// Start the job in the background with a boolean channel for messages
			done := make(chan bool)
			go func() {
				defer close(done)
				w.Logger.Info(w.ctx, fmt.Sprintf(descriptors.F_START_JOB, job.Id))
				job.Job(w.ctx)
				done <- true
			}()

			// Listen for the cancellation signal for cancellation of the job
			select {
			case <-w.ctx.Done():
				w.Logger.Warning(w.ctx, fmt.Sprintf(descriptors.F_SHUTTING_DOWN_JOB, job.Id))
				return
			case <-done:
				w.Logger.Info(w.ctx, fmt.Sprintf(descriptors.F_EXITED_JOB, job.Id))
				return
			}
		}(w, job)
	}
}

func (w *PrimeWorker) Shutdown() {
	w.Logger.Warning(w.ctx, descriptors.SHUTTING_DOWN)
	w.cancel()
	w.wg.Wait()
}

func (w *PrimeWorker) Wait() {
	w.wg.Wait()
}

func (w *PrimeWorker) DequeueJob(job JobDefination) {
	jobIndex := 0
	// Find the position of the job in the job queue
	for index, j := range w.jobs {
		if j.Id == job.Id {
			jobIndex = index
		}
	}

	w.Lock()
	w.wg.Done()
	// Pop off the job from the job queue
	if jobIndex == len(w.jobs)-1 {
		w.jobs = w.jobs[:jobIndex]
	} else {
		w.jobs[jobIndex] = w.jobs[len(w.jobs)-1]
		w.jobs = w.jobs[:len(w.jobs)-1]
	}

	// Append the finished job to the exited job slice
	w.ExitedJobs = append(w.ExitedJobs, job.Id)

	if w.GetJobCount() == 0 {
		w.running = false
	}
	w.Unlock()
}