filed

Job queue using FUSE

git clone git://mccd.space/filed

manager.go (6008B)

 package main
 
 import (
 	"bytes"
 	"context"
 	"git.sr.ht/~marcc/filed/store"
 	"strings"
 
 	"fmt"
 	"log/slog"
 	"math"
 	"os/exec"
 	"sync"
 	"time"
 )
 
 type RestrictionsArg struct {
 	rwFilesArg []string
 	roFilesArg []string
 	rwDirArg   []string
 	roDirArg   []string
 }
 
 type JobManager struct {
 	store      *store.Store
 	activeJobs sync.Map
 	// Command to use for executing jobs in pending
 	execCmd      string
 	restrictions *RestrictionsArg
 
 	// Used to notify that a job should be run with JobManager.notify
 	notifyChan chan struct{}
 }
 
 type ActiveJob struct {
 	mu     sync.Mutex
 	cancel context.CancelFunc
 	output bytes.Buffer
 }
 
 func NewJobManager(s *store.Store, filedLaunchExecutablePath string, restrictions *Restrictions) *JobManager {
 
 	argRestrictions := RestrictionsArg{}
 	if restrictions != nil {
 		argRestrictions.roFilesArg = toArg(restrictions.roFiles, "-rof")
 		argRestrictions.roDirArg = toArg(restrictions.roDir, "-ro")
 		argRestrictions.rwFilesArg = toArg(restrictions.rwFiles, "-rwf")
 		argRestrictions.rwDirArg = toArg(restrictions.rwDir, "-rw")
 	}
 	slog.Info("restrictions", "rest", argRestrictions)
 
 	return &JobManager{
 		store:        s,
 		execCmd:      filedLaunchExecutablePath,
 		restrictions: &argRestrictions,
 
 		// Buffer size 1 = Ignore more than one signal so we don't process jobs too many times
 		notifyChan: make(chan struct{}, 1),
 	}
 }
 
 type signal struct{}
 
 // Notifies to process jobs
 func (jm *JobManager) notify() {
 	select {
 	case jm.notifyChan <- signal{}:
 	default:
 		// Already signaled to process jobs, ignore.
 	}
 }
 
 func toArg(lst []string, argParam string) []string {
 	var result []string
 
 	for _, v := range lst {
 		result = append(result, argParam, v)
 	}
 
 	return result
 }
 
 func (jm *JobManager) StartWorker(ctx context.Context) {
 	jobs, err := jm.store.ListJobsByState(store.StateRunning)
 	if err != nil {
 		panic(err)
 	}
 
 	for _, job := range jobs {
 		slog.Warn("Found dangling job, moving back to pending", "job", job.ID)
 		err := jm.store.RestartJob(job.ID, job.Output)
 		if err != nil {
 			panic(err)
 		}
 	}
 
 	go func() {
 		jm.notify()
 		for {
 			nextRetryIn := jm.processPendingJobs()
 			var timerChan <-chan time.Time
 			if nextRetryIn > 0 {
 				timerChan = time.After(nextRetryIn)
 			}
 
 			select {
 			case <-ctx.Done():
 				return
 			case <-jm.notifyChan:
 				// Wake up because a job was added or finished
 			case <-timerChan:
 				// Wake up because a backoff period ended
 			}
 
 		}
 	}()
 }
 
 func (jm *JobManager) processPendingJobs() time.Duration {
 	var minWait time.Duration = 0
 	conf := jm.store.GetConfig()
 	activeJobs, err := jm.store.ListJobsByState(store.StateRunning)
 	if err != nil {
 		slog.Error("Worker: Failed to list jobs", "error", err)
 		return minWait
 	}
 	activeJobCount := len(activeJobs)
 
 	jobs, err := jm.store.ListJobsByState(store.StatePending)
 	if err != nil {
 		slog.Error("Worker: Failed to list jobs", "error", err)
 		return minWait
 	}
 
 	for _, job := range jobs {
 		if activeJobCount >= conf.MaxJobCount {
 			break
 		}
 
 		if job.Attempts >= conf.MaxAttempts {
 			slog.Warn("Worker: job exceeded max attempts. Moving to failed", "Job", job.ID)
 			jm.store.FailJob(job.ID)
 			continue
 		}
 		ready, waitTime := jm.isReadyWithWait(&conf, &job)
 		if ready {
 			slog.Info("Worker: Starting job", "id", job.ID, "cmd", job.Command, "attempt", job.Attempts+1)
 			if err := jm.store.AttemptJob(job.ID); err != nil {
 				slog.Error("Worker: Failed to claim job", "id", job.ID, "error", err)
 				continue
 			}
 
 			go jm.runJob(job.ID, job.Command)
 			activeJobCount++
 		} else {
 			if minWait == 0 || waitTime < minWait {
 				minWait = waitTime
 			}
 		}
 	}
 	return minWait
 }
 
 func (jm *JobManager) isReadyWithWait(config *store.Config, job *store.Job) (bool, time.Duration) {
 	currentTime := time.Now().Unix()
 	lastModified := job.UpdatedAt.Unix()
 	waitTimeSec := int64(float64(config.BackoffBase) * math.Pow(float64(config.BackoffMult), float64(job.Attempts)))
 	readyAt := lastModified + waitTimeSec
 
 	remaining := readyAt - currentTime
 	if remaining <= 0 {
 		return true, 0
 	}
 	return false, time.Duration(remaining) * time.Second
 }
 
 func (jm *JobManager) runJob(id, commandStr string) {
 	defer jm.notify()
 	timeout_seconds := jm.store.GetConfig().TimeoutSec
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout_seconds)*time.Second)
 	defer cancel()
 
 	active := &ActiveJob{cancel: cancel}
 	jm.activeJobs.Store(id, active)
 	defer jm.activeJobs.Delete(id)
 
 	args := strings.Fields(commandStr)
 
 	cmd := exec.CommandContext(ctx, jm.execCmd)
 	if jm.restrictions != nil {
 		cmd.Args = append(cmd.Args, jm.restrictions.roFilesArg...)
 		cmd.Args = append(cmd.Args, jm.restrictions.rwFilesArg...)
 		cmd.Args = append(cmd.Args, jm.restrictions.rwDirArg...)
 		cmd.Args = append(cmd.Args, jm.restrictions.roDirArg...)
 	}
 	cmd.Args = append(cmd.Args, "--")
 	cmd.Args = append(cmd.Args, args...)
 
 	writer := &SafeBuffer{target: active}
 	cmd.Stdout = writer
 	cmd.Stderr = writer
 
 	err := cmd.Run()
 
 	jobOutput := active.output.Bytes()
 
 	exitCode := cmd.ProcessState.ExitCode()
 
 	if err != nil || exitCode != 0 {
 		slog.Warn("Worker: Job failed", "id", id, "error", err, "exitCode", exitCode)
 		errMsg := fmt.Sprintf("\n\n[System Error]: %v\n", err)
 		jobOutput = append(jobOutput, []byte(errMsg)...)
 		// This can return an error if the job is cancelled by removing it.
 		// Probably we should handle it more gracefully
 		jm.store.RestartJob(id, jobOutput)
 	} else {
 		slog.Info("Worker: Job completed", "id", id, "exitCode", exitCode)
 		jm.store.CompleteJob(id, jobOutput)
 	}
 }
 
 type SafeBuffer struct {
 	target *ActiveJob
 }
 
 func (sb *SafeBuffer) Write(p []byte) (n int, err error) {
 	sb.target.mu.Lock()
 	defer sb.target.mu.Unlock()
 	return sb.target.output.Write(p)
 }
 
 func (jm *JobManager) GetLiveLog(id string) []byte {
 	val, ok := jm.activeJobs.Load(id)
 	if !ok {
 		return nil
 	}
 	job := val.(*ActiveJob)
 	job.mu.Lock()
 	defer job.mu.Unlock()
 	b := make([]byte, job.output.Len())
 	copy(b, job.output.Bytes())
 	return b
 }