From 82eb111c7f97a1cd0a40df814df209940444090d Mon Sep 17 00:00:00 2001
From: Tomas Valenta <valenta.and.thomas@gmail.com>
Date: Fri, 29 Sep 2023 18:09:47 +0200
Subject: [PATCH] Change build exit check

---
 packages/env-build-task-driver/internal/handle.go |  8 ++++++++
 packages/env-build-task-driver/internal/task.go   | 10 ++++++----
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/packages/env-build-task-driver/internal/handle.go b/packages/env-build-task-driver/internal/handle.go
index c046f371d..5788a0715 100644
--- a/packages/env-build-task-driver/internal/handle.go
+++ b/packages/env-build-task-driver/internal/handle.go
@@ -27,6 +27,7 @@ type taskHandle struct {
 
 	exited chan struct{}
 
+	ctx    context.Context
 	cancel context.CancelFunc
 
 	mu sync.RWMutex
@@ -47,6 +48,13 @@ func (h *taskHandle) TaskStatus() *drivers.TaskStatus {
 	}
 }
 
+func (h *taskHandle) IsRunning() bool {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	return h.procState == drivers.TaskStateRunning
+}
+
 func (h *taskHandle) run(ctx context.Context, tracer trace.Tracer, docker *client.Client) {
 	childCtx, childSpan := tracer.Start(ctx, "run")
 	defer childSpan.End()
diff --git a/packages/env-build-task-driver/internal/task.go b/packages/env-build-task-driver/internal/task.go
index 438fc728e..7ddb196d4 100644
--- a/packages/env-build-task-driver/internal/task.go
+++ b/packages/env-build-task-driver/internal/task.go
@@ -103,6 +103,7 @@ func (d *Driver) StartTask(cfg *drivers.TaskConfig) (*drivers.TaskHandle, *drive
 		env:        &env,
 		exited:     make(chan struct{}),
 		cancel:     cancel,
+		ctx:        cancellableBuildContext,
 	}
 
 	driverState := TaskState{
@@ -156,16 +157,13 @@ func (d *Driver) WaitTask(ctx context.Context, taskID string) (<-chan *drivers.E
 func (d *Driver) handleWait(ctx context.Context, handle *taskHandle, ch chan *drivers.ExitResult) {
 	defer close(ch)
 
-	ticker := time.NewTicker(2 * time.Second)
-	defer ticker.Stop()
-
 	for {
 		select {
 		case <-ctx.Done():
 			return
 		case <-d.ctx.Done():
 			return
-		case <-ticker.C:
+		case <-handle.ctx.Done():
 			s := handle.TaskStatus()
 			if s.State == drivers.TaskStateExited {
 				ch <- handle.exitResult
@@ -191,6 +189,10 @@ func (d *Driver) DestroyTask(taskID string, force bool) error {
 		return drivers.ErrTaskNotFound
 	}
 
+	if handle.IsRunning() && !force {
+		return errors.New("task is still running")
+	}
+
 	handle.cancel()
 	d.tasks.Delete(taskID)