feat(wip): cancel tree runners

pkg/runner/tree_runner.go

@@ -1,7 +1,7 @@
 package runner
 
 import (
-	"errors"
+	"fmt"
 	"sync"
 	"sync/atomic"
 
@@ -31,7 +31,7 @@ type TaskTreeRunner struct {
 	// cancel is a channel that is used to cancel the execution of the task tree.
 	cancel chan bool
 	// cancelComplete is a channel that is used to signal that the cancel operation has completed.
-	cancelComplete chan error
+	cancelComplete chan bool
 	// hasError is a flag that indicates whether an error occurred during the execution of the task tree.
 	hasError atomic.Bool
 
@@ -54,7 +54,7 @@ func NewTaskTreeRunner(root *tasks.TaskTreeNode, p ConcurrencyProvider) *TaskTre
 		updates:        make(chan *TreeStatusNode, 1000),
 		wasCanceled:    atomic.Bool{},
 		cancel:         make(chan bool),
-		cancelComplete: make(chan error),
+		cancelComplete: make(chan bool),
 
 		mutex: sync.RWMutex{},
 	}
@@ -68,10 +68,10 @@ func (r *TaskTreeRunner) Status() *TreeStatusNode {
 	return r.statusTree
 }
 
-func (r *TaskTreeRunner) Cancel() error {
+func (r *TaskTreeRunner) Cancel() {
 	r.cancel <- true
 	close(r.cancel)
-	return <-r.cancelComplete
+	<-r.cancelComplete
 }
 
 func (r *TaskTreeRunner) updateTaskStatus(node *tasks.TaskTreeNode, status TaskStatus) {
@@ -93,20 +93,14 @@ func (r *TaskTreeRunner) setError(node *tasks.TaskTreeNode, err error) {
 
 func (r *TaskTreeRunner) Start() error {
 	done := make(chan bool, 1)
-	errs := []error{}
+	errs := map[string]error{}
 	errMu := sync.Mutex{}
 	wg := sync.WaitGroup{}
 
 	// cleanup
 	defer func() {
 		if r.wasCanceled.Load() {
-			errMu.Lock()
-			if len(errs) > 0 {
-				r.cancelComplete <- errors.Join(errs...)
-			} else {
-				r.cancelComplete <- nil
-			}
-			errMu.Unlock()
+			r.cancelComplete <- true
 		}
 		close(r.updates)
 		close(r.cancelComplete)
@@ -129,15 +123,13 @@ func (r *TaskTreeRunner) Start() error {
 			r.mutex.Unlock()
 			go func(task *tasks.TaskTreeNode, cancel <-chan bool) {
 				// acquire a ticket to run the task
-				// fmt.Println("wait", task.NodeID())
 				ticket := r.tickets.Acquire()
-				// fmt.Println("start", task.NodeID())
 				r.updateTaskStatus(task, StatusRunning)
 				if err := task.Main.Run(cancel); err != nil {
-					// fmt.Println("error", task.NodeID(), err)
 					errMu.Lock()
-					errs = append(errs, err)
-					if !r.hasError.Load() {
+					errs[task.NodeID()] = err
+					fmt.Println(err)
+					if !r.hasError.Load() && !r.wasCanceled.Load() {
 						// first node erroring -> close the channel
 						close(r.scheduledNodes)
 					}
@@ -156,7 +148,6 @@ func (r *TaskTreeRunner) Start() error {
 				r.mutex.Lock()
 				delete(r.forwardCancel, task.NodeID())
 				r.mutex.Unlock()
-				// fmt.Println("release ticket", task.NodeID())
 				// release the ticket to be used by another channel
 				r.tickets.Release(ticket)
 				wg.Done()
@@ -172,8 +163,8 @@ func (r *TaskTreeRunner) Start() error {
 		case <-r.cancel:
 			// run was canceled - forward cancel to all tasks
 			r.wasCanceled.Store(true)
+			close(r.scheduledNodes)
 			for _, cancel := range r.forwardCancel {
-				// fmt.Println("cancel", id)
 				cancel <- true
 			}
 			return
@@ -195,8 +186,12 @@ func (r *TaskTreeRunner) Start() error {
 	done <- true
 	close(done)
 
+	if r.wasCanceled.Load() {
+		return tasks.ErrCancelled
+	}
+
 	if len(errs) > 0 {
-		return errors.Join(errs...)
+		return tasks.NewMultiTaskError(errs)
 	}
 	return nil
 }
@@ -206,6 +201,8 @@ func (r *TaskTreeRunner) scheduleNext(node *tasks.TaskTreeNode) {
 		return // fail fast
 	}
 
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
 	statusNode := findStatus(r.statusTree, node)
 	if allDone(r.statusTree) {
 		close(r.scheduledNodes)

pkg/tasks/command_task.go

@@ -87,7 +87,7 @@ func (ct commandTask) Run(cancel <-chan bool) error {
 		} else {
 			err = exec.Command("pkill", "-P", strconv.Itoa(ct.cmd.Process.Pid)).Run()
 		}
-		fmt.Println("killed", ct.cmd.Process.Pid, err)
+		// err = syscall.Kill(ct.cmd.Process.Pid, syscall.SIGKILL)
 		if err != nil {
 			// fall back to builtin kill
 			if err := ct.cmd.Process.Kill(); err != nil {

pkg/tasks/error.go

@@ -0,0 +1,29 @@
+package tasks
+
+import (
+	"errors"
+	"fmt"
+	"strings"
+)
+
+var (
+	ErrCancelled = errors.New("task cancelled")
+)
+
+type MultiTaskError struct {
+	Errors map[string]error
+}
+
+func (mte MultiTaskError) Error() string {
+	nodeIds := make([]string, 0, len(mte.Errors))
+	for id := range mte.Errors {
+		nodeIds = append(nodeIds, id)
+	}
+	return fmt.Sprintf("tasks %s failed", strings.Join(nodeIds, ", "))
+}
+
+func NewMultiTaskError(errors map[string]error) MultiTaskError {
+	return MultiTaskError{
+		Errors: errors,
+	}
+}

pkg/ui/quite_runner.go

@@ -62,10 +62,7 @@ func (m *quiteView) setupInterruptHandler() {
 	go func() {
 		for range c {
 			for _, runner := range m.runners {
-				err := runner.Cancel()
-				m.mu.Lock()
-				m.errs = append(m.errs, err)
-				m.mu.Unlock()
+				runner.Cancel()
 			}
 			break
 		}

pkg/ui/static_tree.go

@@ -5,7 +5,6 @@ import (
 	"os"
 	"os/signal"
 	"sync"
-	"time"
 
 	"github.com/zwoo-hq/zwooc/pkg/tasks"
 )
@@ -27,10 +26,10 @@ func newStaticTreeRunner(forest tasks.Collection, provider SimpleStatusProvider,
 
 	fmt.Printf("%s - %s\n", zwoocBranding, forest.GetName())
 	model.setupInterruptHandler()
-	execStart := time.Now()
-	hasError := false
+	// execStart := time.Now()
+	// hasError := false
 
-	start := time.Now()
+	// start := time.Now()
 	outputs := map[string]*tasks.CommandCapturer{}
 
 	// setup task pipes
@@ -58,7 +57,7 @@ func newStaticTreeRunner(forest tasks.Collection, provider SimpleStatusProvider,
 
 	// wait until everything is completed
 	model.wg.Wait()
-	execEnd := time.Now()
+	// execEnd := time.Now()
 
 	// TODO: add option in provider to find all errors
 	// if model.err != nil {
@@ -110,6 +109,7 @@ func (m *staticTreeView) ReceiveUpdates(c <-chan StatusUpdate, prefix string) {
 			fmt.Printf("%s %s %s\n", prefix, node.NodeID, canceledStyle.Render("was canceled"))
 		}
 	}
+	fmt.Println("updates done")
 	m.wg.Done()
 }
 
@@ -120,6 +120,7 @@ func (m *staticTreeView) WaitForDone() {
 		m.err = err
 		m.mu.Unlock()
 	}
+	m.wg.Done()
 }
 
 func (m *staticTreeView) printFinalStatus() {

pkg/ui/tree_progress.go

@@ -1,6 +1,7 @@
 package ui
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"os/signal"
@@ -49,6 +50,9 @@ func NewTreeProgressView(forest tasks.Collection, status SimpleStatusProvider, o
 	}
 
 	// TODO: done -display cancel or error or success
+	if errors.Is(model.err, tasks.ErrCancelled) {
+		fmt.Println("cancelled")
+	}
 
 	fmt.Println("done!!")
 	return nil

pkg/ui/ui.go

@@ -10,16 +10,14 @@ func NewRunner(forest tasks.Collection, provider SimpleStatusProvider, options V
 	}
 
 	if options.DisableTUI {
-		// TODO: use provided runner
-		newStaticTreeRunner(forest, options)
+		newStaticTreeRunner(forest, provider, options)
 		return
 	}
 
 	// try interactive view
 	if err := NewTreeProgressView(forest, provider, options); err != nil {
 		// fall back to static view
-		// TODO: use provided runner
-		newStaticTreeRunner(forest, options)
+		newStaticTreeRunner(forest, provider, options)
 	}
 }
 
@@ -61,18 +59,22 @@ func (g SimpleStatusProvider) Start() {
 	close(g.start)
 }
 
-func (g SimpleStatusProvider) Cancel() {
+func (g *SimpleStatusProvider) Cancel() {
 	if !g.wasCanceled {
+		g.wasCanceled = true
 		g.cancel <- struct{}{}
 		close(g.cancel)
-		g.wasCanceled = true
 	}
 }
 
 func (g SimpleStatusProvider) UpdateStatus(update StatusUpdate) {
 	g.status <- update
 }
 
+func (g SimpleStatusProvider) CloseUpdates() {
+	close(g.status)
+}
+
 func (g SimpleStatusProvider) Done(err error) {
 	g.done <- err
 	close(g.done)

pkg/zwooc/ui_adapter.go

@@ -1,6 +1,8 @@
 package zwooc
 
 import (
+	"sync"
+
 	"github.com/zwoo-hq/zwooc/pkg/config"
 	"github.com/zwoo-hq/zwooc/pkg/runner"
 	"github.com/zwoo-hq/zwooc/pkg/tasks"
@@ -49,15 +51,23 @@ func createForestRunner(forest tasks.Collection, maxConcurrency int) ui.SimpleSt
 	})
 
 	// forward updates
+	updatesWg := sync.WaitGroup{}
 	for _, r := range runners {
 		currentRunner := r
+		updatesWg.Add(1)
 		go func() {
 			for update := range currentRunner.Updates() {
 				statusProvider.UpdateStatus(runnerToStatusProvider(update))
 			}
+			updatesWg.Done()
 		}()
 	}
 
+	go func() {
+		updatesWg.Wait()
+		statusProvider.CloseUpdates()
+	}()
+
 	return statusProvider
 }