1 files changed, 12 insertions, 6 deletions
diff --git a/vendor/github.com/Microsoft/hcsshim/internal/hcs/waithelper.go b/vendor/github.com/Microsoft/hcsshim/internal/hcs/waithelper.go
index 91e212c57..f07f532c1 100644
--- a/vendor/github.com/Microsoft/hcsshim/internal/hcs/waithelper.go
+++ b/vendor/github.com/Microsoft/hcsshim/internal/hcs/waithelper.go
@@ -1,28 +1,34 @@
 package hcs
 
 import (
+	"context"
 	"time"
 
-	"github.com/sirupsen/logrus"
+	"github.com/Microsoft/hcsshim/internal/log"
 )
 
-func processAsyncHcsResult(err error, resultp *uint16, callbackNumber uintptr, expectedNotification hcsNotification, timeout *time.Duration) ([]ErrorEvent, error) {
-	events := processHcsResult(resultp)
+func processAsyncHcsResult(ctx context.Context, err error, resultJSON string, callbackNumber uintptr, expectedNotification hcsNotification, timeout *time.Duration) ([]ErrorEvent, error) {
+	events := processHcsResult(ctx, resultJSON)
 	if IsPending(err) {
-		return nil, waitForNotification(callbackNumber, expectedNotification, timeout)
+		return nil, waitForNotification(ctx, callbackNumber, expectedNotification, timeout)
 	}
 
 	return events, err
 }
 
-func waitForNotification(callbackNumber uintptr, expectedNotification hcsNotification, timeout *time.Duration) error {
+func waitForNotification(ctx context.Context, callbackNumber uintptr, expectedNotification hcsNotification, timeout *time.Duration) error {
 	callbackMapLock.RLock()
+	if _, ok := callbackMap[callbackNumber]; !ok {
+		callbackMapLock.RUnlock()
+		log.G(ctx).WithField("callbackNumber", callbackNumber).Error("failed to waitForNotification: callbackNumber does not exist in callbackMap")
+		return ErrHandleClose
+	}
 	channels := callbackMap[callbackNumber].channels
 	callbackMapLock.RUnlock()
 
 	expectedChannel := channels[expectedNotification]
 	if expectedChannel == nil {
-		logrus.Errorf("unknown notification type in waitForNotification %x", expectedNotification)
+		log.G(ctx).WithField("type", expectedNotification).Error("unknown notification type in waitForNotification")
 		return ErrInvalidNotificationType
 	}