Repository: bacalhau-project/bacalhau (language: python)

Issue:
Incorrect state when cancelling job
When calling `bacalhau cancel` on a job, the logs show the following info, but two errors to do with the state of the job.
It's almost like the job is being cancelled twice with the second entry complaining ...

```
19:13:12.103 | INF pkg/publicapi/middleware/logger.go:50 > 9128751d-7748-40a8-9c0c-c9c7ad178081] [Method:POST] [NodeID:node-0] [Referer:] [RemoteAddr:127.0.0.1] [RequestID:] [Size:702] [StatusCode:200] [URI:/api/v1/requester/cancel] [UserAgent:Go-http-client/1.1]
19:13:12.109 | INF pkg/executor/docker/handler.go:172 > killing the container [container:dc451521008e2437d6d54e8058d3dc5909d9031c6d717f188c573487535c73f8] [execution:e-50dd14d0-c873-43b7-a49a-bda5f86889ce] [job:9128751d-7748-40a8-9c0c-c9c7ad178081]
...
19:13:13.233 | INF pkg/executor/docker/handler.go:189 > removing container [container:dc451521008e2437d6d54e8058d3dc5909d9031c6d717f188c573487535c73f8] [execution:e-50dd14d0-c873-43b7-a49a-bda5f86889ce] [job:9128751d-7748-40a8-9c0c-c9c7ad178081]
19:13:13.242 | INF pkg/executor/docker/executor.go:217 > received results from execution [executionID:e-50dd14d0-c873-43b7-a49a-bda5f86889ce]
19:13:13.242 | INF pkg/compute/executor.go:195 > cleaning up execution [NodeID:node-1] [execution:e-50dd14d0-c873-43b7-a49a-bda5f86889ce] [job:9128751d-7748-40a8-9c0c-c9c7ad178081]
....
19:13:13.242 | ERR pkg/compute/executor.go:480 > Publishing execution e-50dd14d0-c873-43b7-a49a-bda5f86889ce failed error="execution e-50dd14d0-c873-43b7-a49a-bda5f86889ce is in state Cancelled but expected one of [Running]" [NodeID:node-1] [execution:e-50dd14d0-c873-43b7-a49a-bda5f86889ce] [job:9128751d-7748-40a8-9c0c-c9c7ad178081]
....
19:13:13.242 | ERR pkg/compute/executor.go:488 > Failed to update execution (e-50dd14d0-c873-43b7-a49a-bda5f86889ce) state to failed: execution e-50dd14d0-c873-43b7-a49a-bda5f86889ce is in terminal state Cancelled and cannot transition to Failed error="execution e-50dd14d0-c873-43b7-a49a-bda5f86889ce is in terminal state Cancelled and cannot transition to Failed" [NodeID:node-1] [execution:e-50dd14d0-c873-43b7-a49a-bda5f86889ce] [job:9128751d-7748-40a8-9c0c-c9c7ad178081]
```


Discussion / hints:
This is still the case with 1.4.0. To clarify, the job state is correct, but the logs on the compute nodes are the ones misleading:

