Merge pull request #44 from Random-Liu/stop-on-cancellation

Stop waiting on cancellation.

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:

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)
 	}
 }