handle some edge cases

Author: zacharyblasczyk

cmd/ctrlc/root/run/exec/exec.go

@@ -1,11 +1,13 @@
 package exec
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/signal"
 	"runtime"
 	"syscall"
+	"time"
 
 	"github.com/MakeNowJust/heredoc/v2"
 	"github.com/charmbracelet/log"
@@ -17,8 +19,9 @@ import (
 
 func NewRunExecCmd() *cobra.Command {
 	var (
-		name         string
-		jobAgentType = "exec-bash"
+		name             string
+		jobAgentType     = "exec-bash"
+		maxJobTimeoutMin int
 	)
 
 	if runtime.GOOS == "windows" {
@@ -46,6 +49,15 @@ func NewRunExecCmd() *cobra.Command {
 				return fmt.Errorf("workspace is required")
 			}
 
+			maxJobTimeoutMin := viper.GetInt("max-job-timeout")
+			if maxJobTimeoutMin <= 0 {
+				maxJobTimeoutMin = 30
+			}
+
+			// Create a context with timeout for job execution
+			ctx, cancel := context.WithTimeout(context.Background(), time.Duration(maxJobTimeoutMin)*time.Minute)
+			defer cancel()
+
 			runner := NewExecRunner(client)
 			jobAgentConfig := api.UpsertJobAgentJSONRequestBody{
 				Name:        name,
@@ -69,10 +81,11 @@ func NewRunExecCmd() *cobra.Command {
 				<-c
 				log.Info("Shutting down gracefully...")
 				runner.ExitAll(true)
+				cancel() // Cancel the context when shutting down
 			}()
 
-			// Run job check - AddIntervalSupport will handle repeated execution
-			if err := ja.RunQueuedJobs(); err != nil {
+			// Run job check with timeout context - AddIntervalSupport will handle repeated execution
+			if err := ja.RunQueuedJobs(ctx); err != nil {
 				return fmt.Errorf("failed to run queued jobs: %w", err)
 			}
 
@@ -81,6 +94,7 @@ func NewRunExecCmd() *cobra.Command {
 	}
 
 	cmd.Flags().StringVar(&name, "name", "", "Name of the job agent")
+	cmd.Flags().IntVar(&maxJobTimeoutMin, "max-job-timeout", 0, "Maximum job execution time in minutes (default 30)")
 	cmd.MarkFlagRequired("name")
 	cmd.MarkFlagRequired("workspace")
 	return cmd

cmd/ctrlc/root/run/exec/runner.go

@@ -52,6 +52,20 @@ func NewExecRunner(client *api.ClientWithResponses) *ExecRunner {
 		sig := <-c
 		log.Info("Received shutdown signal, terminating all jobs", "signal", sig)
 
+		// Set up a short timeout for handling additional signals
+		timeout := make(chan bool, 1)
+		go func() {
+			time.Sleep(5 * time.Second)
+			timeout <- true
+		}()
+
+		// Handle additional signals
+		go func() {
+			sig := <-c
+			log.Warn("Received second signal, forcing immediate exit", "signal", sig)
+			os.Exit(1)
+		}()
+
 		// Update all job statuses before exiting
 		runner.mu.Lock()
 		for _, runningJob := range runner.runningJobs {
@@ -67,14 +81,26 @@ func NewExecRunner(client *api.ClientWithResponses) *ExecRunner {
 		// Now terminate all processes
 		runner.ExitAll(true)
 		log.Info("Shutdown complete, exiting")
-		os.Exit(0)
 	}()
 
 	return runner
 }
 
+// cleanupJob safely removes a job from the runningJobs map with proper mutex handling
+func (r *ExecRunner) cleanupJob(handle string) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if job, exists := r.runningJobs[handle]; exists {
+		log.Debug("Job cleanup complete", "id", job.jobID)
+		delete(r.runningJobs, handle)
+	}
+}
+
 // Start creates a temporary script file, starts the process, and updates job status when the process completes.
 func (r *ExecRunner) Start(ctx context.Context, job *api.JobWithDetails) (api.JobStatus, error) {
+	// Ensure job has the current context
+	job = job.WithContext(ctx)
+
 	// Template the script using the job
 	script, err := job.TemplateJobDetails()
 	if err != nil {
@@ -148,12 +174,7 @@ func (r *ExecRunner) Start(ctx context.Context, job *api.JobWithDetails) (api.Jo
 	// Spawn a goroutine to wait for the process to finish and update the job status
 	go func(handle, scriptPath string) {
 		defer os.Remove(scriptPath)
-		defer func() {
-			r.mu.Lock()
-			delete(r.runningJobs, handle)
-			r.mu.Unlock()
-			log.Debug("Job cleanup complete", "id", runningJob.jobID)
-		}()
+		defer r.cleanupJob(handle)
 
 		log.Debug("Waiting for command to complete", "id", runningJob.jobID, "handle", handle)
 		err := cmd.Wait()

internal/api/job.go

@@ -14,18 +14,25 @@ type JobWithDetails struct {
 	Details    map[string]interface{}
 	Client     *ClientWithResponses
 	ExternalID string
+	ctx        context.Context
 }
 
 // NewJobWithDetails creates a new JobWithDetails with the client and job data
 func NewJobWithDetails(client *ClientWithResponses, job Job) (*JobWithDetails, error) {
+	return NewJobWithDetailsContext(context.Background(), client, job)
+}
+
+// NewJobWithDetailsContext creates a new JobWithDetails with a specific context
+func NewJobWithDetailsContext(ctx context.Context, client *ClientWithResponses, job Job) (*JobWithDetails, error) {
 	j := &JobWithDetails{
 		Job:    job,
 		Client: client,
+		ctx:    ctx,
 	}
 
 	// Fetch job details
 	var err error
-	j.Details, err = j.GetJobDetails(context.Background())
+	j.Details, err = j.GetJobDetails()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get job details: %w", err)
 	}
@@ -34,8 +41,8 @@ func NewJobWithDetails(client *ClientWithResponses, job Job) (*JobWithDetails, e
 }
 
 // GetJobDetails retrieves job details for templating
-func (j *JobWithDetails) GetJobDetails(ctx context.Context) (map[string]interface{}, error) {
-	resp, err := j.Client.GetJobWithResponse(ctx, j.Id.String())
+func (j *JobWithDetails) GetJobDetails() (map[string]interface{}, error) {
+	resp, err := j.Client.GetJobWithResponse(j.ctx, j.Id.String())
 	if err != nil {
 		return nil, fmt.Errorf("failed to get job details: %w", err)
 	}
@@ -54,6 +61,14 @@ func (j *JobWithDetails) GetJobDetails(ctx context.Context) (map[string]interfac
 	return details, nil
 }
 
+// WithContext returns a copy of the job with a new context
+func (j *JobWithDetails) WithContext(ctx context.Context) *JobWithDetails {
+	// Create a shallow copy with the new context
+	jobCopy := *j
+	jobCopy.ctx = ctx
+	return &jobCopy
+}
+
 // TemplateJobDetails applies the job details template to the script
 func (j *JobWithDetails) TemplateJobDetails() (string, error) {
 	// Extract script from JobAgentConfig
@@ -89,7 +104,7 @@ func (j *JobWithDetails) UpdateStatus(status JobStatus, message string) error {
 		body.ExternalId = &j.ExternalID
 	}
 
-	resp, err := j.Client.UpdateJobWithResponse(context.Background(), j.Id.String(), body)
+	resp, err := j.Client.UpdateJobWithResponse(j.ctx, j.Id.String(), body)
 	if err != nil {
 		return fmt.Errorf("failed to update job status: %w", err)
 	}

pkg/jobagent/runner.go

@@ -47,8 +47,8 @@ type JobAgent struct {
 // RunQueuedJobs retrieves and executes any queued jobs for this agent.
 // For each job, it starts execution using the runner's Start method, which
 // will update the job status when the job completes.
-func (a *JobAgent) RunQueuedJobs() error {
-	jobs, err := a.client.GetNextJobsWithResponse(context.Background(), a.id)
+func (a *JobAgent) RunQueuedJobs(ctx context.Context) error {
+	jobs, err := a.client.GetNextJobsWithResponse(ctx, a.id)
 	if err != nil {
 		return err
 	}
@@ -65,7 +65,7 @@ func (a *JobAgent) RunQueuedJobs() error {
 		// 	continue
 		// }
 
-		job, err := api.NewJobWithDetails(a.client, apiJob)
+		job, err := api.NewJobWithDetailsContext(ctx, a.client, apiJob)
 		if err != nil {
 			log.Error("Failed to create job with details", "error", err, "jobId", apiJob.Id.String())
 			continue
@@ -76,7 +76,7 @@ func (a *JobAgent) RunQueuedJobs() error {
 			defer wg.Done()
 
 			// Start the job - status updates happen inside Start
-			status, err := a.runner.Start(context.Background(), job)
+			status, err := a.runner.Start(ctx, job)
 
 			if err != nil {
 				log.Error("Failed to start job", "error", err, "jobId", job.Id.String())