瀏覽代碼

Merge branch 'main' into next

jamesread 5 月之前
父節點
當前提交
f06e56857e
共有 2 個文件被更改,包括 70 次插入11 次删除
  1. 4 11
      service/internal/executor/executor.go
  2. 66 0
      service/internal/executor/timeout_context.go

+ 4 - 11
service/internal/executor/executor.go

@@ -827,7 +827,7 @@ func buildEnv(args map[string]string) []string {
 }
 
 func stepExec(req *ExecutionRequest) bool {
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(req.Binding.Action.Timeout)*time.Second)
+	ctx, cancel := newTimeoutContext(context.Background(), time.Duration(req.Binding.Action.Timeout)*time.Second, req.executor)
 	defer cancel()
 	streamer := &OutputStreamer{Req: req}
 	cmd := buildCommand(ctx, req)
@@ -839,6 +839,7 @@ func stepExec(req *ExecutionRequest) bool {
 	prepareCommand(cmd, streamer, req)
 	runerr := cmd.Start()
 	req.logEntry.Process = cmd.Process
+	ctx.setProcess(cmd.Process)
 	waiterr := cmd.Wait()
 	req.logEntry.ExitCode = int32(cmd.ProcessState.ExitCode())
 	req.logEntry.Output = streamer.String()
@@ -851,15 +852,6 @@ func stepExec(req *ExecutionRequest) bool {
 			"actionTitle": req.logEntry.ActionTitle,
 		}).Warnf("Action timed out")
 
-		// The context timeout should kill the process, but let's make sure.
-		err := req.executor.Kill(req.logEntry)
-
-		if err != nil {
-			log.WithFields(log.Fields{
-				"actionTitle": req.logEntry.ActionTitle,
-			}).Warnf("could not kill process: %v", err)
-		}
-
 		req.logEntry.TimedOut = true
 		req.logEntry.Output += "OliveTin::timeout - this action timed out after " + fmt.Sprintf("%v", req.Binding.Action.Timeout) + " seconds. If you need more time for this action, set a longer timeout. See https://docs.olivetin.app/action_customization/timeouts.html for more help."
 	}
@@ -888,7 +880,7 @@ func stepExecAfter(req *ExecutionRequest) bool {
 		return true
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(req.Binding.Action.Timeout)*time.Second)
+	ctx, cancel := newTimeoutContext(context.Background(), time.Duration(req.Binding.Action.Timeout)*time.Second, req.executor)
 	defer cancel()
 
 	var stdout bytes.Buffer
@@ -917,6 +909,7 @@ func stepExecAfter(req *ExecutionRequest) bool {
 	cmd.Env = buildEnv(args)
 
 	runerr := cmd.Start()
+	ctx.setProcess(cmd.Process)
 
 	waiterr := cmd.Wait()
 

+ 66 - 0
service/internal/executor/timeout_context.go

@@ -0,0 +1,66 @@
+package executor
+
+import (
+	"context"
+	"os"
+	"sync"
+	"time"
+
+	log "github.com/sirupsen/logrus"
+)
+
+// timeoutContext is a custom context that kills the process group when cancelled due to timeout.
+type timeoutContext struct {
+	context.Context
+	cancel    context.CancelFunc
+	process   *os.Process
+	executor  *Executor
+	processMu sync.Mutex
+}
+
+// newTimeoutContext creates a context that will kill the process group when the timeout expires.
+func newTimeoutContext(parent context.Context, timeout time.Duration, executor *Executor) (*timeoutContext, context.CancelFunc) {
+	ctx, cancel := context.WithTimeout(parent, timeout)
+	tc := &timeoutContext{
+		Context:  ctx,
+		cancel:   cancel,
+		executor: executor,
+	}
+
+	// Start a goroutine that kills the process group when the context is cancelled
+	go func() {
+		<-ctx.Done()
+		if ctx.Err() == context.DeadlineExceeded {
+			tc.processMu.Lock()
+			process := tc.process
+			tc.processMu.Unlock()
+
+			if process != nil {
+				logEntry := &InternalLogEntry{Process: process}
+				if err := executor.Kill(logEntry); err != nil {
+					log.WithFields(log.Fields{
+						"error": err,
+					}).Warnf("Failed to kill process group on timeout")
+				}
+			}
+		}
+	}()
+
+	return tc, cancel
+}
+
+func (tc *timeoutContext) setProcess(process *os.Process) {
+	tc.processMu.Lock()
+	tc.process = process
+	tc.processMu.Unlock()
+
+	// If deadline already expired before process was set, kill now
+	if tc.Context.Err() == context.DeadlineExceeded && process != nil {
+		logEntry := &InternalLogEntry{Process: process}
+		if err := tc.executor.Kill(logEntry); err != nil {
+			log.WithFields(log.Fields{
+				"error": err,
+			}).Warnf("Failed to kill process group on timeout (late registration)")
+		}
+	}
+}