From f10695300a166d467f397314aa2381b431eb022a Mon Sep 17 00:00:00 2001
From: Brian Goff <cpuguy83@gmail.com>
Date: Wed, 16 Nov 2016 22:08:37 -0500
Subject: [PATCH] Use WaitWithContext for WaitStop

Signed-off-by: Brian Goff <cpuguy83@gmail.com>
Upstream-commit: 36d6d76a41be4973eed98f64a565f8cf92dc16e0
Component: engine
---
 components/engine/container/state.go | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/components/engine/container/state.go b/components/engine/container/state.go
index 4dd2ecec69..ff080f1bfc 100644
--- a/components/engine/container/state.go
+++ b/components/engine/container/state.go
@@ -177,26 +177,24 @@ func wait(waitChan <-chan struct{}, timeout time.Duration) error {
 // immediately. If you want wait forever you must supply negative timeout.
 // Returns exit code, that was passed to SetStopped
 func (s *State) WaitStop(timeout time.Duration) (int, error) {
-	s.Lock()
-	if !s.Running {
-		exitCode := s.ExitCodeValue
-		s.Unlock()
-		return exitCode, nil
+	ctx := context.Background()
+	if timeout >= 0 {
+		var cancel func()
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+		defer cancel()
 	}
-	waitChan := s.waitChan
-	s.Unlock()
-	if err := wait(waitChan, timeout); err != nil {
+	if err := s.WaitWithContext(ctx); err != nil {
+		if status, ok := err.(*StateStatus); ok {
+			return status.ExitCode(), nil
+		}
 		return -1, err
 	}
-	s.Lock()
-	defer s.Unlock()
-	return s.ExitCode(), nil
+	return 0, nil
 }
 
 // WaitWithContext waits for the container to stop. Optional context can be
 // passed for canceling the request.
 func (s *State) WaitWithContext(ctx context.Context) error {
-	// todo(tonistiigi): make other wait functions use this
 	s.Lock()
 	if !s.Running {
 		state := newStateStatus(s.ExitCode(), s.Error())