From 84390cc6a799616202faf27b3c38e71940394274 Mon Sep 17 00:00:00 2001
From: Lantao Liu <lantaol@google.com>
Date: Tue, 23 May 2017 05:47:48 +0000
Subject: [PATCH] Stop waiting on cancellation.

Signed-off-by: Lantao Liu <lantaol@google.com>
---
 pkg/server/container_stop.go      |  8 +++++---
 pkg/server/container_stop_test.go | 24 ++++++++++++++++++++++--
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/pkg/server/container_stop.go b/pkg/server/container_stop.go
index c686981c0..6f2a5d974 100644
--- a/pkg/server/container_stop.go
+++ b/pkg/server/container_stop.go
@@ -76,7 +76,7 @@ func (c *criContainerdService) StopContainer(ctx context.Context, r *runtime.Sto
 		return nil, fmt.Errorf("failed to stop container %q: %v", id, err)
 	}
 
-	err = c.waitContainerStop(id, time.Duration(r.GetTimeout())*time.Second)
+	err = c.waitContainerStop(ctx, id, time.Duration(r.GetTimeout())*time.Second)
 	if err == nil {
 		return &runtime.StopContainerResponse{}, nil
 	}
@@ -92,7 +92,7 @@ func (c *criContainerdService) StopContainer(ctx context.Context, r *runtime.Sto
 	}
 
 	// Wait forever until container stop is observed by event monitor.
-	if err := c.waitContainerStop(id, killContainerTimeout); err != nil {
+	if err := c.waitContainerStop(ctx, id, killContainerTimeout); err != nil {
 		return nil, fmt.Errorf("error occurs during waiting for container %q to stop: %v",
 			id, err)
 	}
@@ -100,7 +100,7 @@ func (c *criContainerdService) StopContainer(ctx context.Context, r *runtime.Sto
 }
 
 // waitContainerStop polls container state until timeout exceeds or container is stopped.
-func (c *criContainerdService) waitContainerStop(id string, timeout time.Duration) error {
+func (c *criContainerdService) waitContainerStop(ctx context.Context, id string, timeout time.Duration) error {
 	ticker := time.NewTicker(stopCheckPollInterval)
 	defer ticker.Stop()
 	timeoutTimer := time.NewTimer(timeout)
@@ -122,6 +122,8 @@ func (c *criContainerdService) waitContainerStop(id string, timeout time.Duratio
 			return nil
 		}
 		select {
+		case <-ctx.Done():
+			return fmt.Errorf("wait container %q is cancelled", id)
 		case <-timeoutTimer.C:
 			return fmt.Errorf("wait container %q stop timeout", id)
 		case <-ticker.C:
diff --git a/pkg/server/container_stop_test.go b/pkg/server/container_stop_test.go
index 31db1bab8..6c00a7964 100644
--- a/pkg/server/container_stop_test.go
+++ b/pkg/server/container_stop_test.go
@@ -33,9 +33,10 @@ import (
 
 func TestWaitContainerStop(t *testing.T) {
 	id := "test-id"
-	timeout := 2 * stopCheckPollInterval
 	for desc, test := range map[string]struct {
 		metadata  *metadata.ContainerMetadata
+		cancel    bool
+		timeout   time.Duration
 		expectErr bool
 	}{
 		"should return error if timeout exceeds": {
@@ -44,10 +45,22 @@ func TestWaitContainerStop(t *testing.T) {
 				CreatedAt: time.Now().UnixNano(),
 				StartedAt: time.Now().UnixNano(),
 			},
+			timeout:   2 * stopCheckPollInterval,
+			expectErr: true,
+		},
+		"should return error if context is cancelled": {
+			metadata: &metadata.ContainerMetadata{
+				ID:        id,
+				CreatedAt: time.Now().UnixNano(),
+				StartedAt: time.Now().UnixNano(),
+			},
+			timeout:   time.Hour,
+			cancel:    true,
 			expectErr: true,
 		},
 		"should not return error if container is removed before timeout": {
 			metadata:  nil,
+			timeout:   time.Hour,
 			expectErr: false,
 		},
 		"should not return error if container is stopped before timeout": {
@@ -57,6 +70,7 @@ func TestWaitContainerStop(t *testing.T) {
 				StartedAt:  time.Now().UnixNano(),
 				FinishedAt: time.Now().UnixNano(),
 			},
+			timeout:   time.Hour,
 			expectErr: false,
 		},
 	} {
@@ -64,7 +78,13 @@ func TestWaitContainerStop(t *testing.T) {
 		if test.metadata != nil {
 			assert.NoError(t, c.containerStore.Create(*test.metadata))
 		}
-		err := c.waitContainerStop(id, timeout)
+		ctx := context.Background()
+		if test.cancel {
+			cancelledCtx, cancel := context.WithCancel(ctx)
+			cancel()
+			ctx = cancelledCtx
+		}
+		err := c.waitContainerStop(ctx, id, test.timeout)
 		assert.Equal(t, test.expectErr, err != nil, desc)
 	}
 }