This is a description of a job that has been stopped
```
→    bacalhau job describe j-e209ab04-285c-4930-b908-e8d0e2e6772a
ID            = j-e209ab04-285c-4930-b908-e8d0e2e6772a
Name          = j-e209ab04-285c-4930-b908-e8d0e2e6772a
Namespace     = default
Type          = batch
State         = Stopped
Message       = Job requested to stop before completion
Count         = 1
Created Time  = 2024-06-26 13:27:08
Modified Time = 2024-06-26 13:27:35
Version       = 0

Summary

Job History
 TIME                 REV.  STATE    TOPIC       EVENT
 2024-06-26 15:27:08  1     Pending  Submission  Job submitted
 2024-06-26 15:27:10  2     Running
 2024-06-26 15:27:35  3     Stopped  Scheduling  Job requested to stop before completion

Executions
 ID          NODE ID     STATE      DESIRED  REV.  CREATED  MODIFIED  COMMENT
 e-fdf5f485  n-6e8998ad  Cancelled  Stopped  7     37s ago  10s ago   Accepted job

Execution e-fdf5f485 History
 TIME                 REV.  STATE              TOPIC            EVENT
 2024-06-26 15:27:08  1     New
 2024-06-26 15:27:08  2     AskForBid
 2024-06-26 15:27:10  3     AskForBidAccepted  Requesting Node  Accepted job
 2024-06-26 15:27:10  4     AskForBidAccepted
 2024-06-26 15:27:10  5     BidAccepted
 2024-06-26 15:27:35  6     BidAccepted        Scheduling       Execution stop requested because job has been stopped
 2024-06-26 15:27:35  7     Cancelled
```
This is happening, mainly because we wait for the job execution to complete, and eventually transition the job state to completion. However, this is problematic if we cancel the job in between. As a cancellation of job would be like normal closure of execution, and eventually the [`Run`](https://github.com/bacalhau-project/bacalhau/blob/main/pkg/compute/executor.go#L294) function would try to mark the job as `Completed`, when it has already been `Cancelled`. 

This is the cause of first error log we see in compute node logs. At the same time in this [`Run`](https://github.com/bacalhau-project/bacalhau/blob/main/pkg/compute/executor.go#L305) function we try to mark the job as `Failed`, if there is some error(  the error raised when we try to transition it to complete, when it is already marked cancelled). This is the cause of second log. 

There are couple of things we can do here 

- Before marking the job as completed, we take a look that if the job is already cancelled we just exit.
- Ignore the error in case the job is already in one of the terminal state. Ideally, if job has already reached to a terminal state, we should not overwrite it. Meaning we can log it as warn but no need to log it as error. We can simply ignore it. 

Let me know what you folks thing @wdbaruni  @frrist  !
> Before marking the job as completed, we take a look that if the job is already cancelled we just exit.

There will always be race conditions if you don't check and update atomically and within a transaction. You can minimize the window, but you can never prevent it. That is why with have an update condition at the execution store level that verifies the entry in boltdb is in a desired state before updating

> Ignore the error in case the job is already in one of the terminal state. Ideally, if job has already reached to a terminal state, we should not overwrite it. Meaning we can log it as warn but no need to log it as error. We can simply ignore it.

Yeah thats what the update condition is doing where we don't overwrite the state if it is terminal, and we return an error instead. What we should be doing is define more error types at the execution store level, have a specific error type such as `ErrAlreadyTerminal` or something like that, and then have the executor ignore logging the error

In addition to that, we should evaluate not calling `handleFailure` inside [Run](https://github.com/bacalhau-project/bacalhau/blob/ec2e2d69eb091605fd69da6ba9d9e2039c08eb0c/pkg/compute/executor.go#L305) method if the context was cancelled, but we need to evaluate what other cases the context can be cancelled and whether we should fail the execution in those cases or not

I think having more error types for now should be more than enough

Fix patch:
diff --git a/pkg/compute/endpoint.go b/pkg/compute/endpoint.go
index 3473901394..b1e62e864a 100644
--- a/pkg/compute/endpoint.go
+++ b/pkg/compute/endpoint.go
@@ -140,13 +140,6 @@ func (s BaseEndpoint) CancelExecution(ctx context.Context, request CancelExecuti
 			localExecutionState.Execution.ID, localExecutionState.State)
 	}
 
-	if localExecutionState.State.IsExecuting() {
-		err = s.executor.Cancel(ctx, localExecutionState)
-		if err != nil {
-			return CancelExecutionResponse{}, err
-		}
-	}
-
 	err = s.executionStore.UpdateExecutionState(ctx, store.UpdateExecutionStateRequest{
 		ExecutionID: request.ExecutionID,
 		NewState:    store.ExecutionStateCancelled,
@@ -155,13 +148,21 @@ func (s BaseEndpoint) CancelExecution(ctx context.Context, request CancelExecuti
 	if err != nil {
 		return CancelExecutionResponse{}, err
 	}
+
+	if localExecutionState.State.IsExecuting() {
+		err = s.executor.Cancel(ctx, localExecutionState)
+		if err != nil {
+			return CancelExecutionResponse{}, err
+		}
+	}
 	return CancelExecutionResponse{
 		ExecutionMetadata: NewExecutionMetadata(localExecutionState.Execution),
 	}, nil
 }
 
 func (s BaseEndpoint) ExecutionLogs(ctx context.Context, request ExecutionLogsRequest) (
-	<-chan *concurrency.AsyncResult[models.ExecutionLog], error) {
+	<-chan *concurrency.AsyncResult[models.ExecutionLog], error,
+) {
 	return s.logServer.GetLogStream(ctx, executor.LogStreamRequest{
 		ExecutionID: request.ExecutionID,
 		Tail:        request.Tail,
diff --git a/pkg/compute/executor.go b/pkg/compute/executor.go
index 567467fc31..332bd5908f 100644
--- a/pkg/compute/executor.go
+++ b/pkg/compute/executor.go
@@ -20,7 +20,7 @@ import (
 	"github.com/bacalhau-project/bacalhau/pkg/system"
 )
 
-const StorageDirectoryPerms = 0755
+const StorageDirectoryPerms = 0o755
 
 type BaseExecutorParams struct {
 	ID                     string
@@ -66,7 +66,8 @@ func prepareInputVolumes(
 	ctx context.Context,
 	strgprovider storage.StorageProvider,
 	storageDirectory string, inputSources ...*models.InputSource) (
-	[]storage.PreparedStorage, func(context.Context) error, error) {
+	[]storage.PreparedStorage, func(context.Context) error, error,
+) {
 	inputVolumes, err := storage.ParallelPrepareStorage(ctx, strgprovider, storageDirectory, inputSources...)
 	if err != nil {
 		return nil, nil, err
@@ -80,7 +81,8 @@ func prepareWasmVolumes(
 	ctx context.Context,
 	strgprovider storage.StorageProvider,
 	storageDirectory string, wasmEngine wasmmodels.EngineSpec) (
-	map[string][]storage.PreparedStorage, func(context.Context) error, error) {
+	map[string][]storage.PreparedStorage, func(context.Context) error, error,
+) {
 	importModuleVolumes, err := storage.ParallelPrepareStorage(ctx, strgprovider, storageDirectory, wasmEngine.ImportModules...)
 	if err != nil {
 		return nil, nil, err
@@ -300,7 +302,7 @@ func (e *BaseExecutor) Run(ctx context.Context, state store.LocalExecutionState)
 	stopwatch := telemetry.Timer(ctx, jobDurationMilliseconds, state.Execution.Job.MetricAttributes()...)
 	topic := EventTopicExecutionRunning
 	defer func() {
-		if err != nil {
+		if err != nil && err.Error() != executor.ErrAlreadyCancelled.Error() {
 			e.handleFailure(ctx, state, err, topic)
 		}
 		dur := stopwatch()
@@ -352,7 +354,7 @@ func (e *BaseExecutor) Run(ctx context.Context, state store.LocalExecutionState)
 		return err
 	}
 	if result.ErrorMsg != "" {
-		return fmt.Errorf("execution error: %s", result.ErrorMsg)
+		return fmt.Errorf(result.ErrorMsg)
 	}
 	jobsCompleted.Add(ctx, 1)
 
@@ -416,7 +418,8 @@ func (e *BaseExecutor) Run(ctx context.Context, state store.LocalExecutionState)
 
 // Publish the result of an execution after it has been verified.
 func (e *BaseExecutor) publish(ctx context.Context, localExecutionState store.LocalExecutionState,
-	resultFolder string) (publishedResult models.SpecConfig, err error) {
+	resultFolder string,
+) (publishedResult models.SpecConfig, err error) {
 	execution := localExecutionState.Execution
 	log.Ctx(ctx).Debug().Msgf("Publishing execution %s", execution.ID)
 
diff --git a/pkg/executor/docker/executor.go b/pkg/executor/docker/executor.go
index cd90f9009d..31d7fe63aa 100644
--- a/pkg/executor/docker/executor.go
+++ b/pkg/executor/docker/executor.go
@@ -124,7 +124,6 @@ func (e *Executor) Start(ctx context.Context, request *executor.RunCommandReques
 	// bacalhau execution label _before_ we do anything else.  If we are able to find one then we
 	// will use that container in the executionHandler that we create.
 	containerID, err := e.FindRunningContainer(ctx, request.ExecutionID)
-
 	if err != nil {
 		// Unable to find a running container for this execution, we will instead check for a handler, and
 		// failing that will create a new container.
@@ -153,6 +152,7 @@ func (e *Executor) Start(ctx context.Context, request *executor.RunCommandReques
 		containerID = jobContainer.ID
 	}
 
+	childCtx, cancel := context.WithCancelCause(ctx)
 	handler := &executionHandler{
 		client: e.client,
 		logger: log.With().
@@ -169,12 +169,13 @@ func (e *Executor) Start(ctx context.Context, request *executor.RunCommandReques
 		waitCh:      make(chan bool),
 		activeCh:    make(chan bool),
 		running:     atomic.NewBool(false),
+		cancelFunc:  cancel,
 	}
 
 	// register the handler for this executionID
 	e.handlers.Put(request.ExecutionID, handler)
 	// run the container.
-	go handler.run(ctx)
+	go handler.run(childCtx)
 	return nil
 }
 
@@ -234,7 +235,8 @@ func (e *Executor) Cancel(ctx context.Context, executionID string) error {
 	if !found {
 		return fmt.Errorf("canceling execution (%s): %w", executionID, executor.ErrNotFound)
 	}
-	return handler.kill(ctx)
+	handler.cancelFunc(executor.ErrAlreadyCancelled)
+	return nil
 }
 
 // GetLogStream provides a stream of output logs for a specific execution.
@@ -438,7 +440,8 @@ func configureDevices(ctx context.Context, resources *models.Resources) ([]conta
 }
 
 func makeContainerMounts(
-	ctx context.Context, inputs []storage.PreparedStorage, outputs []*models.ResultPath, resultsDir string) ([]mount.Mount, error) {
+	ctx context.Context, inputs []storage.PreparedStorage, outputs []*models.ResultPath, resultsDir string,
+) ([]mount.Mount, error) {
 	// the actual mounts we will give to the container
 	// these are paths for both input and output data
 	var mounts []mount.Mount
diff --git a/pkg/executor/docker/handler.go b/pkg/executor/docker/handler.go
index a14bc028fc..e1a1e54c9d 100644
--- a/pkg/executor/docker/handler.go
+++ b/pkg/executor/docker/handler.go
@@ -44,6 +44,8 @@ type executionHandler struct {
 	waitCh chan bool
 	// true until the run method returns
 	running *atomic.Bool
+	// cancel function
+	cancelFunc context.CancelCauseFunc
 
 	//
 	// results
@@ -89,9 +91,12 @@ func (h *executionHandler) run(ctx context.Context) {
 	select {
 	case <-ctx.Done():
 		// failure case, the context has been canceled. We are aborting this execution
-		reason := fmt.Errorf("context canceled while waiting on container status: %w", ctx.Err())
-		h.logger.Err(reason).Msg("cancel waiting on container status")
-		h.result = executor.NewFailedResult(reason.Error())
+		cause := context.Cause(ctx)
+		if cause == nil {
+			cause = fmt.Errorf("context canceled while waiting on container status: %w", ctx.Err())
+		}
+		h.logger.Err(cause).Msg("cancel waiting on container status")
+		h.result = executor.NewFailedResult(cause.Error())
 		// the context was canceled, bail.
 		return
 	case err := <-errCh:
diff --git a/pkg/executor/types.go b/pkg/executor/types.go
index f87daa8d32..f4197876f7 100644
--- a/pkg/executor/types.go
+++ b/pkg/executor/types.go
@@ -79,6 +79,8 @@ var (
 	// ErrAlreadyStarted is returned when trying to start an already started execution.
 	ErrAlreadyStarted = fmt.Errorf("execution already started")
 
+	ErrAlreadyCancelled = fmt.Errorf("execution already cancelled")
+
 	// ErrAlreadyComplete is returned when action is attempted on an execution that is already complete.
 	ErrAlreadyComplete = fmt.Errorf("execution already complete")
 


Test patch:
diff --git a/pkg/executor/docker/executor_test.go b/pkg/executor/docker/executor_test.go
index 5287a05b44..f3a5273f89 100644
--- a/pkg/executor/docker/executor_test.go
+++ b/pkg/executor/docker/executor_test.go
@@ -1,4 +1,4 @@
-//go:build unit || !integration
+//go:build integration || !unit
 
 package docker
 
@@ -285,6 +285,7 @@ func (s *ExecutorTestSuite) TestDockerResourceLimitsMemory() {
 		require.LessOrEqual(s.T(), diff, 4096, "the difference between the container reported memory and the configured limit exceeds the page size")
 	}
 }
+
 func (s *ExecutorTestSuite) TestDockerNetworkingFull() {
 	task := mock.TaskBuilder().
 		Network(models.NewNetworkConfigBuilder().
@@ -401,6 +402,54 @@ func (s *ExecutorTestSuite) TestDockerNetworkingFiltersHTTPS() {
 	require.Empty(s.T(), result.STDOUT)
 }
 
+func (s *ExecutorTestSuite) TestDockerExecutionCancellation() {
+	resultC := make(chan *models.RunCommandResult, 1)
+	errC := make(chan error, 1)
+	executionID := uuid.New().String()
+
+	es, err := dockermodels.NewDockerEngineBuilder("ubuntu").
+		WithEntrypoint("bash", "-c", "sleep 30").
+		Build()
+
+	s.Require().NoError(err)
+
+	task := mock.TaskBuilder().
+		Engine(es).
+		BuildOrDie()
+
+	jobCtx := context.Background()
+	go func() {
+		result, err := s.runJobWithContext(jobCtx, task, executionID)
+		if err != nil {
+			errC <- err
+		} else {
+			resultC <- result
+		}
+	}()
+
+	s.Eventually(func() bool {
+		handler, ok := s.executor.handlers.Get(executionID)
+		return ok && handler.active()
+	}, time.Second*10, time.Millisecond*100, "Could not find a running container")
+
+	// This is important to do. In our docker executor, we set active to true, before calling the docker client with ContainerStart
+	// Hence there is a bit of time before the container actually gets started. The correct way of identifying that whether
+	// a contianer has started or not is via activeCh. We want to make sure that contianer is started before canceling the execution.
+	handler, _ := s.executor.handlers.Get(executionID)
+	<-handler.activeCh
+
+	err = s.executor.Cancel(jobCtx, executionID)
+	s.Require().NoError(err)
+
+	select {
+	case err := <-errC:
+		s.Require().Failf("Executor run should have returned a result, but instead returned err: %w", err.Error())
+	case result := <-resultC:
+		s.Require().NotNil(result)
+		s.Require().Equal(executor.ErrAlreadyCancelled.Error(), result.ErrorMsg)
+	}
+}
+
 func (s *ExecutorTestSuite) TestDockerNetworkingAppendsHTTPHeader() {
 	s.server.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		_, err := w.Write([]byte(r.Header.Get("X-Bacalhau-Job-ID")))
diff --git a/pkg/nats/stream/streaming_client_interactions_test.go b/pkg/nats/stream/streaming_client_interactions_test.go
index eac5931613..a0d86aff87 100644
--- a/pkg/nats/stream/streaming_client_interactions_test.go
+++ b/pkg/nats/stream/streaming_client_interactions_test.go
@@ -5,17 +5,20 @@ package stream
 import (
 	"context"
 	"encoding/json"
+	"testing"
+	"time"
+
 	"github.com/bacalhau-project/bacalhau/pkg/lib/network"
 	nats_helper "github.com/bacalhau-project/bacalhau/pkg/nats"
 	"github.com/nats-io/nats-server/v2/server"
 	"github.com/nats-io/nats.go"
 	"github.com/stretchr/testify/suite"
-	"testing"
-	"time"
 )
 
-const subjectName = "topic.stream"
-const testString = "Hello from bacalhau"
+const (
+	subjectName = "topic.stream"
+	testString  = "Hello from bacalhau"
+)
 
 type StreamingClientInteractionTestSuite struct {
 	suite.Suite
@@ -28,7 +31,6 @@ type StreamingClientInteractionTestSuite struct {
 
 type testData struct {
 	contextCancelled    bool
-	streamReplySub      string
 	heartBeatRequestSub string
 }
 
@@ -37,7 +39,6 @@ func (s *StreamingClientInteractionTestSuite) SetupSuite() {
 	s.natServer = s.createNatsServer()
 	s.pc = s.createProducerClient()
 	s.cc = s.createConsumerClient()
-
 }
 
 func (s *StreamingClientInteractionTestSuite) TearDownSuite() {
@@ -47,7 +48,6 @@ func (s *StreamingClientInteractionTestSuite) TearDownSuite() {
 }
 
 func (s *StreamingClientInteractionTestSuite) createNatsServer() *server.Server {
-	ctx := context.Background()
 	port, err := network.GetFreePort()
 	s.Require().NoError(err)
 
@@ -55,10 +55,11 @@ func (s *StreamingClientInteractionTestSuite) createNatsServer() *server.Server
 		Port: port,
 	}
 
-	ns, err := nats_helper.NewServerManager(ctx, nats_helper.ServerManagerParams{
+	ns, err := nats_helper.NewServerManager(s.ctx, nats_helper.ServerManagerParams{
 		Options: &serverOpts,
 	})
 	s.Require().NoError(err)
+
 	return ns.Server
 }
 
@@ -69,47 +70,49 @@ func (s *StreamingClientInteractionTestSuite) createProducerClient() *ProducerCl
 	pc, err := NewProducerClient(s.ctx, ProducerClientParams{
 		Conn: clientManager.Client,
 		Config: StreamProducerClientConfig{
-			HeartBeatIntervalDuration:        100 * time.Millisecond,
-			HeartBeatRequestTimeout:          50 * time.Millisecond,
-			StreamCancellationBufferDuration: 100 * time.Millisecond,
+			HeartBeatIntervalDuration:        200 * time.Millisecond,
+			HeartBeatRequestTimeout:          100 * time.Millisecond,
+			StreamCancellationBufferDuration: 200 * time.Millisecond,
 		},
 	})
 
 	s.Require().NoError(err)
+	s.Eventually(func() bool {
+		return pc.Conn.IsConnected()
+	}, 500*time.Millisecond, 10*time.Millisecond)
 	return pc
 }
 
 func (s *StreamingClientInteractionTestSuite) createConsumerClient() *ConsumerClient {
-
 	clientManager, err := nats_helper.NewClientManager(s.ctx, s.natServer.ClientURL(), nats.Name("streaming-test"))
 	s.Require().NoError(err)
 
 	cc, err := NewConsumerClient(ConsumerClientParams{
 		Conn: clientManager.Client,
 		Config: StreamConsumerClientConfig{
-			StreamCancellationBufferDuration: 50 * time.Millisecond,
+			StreamCancellationBufferDuration: 100 * time.Millisecond,
 		},
 	})
 
 	s.Require().NoError(err)
+	s.Eventually(func() bool {
+		return cc.Conn.IsConnected()
+	}, 500*time.Millisecond, 10*time.Millisecond)
 	return cc
 }
 
-func TestStreamingClientTestSuit(t *testing.T) {
+func TestStreamingClientTestSuite(t *testing.T) {
 	suite.Run(t, new(StreamingClientInteractionTestSuite))
 }
 
 func (s *StreamingClientInteractionTestSuite) TestStreamConsumerClientGoingDown() {
-
 	// Set up for the test
+	ctx := context.Background()
 	td := &testData{}
-	clientManager, err := nats_helper.NewClientManager(s.ctx, s.natServer.ClientURL(), nats.Name("stream-testing-consumer-going-down"))
-	s.Require().NoError(err)
 
 	// Produce some data once asked for
-	ctx, cancel := context.WithCancel(s.ctx)
-	_, err = clientManager.Client.Subscribe(subjectName, func(msg *nats.Msg) {
-
+	ctx, cancel := context.WithCancel(ctx)
+	_, err := s.pc.Conn.Subscribe(subjectName, func(msg *nats.Msg) {
 		s.Require().NotNil(msg)
 
 		var streamRequest Request
@@ -148,38 +151,33 @@ func (s *StreamingClientInteractionTestSuite) TestStreamConsumerClientGoingDown(
 					sMsgData, err := json.Marshal(sMsg)
 					s.Require().NoError(err)
 
-					clientManager.Client.Publish(msg.Reply, sMsgData)
+					s.pc.Conn.Publish(msg.Reply, sMsgData)
 				}
 			}
-
 		}()
 	})
 	s.Require().NoError(err)
+
+	// Wait for some time before we publish a request.
+	time.Sleep(1 * time.Second)
+
 	data, err := json.Marshal(testString)
 	s.Require().NoError(err)
 
 	_, err = s.cc.OpenStream(s.ctx, subjectName, data)
 	s.Require().NoError(err)
 
-	// Close the Consumer Client After Certain Time
-	go func() {
-
-		ticker := time.NewTicker(1 * time.Second)
-		defer ticker.Stop()
-
-		for {
-			select {
-			case <-ticker.C:
-				s.cc.Conn.Close()
-				return
-			}
-		}
+	// Wait for the message to be received by subscriber above
+	s.Eventually(func() bool {
+		return td.heartBeatRequestSub != ""
+	}, time.Second*10, time.Millisecond*100, "Streaming request yet not received")
 
-	}()
+	// Close the Consumer Client After Certain Time
+	time.Sleep(time.Second * 1)
+	s.cc.Conn.Close()
 
 	// Validate that producer client does the cleanup
 	s.Eventually(func() bool {
 		return td.contextCancelled
-	}, 1800*time.Millisecond, 100*time.Millisecond)
-
+	}, 2800*time.Millisecond, 50*time.Millisecond)
 }
diff --git a/pkg/test/compute/cancelled_docker_execution_test.go b/pkg/test/compute/cancelled_docker_execution_test.go
new file mode 100644
index 0000000000..e3b0b78750
--- /dev/null
+++ b/pkg/test/compute/cancelled_docker_execution_test.go
@@ -0,0 +1,209 @@
+//go:build integration || !unit
+
+package compute
+
+import (
+	"context"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/bacalhau-project/bacalhau/pkg/compute"
+	"github.com/bacalhau-project/bacalhau/pkg/docker"
+	dockermodels "github.com/bacalhau-project/bacalhau/pkg/executor/docker/models"
+	noop_publisher "github.com/bacalhau-project/bacalhau/pkg/publisher/noop"
+	noop_storage "github.com/bacalhau-project/bacalhau/pkg/storage/noop"
+	"github.com/bacalhau-project/bacalhau/pkg/test/mock"
+
+	"github.com/bacalhau-project/bacalhau/pkg/authz"
+	"github.com/bacalhau-project/bacalhau/pkg/compute/store"
+	"github.com/bacalhau-project/bacalhau/pkg/compute/store/boltdb"
+	"github.com/bacalhau-project/bacalhau/pkg/compute/store/resolver"
+	"github.com/bacalhau-project/bacalhau/pkg/config/configenv"
+	executor_common "github.com/bacalhau-project/bacalhau/pkg/executor"
+	dockerexecutor "github.com/bacalhau-project/bacalhau/pkg/executor/docker"
+	"github.com/bacalhau-project/bacalhau/pkg/lib/provider"
+	"github.com/bacalhau-project/bacalhau/pkg/models"
+	"github.com/bacalhau-project/bacalhau/pkg/node"
+	"github.com/bacalhau-project/bacalhau/pkg/publicapi"
+	"github.com/bacalhau-project/bacalhau/pkg/publisher"
+	"github.com/bacalhau-project/bacalhau/pkg/setup"
+	"github.com/bacalhau-project/bacalhau/pkg/storage"
+	"github.com/labstack/echo/v4"
+	"github.com/stretchr/testify/suite"
+)
+
+type CancelDockerExecutionSuite struct {
+	suite.Suite
+	node             *node.Compute
+	store            store.ExecutionStore
+	callbackStore    *CallbackStore
+	bidChannel       chan compute.BidResult
+	failureChannel   chan compute.ComputeError
+	completedChannel chan compute.RunResult
+}
+
+func TestCancelDockerExecutionSuite(t *testing.T) {
+	suite.Run(t, new(CancelDockerExecutionSuite))
+}
+
+func (s *CancelDockerExecutionSuite) TestCancelledDockerExeuction() {
+	docker.MustHaveDocker(s.T())
+	ctx := context.Background()
+
+	_, c := setup.SetupBacalhauRepoForTesting(s.T())
+
+	executionStore, err := boltdb.NewStore(context.Background(), filepath.Join(s.T().TempDir(), "executions.db"))
+	s.store = executionStore
+	s.Require().NoError(err)
+
+	s.callbackStore = &CallbackStore{}
+	s.callbackStore.GetExecutionFn = s.store.GetExecution
+	s.callbackStore.GetExecutionsFn = s.store.GetExecutions
+	s.callbackStore.GetLiveExecutionsFn = s.store.GetLiveExecutions
+	s.callbackStore.GetExecutionHistoryFn = s.store.GetExecutionHistory
+	s.callbackStore.CreateExecutionFn = s.store.CreateExecution
+	s.callbackStore.DeleteExecutionFn = s.store.DeleteExecution
+	s.callbackStore.GetExecutionCountFn = s.store.GetExecutionCount
+	s.callbackStore.CloseFn = s.store.Close
+
+	// This test we are testing that when an execution is cancelled, the executor, no longer
+	// tries to mark the execution as completed. The sync group here is 3, mainly because
+	// the execution state will go to
+	// 1. ExecutionStateBidAccepted
+	// 2. ExecutionStateRunning
+	// 3. ExecutionStateCancelled
+	// If any more requests come in to update execution state, the wait group would go negative indicating some other state update is
+	// being called to.
+	exeuctionStateUpdateWg := sync.WaitGroup{}
+	exeuctionStateUpdateWg.Add(3)
+	s.callbackStore.UpdateExecutionStateFn = func(ctx context.Context, request store.UpdateExecutionStateRequest) error {
+		defer exeuctionStateUpdateWg.Done()
+		return s.store.UpdateExecutionState(ctx, request)
+	}
+
+	computeConfig, err := node.NewComputeConfigWith(c.Node.ComputeStoragePath, node.ComputeConfigParams{
+		TotalResourceLimits: models.Resources{
+			CPU: 2,
+		},
+		ExecutionStore: s.callbackStore,
+	})
+	s.Require().NoError(err)
+
+	dockerExecutor, err := dockerexecutor.NewExecutor("compute-node-docker-executor-test", configenv.Testing.Node.Compute.ManifestCache)
+	s.Require().NoError(err)
+
+	apiServer, err := publicapi.NewAPIServer(publicapi.ServerParams{
+		Router:     echo.New(),
+		Address:    "0.0.0.0",
+		Port:       0,
+		Config:     publicapi.DefaultConfig(),
+		Authorizer: authz.AlwaysAllow,
+	})
+
+	s.NoError(err)
+
+	storagePath := s.T().TempDir()
+	repoPath := s.T().TempDir()
+	s.bidChannel = make(chan compute.BidResult, 1)
+	s.completedChannel = make(chan compute.RunResult)
+	s.failureChannel = make(chan compute.ComputeError)
+
+	noopstorage := noop_storage.NewNoopStorage()
+	callback := compute.CallbackMock{
+		OnBidCompleteHandler: func(ctx context.Context, result compute.BidResult) {
+			s.bidChannel <- result
+		},
+		OnRunCompleteHandler: func(ctx context.Context, result compute.RunResult) {
+			s.completedChannel <- result
+		},
+		OnComputeFailureHandler: func(ctx context.Context, err compute.ComputeError) {
+			s.failureChannel <- err
+		},
+	}
+
+	s.node, err = node.NewComputeNode(
+		context.Background(),
+		"test",
+		apiServer,
+		computeConfig,
+		storagePath,
+		repoPath,
+		provider.NewNoopProvider[storage.Storage](noopstorage),
+		provider.NewMappedProvider(map[string]executor_common.Executor{
+			models.EngineDocker: dockerExecutor,
+		}),
+		provider.NewNoopProvider[publisher.Publisher](noop_publisher.NewNoopPublisher()),
+		callback,
+		ManagementEndpointMock{},
+		map[string]string{},
+		HeartbeatClientMock{},
+	)
+
+	s.Require().NoError(err)
+	stateResolver := *resolver.NewStateResolver(resolver.StateResolverParams{
+		ExecutionStore: s.node.ExecutionStore,
+	})
+
+	es, err := dockermodels.NewDockerEngineBuilder("ubuntu").
+		WithEntrypoint("bash", "-c", "sleep 10").
+		Build()
+
+	s.Require().NoError(err)
+	s.Require().NoError(err)
+	task := mock.TaskBuilder().
+		Engine(es).
+		BuildOrDie()
+	job := mock.Job()
+	job.Tasks[0] = task
+	job.Normalize()
+
+	execution := mock.ExecutionForJob(job)
+	executionID := s.prepareAndAskForBid(ctx, execution)
+
+	_, err = s.node.LocalEndpoint.BidAccepted(ctx, compute.BidAcceptedRequest{ExecutionID: executionID})
+	s.Require().NoError(err)
+
+	err = stateResolver.Wait(ctx, executionID, resolver.CheckForState(store.ExecutionStateBidAccepted))
+	s.NoError(err)
+
+	err = stateResolver.Wait(ctx, executionID, resolver.CheckForState(store.ExecutionStateRunning))
+	s.NoError(err)
+
+	// We need to wait for the container to become active, before we cancel the execution.
+	time.Sleep(time.Second * 2)
+	_, err = s.node.LocalEndpoint.CancelExecution(ctx, compute.CancelExecutionRequest{
+		ExecutionID: executionID,
+	})
+	s.NoError(err)
+	err = stateResolver.Wait(ctx, executionID, resolver.CheckForState(store.ExecutionStateCancelled))
+	s.NoError(err)
+	exeuctionStateUpdateWg.Wait()
+}
+
+func (s *CancelDockerExecutionSuite) askForBid(ctx context.Context, execution *models.Execution) compute.BidResult {
+	_, err := s.node.LocalEndpoint.AskForBid(ctx, compute.AskForBidRequest{
+		RoutingMetadata: compute.RoutingMetadata{
+			TargetPeerID: s.node.ID,
+			SourcePeerID: s.node.ID,
+		},
+		Execution:       execution,
+		WaitForApproval: true,
+	})
+	s.NoError(err)
+
+	select {
+	case result := <-s.bidChannel:
+		return result
+	case <-time.After(5 * time.Second):
+		s.FailNow("did not receive a bid response")
+		return compute.BidResult{}
+	}
+}
+
+func (s *CancelDockerExecutionSuite) prepareAndAskForBid(ctx context.Context, execution *models.Execution) string {
+	result := s.askForBid(ctx, execution)
+	s.True(result.Accepted)
+	return result.ExecutionID
+}