diff --git a/internal/docker/docker.go b/internal/docker/docker.go
index d171038..9d4e300 100644
--- a/internal/docker/docker.go
+++ b/internal/docker/docker.go
@@ -51,12 +51,12 @@ func NewClient(endpoint string) (*Client, error) {
 	}, nil
 }
 
-func (c *Client) ListContainersByLabels(labels map[string]string) ([]Container, error) {
+func (c *Client) ListContainersByLabels(ctx context.Context, labels map[string]string) ([]Container, error) {
 	filters := make(map[string][]string)
 	for k, v := range labels {
 		filters["label"] = append(filters["label"], fmt.Sprintf("%s=%s", k, v))
 	}
-	containers, err := c.api.ListContainers(docker.ListContainersOptions{Filters: filters})
+	containers, err := c.api.ListContainers(docker.ListContainersOptions{Filters: filters, Context: ctx})
 	if err != nil {
 		return nil, err
 	}
diff --git a/internal/docker/docker_test.go b/internal/docker/docker_test.go
index 438643c..34a05a7 100644
--- a/internal/docker/docker_test.go
+++ b/internal/docker/docker_test.go
@@ -4,6 +4,7 @@
 package docker
 
 import (
+	"context"
 	"testing"
 
 	"github.com/fsouza/go-dockerclient"
@@ -48,7 +49,7 @@ func (s *S) TestGetContainersByLabel(c *check.C) {
 	err = client.api.StartContainer(cont2.ID, nil)
 	c.Assert(err, check.IsNil)
 
-	containers, err := client.ListContainersByLabels(map[string]string{"A": "VA", "B": "VB"})
+	containers, err := client.ListContainersByLabels(context.Background(), map[string]string{"A": "VA", "B": "VB"})
 	c.Assert(err, check.IsNil)
 	c.Assert(containers, check.DeepEquals, []Container{{ID: cont.ID}})
 }
diff --git a/main.go b/main.go
index 75806f5..52d4eda 100644
--- a/main.go
+++ b/main.go
@@ -51,37 +51,18 @@ func main() {
 			fatal("failed to create docker client: %v", err)
 
 		}
-		hostname, err := os.Hostname()
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+		mainContainer, err := getMainContainer(ctx, dockerClient)
+		cancel()
 		if err != nil {
-			fatal("failed to get hostname: %v", err)
+			fatal("failed to get main container: %v", err)
 		}
-		timeout := time.After(time.Second * 30)
-		var containers []docker.Container
-	loop:
-		for {
-			containers, err = dockerClient.ListContainersByLabels(map[string]string{
-				"io.kubernetes.container.name": hostname,
-				"io.kubernetes.pod.name":       hostname,
-			})
-			if err != nil {
-				fatal("failed to get containers: %v", err)
-			}
-			select {
-			case <-timeout:
-				break loop
-			case <-time.After(time.Second * 3):
-				break
-			}
-		}
-		if len(containers) != 1 {
-			fatal("failed to get main container from sidecar, got %v.", containers)
-		}
-		executor = &docker.Executor{Client: dockerClient, ContainerID: containers[0].ID}
+		executor = &docker.Executor{Client: dockerClient, ContainerID: mainContainer.ID}
 		filesystem = &executorFS{executor: executor}
 		defer func() {
 			fmt.Println("---- Building application image ----")
 			imgName := os.Getenv("DEPLOYAGENT_DST_IMAGE")
-			img, err := dockerClient.Commit(context.Background(), containers[0].ID, imgName)
+			img, err := dockerClient.Commit(context.Background(), mainContainer.ID, imgName)
 			if err != nil {
 				fatal("error commiting image %v: %v", imgName, err)
 			}
@@ -119,6 +100,30 @@ func main() {
 	}
 }
 
+func getMainContainer(ctx context.Context, dockerClient *docker.Client) (docker.Container, error) {
+	hostname, err := os.Hostname()
+	if err != nil {
+		return docker.Container{}, fmt.Errorf("failed to get hostname: %v", err)
+	}
+	for {
+		containers, err := dockerClient.ListContainersByLabels(ctx, map[string]string{
+			"io.kubernetes.container.name": hostname,
+			"io.kubernetes.pod.name":       hostname,
+		})
+		if err != nil {
+			return docker.Container{}, fmt.Errorf("failed to get containers: %v", err)
+		}
+		if len(containers) == 1 {
+			return containers[0], nil
+		}
+		select {
+		case <-ctx.Done():
+			return docker.Container{}, ctx.Err()
+		case <-time.After(time.Second * 1):
+		}
+	}
+}
+
 func fatal(format string, v ...interface{}) {
 	var w io.Writer = os.Stderr
 	file, err := os.OpenFile("/dev/termination-log", os.O_WRONLY|os.O_CREATE, 0666)