1 files changed, 19 insertions, 23 deletions
diff --git a/vendor/github.com/Microsoft/hcsshim/internal/hcs/callback.go b/vendor/github.com/Microsoft/hcsshim/internal/hcs/callback.go
index 6d909873c..62ba81751 100644
--- a/vendor/github.com/Microsoft/hcsshim/internal/hcs/callback.go
+++ b/vendor/github.com/Microsoft/hcsshim/internal/hcs/callback.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/Microsoft/hcsshim/internal/interop"
 	"github.com/Microsoft/hcsshim/internal/logfields"
+	"github.com/Microsoft/hcsshim/internal/vmcompute"
 	"github.com/sirupsen/logrus"
 )
 
@@ -88,7 +89,7 @@ type notificationChannel chan error
 
 type notifcationWatcherContext struct {
 	channels notificationChannels
-	handle   hcsCallback
+	handle   vmcompute.HcsCallback
 
 	systemID  string
 	processID int
@@ -98,21 +99,27 @@ type notificationChannels map[hcsNotification]notificationChannel
 
 func newSystemChannels() notificationChannels {
 	channels := make(notificationChannels)
-
-	channels[hcsNotificationSystemExited] = make(notificationChannel, 1)
-	channels[hcsNotificationSystemCreateCompleted] = make(notificationChannel, 1)
-	channels[hcsNotificationSystemStartCompleted] = make(notificationChannel, 1)
-	channels[hcsNotificationSystemPauseCompleted] = make(notificationChannel, 1)
-	channels[hcsNotificationSystemResumeCompleted] = make(notificationChannel, 1)
-
+	for _, notif := range []hcsNotification{
+		hcsNotificationServiceDisconnect,
+		hcsNotificationSystemExited,
+		hcsNotificationSystemCreateCompleted,
+		hcsNotificationSystemStartCompleted,
+		hcsNotificationSystemPauseCompleted,
+		hcsNotificationSystemResumeCompleted,
+	} {
+		channels[notif] = make(notificationChannel, 1)
+	}
 	return channels
 }
 
 func newProcessChannels() notificationChannels {
 	channels := make(notificationChannels)
-
-	channels[hcsNotificationProcessExited] = make(notificationChannel, 1)
-
+	for _, notif := range []hcsNotification{
+		hcsNotificationServiceDisconnect,
+		hcsNotificationProcessExited,
+	} {
+		channels[notif] = make(notificationChannel, 1)
+	}
 	return channels
 }
 
@@ -143,18 +150,7 @@ func notificationWatcher(notificationType hcsNotification, callbackNumber uintpt
 	if context.processID != 0 {
 		log.Data[logfields.ProcessID] = context.processID
 	}
-	log.Debug("")
-
-	// The HCS notification system can grow overtime. We explicitly opt-in to
-	// the notifications we would like to handle, all others we simply return.
-	// This means that as it grows we don't have issues associated with new
-	// notification types the code didn't know about.
-	switch notificationType {
-	case hcsNotificationSystemExited, hcsNotificationSystemCreateCompleted, hcsNotificationSystemStartCompleted, hcsNotificationSystemPauseCompleted, hcsNotificationSystemResumeCompleted:
-	case hcsNotificationProcessExited:
-	default:
-		return 0
-	}
+	log.Debug("HCS notification")
 
 	if channel, ok := context.channels[notificationType]; ok {
 		channel <- result