summaryrefslogtreecommitdiff
path: root/internal/executor/executor.go
diff options
context:
space:
mode:
Diffstat (limited to 'internal/executor/executor.go')
-rw-r--r--internal/executor/executor.go27
1 files changed, 25 insertions, 2 deletions
diff --git a/internal/executor/executor.go b/internal/executor/executor.go
index f6932f4..62fed2e 100644
--- a/internal/executor/executor.go
+++ b/internal/executor/executor.go
@@ -35,6 +35,7 @@ type Pool struct {
mu sync.Mutex
active int
+ cancels map[string]context.CancelFunc // taskID → cancel
resultCh chan *Result
Questions *QuestionRegistry
}
@@ -55,6 +56,7 @@ func NewPool(maxConcurrent int, runner Runner, store *storage.DB, logger *slog.L
runner: runner,
store: store,
logger: logger,
+ cancels: make(map[string]context.CancelFunc),
resultCh: make(chan *Result, maxConcurrent*2),
Questions: NewQuestionRegistry(),
}
@@ -81,6 +83,19 @@ func (p *Pool) Results() <-chan *Result {
return p.resultCh
}
+// Cancel requests cancellation of a running task. Returns false if the task
+// is not currently running in this pool.
+func (p *Pool) Cancel(taskID string) bool {
+ p.mu.Lock()
+ cancel, ok := p.cancels[taskID]
+ p.mu.Unlock()
+ if !ok {
+ return false
+ }
+ cancel()
+ return true
+}
+
// SubmitResume re-queues a blocked task using the provided resume execution.
// The execution must have ResumeSessionID and ResumeAnswer set.
func (p *Pool) SubmitResume(ctx context.Context, t *task.Task, exec *storage.Execution) error {
@@ -230,14 +245,22 @@ func (p *Pool) execute(ctx context.Context, t *task.Task) {
p.logger.Error("failed to update task state", "error", err)
}
- // Apply task timeout.
+ // Apply task timeout and register cancel so callers can stop this task.
var cancel context.CancelFunc
if t.Timeout.Duration > 0 {
ctx, cancel = context.WithTimeout(ctx, t.Timeout.Duration)
} else {
ctx, cancel = context.WithCancel(ctx)
}
- defer cancel()
+ p.mu.Lock()
+ p.cancels[t.ID] = cancel
+ p.mu.Unlock()
+ defer func() {
+ cancel()
+ p.mu.Lock()
+ delete(p.cancels, t.ID)
+ p.mu.Unlock()
+ }()
// Run the task.
err := p.runner.Run(ctx, t, exec)