Fix context usage in shim package

runner/internal/shim/backends/aws.go

@@ -2,6 +2,7 @@ package backends
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
@@ -23,10 +24,10 @@ func NewAWSBackend() *AWSBackend {
 // * Red Hat and CentOS: may increment trailing letters in some versions – not supported.
 // * Other legacy systems: /dev/sda => /dev/sda.
 // More: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/device_naming.html
-func (e *AWSBackend) GetRealDeviceName(volumeID, deviceName string) (string, error) {
+func (e *AWSBackend) GetRealDeviceName(ctx context.Context, volumeID, deviceName string) (string, error) {
 	// Run the lsblk command to get block device information
 	// On AWS, SERIAL contains volume id.
-	cmd := exec.Command("lsblk", "-o", "NAME,SERIAL")
+	cmd := exec.CommandContext(ctx, "lsblk", "-o", "NAME,SERIAL")
 	var out bytes.Buffer
 	cmd.Stdout = &out
 	if err := cmd.Run(); err != nil {
@@ -67,7 +68,7 @@ func (e *AWSBackend) GetRealDeviceName(volumeID, deviceName string) (string, err
 	}
 
 	// Run lsblk again to check for partitions on the base device
-	cmd = exec.Command("lsblk", "-ln", "-o", "NAME", baseDevice)
+	cmd = exec.CommandContext(ctx, "lsblk", "-ln", "-o", "NAME", baseDevice)
 	out.Reset()
 	cmd.Stdout = &out
 	if err := cmd.Run(); err != nil {

runner/internal/shim/backends/base.go

@@ -1,6 +1,8 @@
 package backends
 
+import "context"
+
 type Backend interface {
 	// GetRealDeviceName returns the real device name for the given volume ID and virtual device name.
-	GetRealDeviceName(volumeID, deviceName string) (string, error)
+	GetRealDeviceName(ctx context.Context, volumeID, deviceName string) (string, error)
 }

runner/internal/shim/backends/gcp.go

@@ -1,6 +1,7 @@
 package backends
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -15,7 +16,7 @@ func NewGCPBackend() *GCPBackend {
 }
 
 // GetRealDeviceName resolves device names according to https://cloud.google.com/compute/docs/disks/disk-symlinks
-func (e *GCPBackend) GetRealDeviceName(volumeID, deviceName string) (string, error) {
+func (e *GCPBackend) GetRealDeviceName(ctx context.Context, volumeID, deviceName string) (string, error) {
 	// Try resolving first partition or external volumes
 	realDeviceName, err := os.Readlink(fmt.Sprintf("/dev/disk/by-id/google-%s-part1", deviceName))
 	if err != nil {

runner/internal/shim/docker.go

@@ -535,12 +535,12 @@ func unmountVolumes(ctx context.Context, taskConfig TaskConfig) error {
 	var failed []string
 	for _, volume := range taskConfig.Volumes {
 		mountPoint := getVolumeMountPoint(volume.Name)
-		cmd := exec.Command("mountpoint", mountPoint)
+		cmd := exec.CommandContext(ctx, "mountpoint", mountPoint)
 		if output, err := cmd.CombinedOutput(); err != nil {
 			log.Info(ctx, "skipping", "mountpoint", mountPoint, "output", output)
 			continue
 		}
-		cmd = exec.Command("umount", "-qf", mountPoint)
+		cmd = exec.CommandContext(context.TODO(), "umount", "-qf", mountPoint)
 		if output, err := cmd.CombinedOutput(); err != nil {
 			log.Error(ctx, "failed to unmount", "mountpoint", mountPoint, "output", output)
 			failed = append(failed, mountPoint)
@@ -559,7 +559,7 @@ func formatAndMountVolume(ctx context.Context, volume VolumeInfo) error {
 	if err != nil {
 		return tracerr.Wrap(err)
 	}
-	deviceName, err := backend.GetRealDeviceName(volume.VolumeId, volume.DeviceName)
+	deviceName, err := backend.GetRealDeviceName(ctx, volume.VolumeId, volume.DeviceName)
 	if err != nil {
 		return tracerr.Wrap(err)
 	}
@@ -618,7 +618,7 @@ func prepareInstanceMountPoints(taskConfig TaskConfig) error {
 // Returns true if the file system is created.
 func initFileSystem(ctx context.Context, deviceName string, errorIfNotExists bool) (bool, error) {
 	// Run the lsblk command to get filesystem type
-	cmd := exec.Command("lsblk", "-no", "FSTYPE", deviceName)
+	cmd := exec.CommandContext(ctx, "lsblk", "-no", "FSTYPE", deviceName)
 	var out bytes.Buffer
 	cmd.Stdout = &out
 	if err := cmd.Run(); err != nil {
@@ -636,7 +636,7 @@ func initFileSystem(ctx context.Context, deviceName string, errorIfNotExists boo
 	}
 
 	log.Debug(ctx, "formatting disk with ext4 filesystem...", "device", deviceName)
-	cmd = exec.Command("mkfs.ext4", "-F", deviceName)
+	cmd = exec.CommandContext(ctx, "mkfs.ext4", "-F", deviceName)
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return false, fmt.Errorf("failed to format disk: %w, output: %s", err, string(output))
 	}
@@ -655,7 +655,7 @@ func mountDisk(ctx context.Context, deviceName, mountPoint string, fsRootPerms o
 
 	// Mount the disk to the mount point
 	log.Debug(ctx, "mounting disk...", "device", deviceName, "mountpoint", mountPoint)
-	cmd := exec.Command("mount", deviceName, mountPoint)
+	cmd := exec.CommandContext(ctx, "mount", deviceName, mountPoint)
 	if output, err := cmd.CombinedOutput(); err != nil {
 		return fmt.Errorf("failed to mount disk: %w, output: %s", err, string(output))
 	}