add contexts

beam-cloud · Jan 2, 2025 · e0d6b28 · e0d6b28
1 parent 5feceb1
commit e0d6b28
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/pkg/abstractions/shell/http.go b/pkg/abstractions/shell/http.go
@@ -101,5 +101,7 @@ func (g *shellGroup) ShellConnect(ctx echo.Context) error {
 		return nil
 	case <-clientCtx.Done():
 		return nil
+	case <-g.ss.ctx.Done():
+		return nil
 	}
 }
diff --git a/pkg/abstractions/shell/shell.go b/pkg/abstractions/shell/shell.go
@@ -301,6 +301,8 @@ func (ss *SSHShellService) keepAlive(ctx context.Context, containerId string, do
 
 	for {
 		select {
+		case <-ss.ctx.Done():
+			return
 		case <-ctx.Done():
 			return
 		case <-done:
@@ -317,6 +319,8 @@ func (ss *SSHShellService) waitForContainer(ctx context.Context, containerId str
 
 	for {
 		select {
+		case <-ss.ctx.Done():
+			return nil
 		case <-timeoutCtx.Done():
 			return fmt.Errorf("timed out waiting for container to be available")
 		default: