Merge pull request #42 from Microsoft/NotificationWaiter

New notification design allows for multiple waiters

Author: darstahl

callback.go

@@ -40,39 +40,36 @@ type hcsNotification uint32
 type notificationChannel chan error
 
 type notifcationWatcherContext struct {
-	channel              notificationChannel
-	expectedNotification hcsNotification
-	handle               hcsCallback
+	channels notificationChannels
+	handle   hcsCallback
+}
+
+type notificationChannels map[hcsNotification]notificationChannel
+
+func newChannels() 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)
+	channels[hcsNotificationProcessExited] = make(notificationChannel, 1)
+	channels[hcsNotificationServiceDisconnect] = make(notificationChannel, 1)
+	return channels
 }
 
 func notificationWatcher(notificationType hcsNotification, callbackNumber uintptr, notificationStatus uintptr, notificationData *uint16) uintptr {
-	var (
-		result       error
-		completeWait = false
-	)
+	var result error
+	if int32(notificationStatus) < 0 {
+		result = syscall.Errno(win32FromHresult(notificationStatus))
+	}
 
 	callbackMapLock.RLock()
-	context := callbackMap[callbackNumber]
+	channels := callbackMap[callbackNumber].channels
 	callbackMapLock.RUnlock()
 
-	if notificationType == context.expectedNotification {
-		if int32(notificationStatus) < 0 {
-			result = syscall.Errno(win32FromHresult(notificationStatus))
-		} else {
-			result = nil
-		}
-		completeWait = true
-	} else if notificationType == hcsNotificationSystemExited {
-		result = ErrUnexpectedContainerExit
-		completeWait = true
-	} else if notificationType == hcsNotificationServiceDisconnect {
-		result = ErrUnexpectedProcessAbort
-		completeWait = true
-	}
-
-	if completeWait {
-		context.channel <- result
-	}
+	channels[notificationType] <- result
 
 	return 0
 }

container.go

@@ -26,8 +26,9 @@ type ContainerError struct {
 }
 
 type container struct {
-	handle hcsSystem
-	id     string
+	handle         hcsSystem
+	id             string
+	callbackNumber uintptr
 }
 
 type containerProperties struct {
@@ -61,19 +62,25 @@ func CreateContainer(id string, c *ContainerConfig) (Container, error) {
 	logrus.Debugf(title+" id=%s config=%s", id, configuration)
 
 	var (
-		handle      hcsSystem
 		resultp     *uint16
 		createError error
 	)
 	if hcsCallbacksSupported {
 		var identity syscall.Handle
-		createError = hcsCreateComputeSystem(id, configuration, identity, &handle, &resultp)
+		createError = hcsCreateComputeSystem(id, configuration, identity, &container.handle, &resultp)
+
+		if createError == nil || createError == ErrVmcomputeOperationPending {
+			if err := container.registerCallback(); err != nil {
+				err := &ContainerError{Container: container, Operation: operation, Err: err}
+				logrus.Error(err)
+				return nil, err
+			}
+		}
 	} else {
-		createError = hcsCreateComputeSystemTP5(id, configuration, &handle, &resultp)
+		createError = hcsCreateComputeSystemTP5(id, configuration, &container.handle, &resultp)
 	}
-	container.handle = handle
 
-	err = processAsyncHcsResult(container, createError, resultp, hcsNotificationSystemCreateCompleted, &defaultTimeout)
+	err = processAsyncHcsResult(createError, resultp, container.callbackNumber, hcsNotificationSystemCreateCompleted, &defaultTimeout)
 	if err != nil {
 		err := &ContainerError{Container: container, Operation: operation, ExtraInfo: configuration, Err: err}
 		logrus.Error(err)
@@ -122,7 +129,7 @@ func (container *container) Start() error {
 
 	var resultp *uint16
 	err := hcsStartComputeSystemTP5(container.handle, nil, &resultp)
-	err = processAsyncHcsResult(container, err, resultp, hcsNotificationSystemStartCompleted, &defaultTimeout)
+	err = processAsyncHcsResult(err, resultp, container.callbackNumber, hcsNotificationSystemStartCompleted, &defaultTimeout)
 	if err != nil {
 		err := &ContainerError{Container: container, Operation: operation, Err: err}
 		logrus.Error(err)
@@ -186,7 +193,7 @@ func (container *container) Wait() error {
 	logrus.Debugf(title+" id=%s", container.id)
 
 	if hcsCallbacksSupported {
-		err := registerAndWaitForCallback(container, hcsNotificationSystemExited)
+		err := waitForNotification(container.callbackNumber, hcsNotificationSystemExited, nil)
 		if err != nil {
 			err := &ContainerError{Container: container, Operation: operation, Err: err}
 			logrus.Error(err)
@@ -217,7 +224,7 @@ func (container *container) WaitTimeout(timeout time.Duration) error {
 	logrus.Debugf(title+" id=%s", container.id)
 
 	if hcsCallbacksSupported {
-		err := registerAndWaitForCallbackTimeout(container, hcsNotificationSystemExited, timeout)
+		err := waitForNotification(container.callbackNumber, hcsNotificationSystemExited, &timeout)
 		if err == ErrTimeout {
 			return ErrTimeout
 		} else if err != nil {
@@ -304,7 +311,7 @@ func (container *container) Pause() error {
 
 	var resultp *uint16
 	err := hcsPauseComputeSystemTP5(container.handle, nil, &resultp)
-	err = processAsyncHcsResult(container, err, resultp, hcsNotificationSystemPauseCompleted, &defaultTimeout)
+	err = processAsyncHcsResult(err, resultp, container.callbackNumber, hcsNotificationSystemPauseCompleted, &defaultTimeout)
 	if err != nil {
 		err := &ContainerError{Container: container, Operation: operation, Err: err}
 		logrus.Error(err)
@@ -325,7 +332,7 @@ func (container *container) Resume() error {
 	)
 
 	err := hcsResumeComputeSystemTP5(container.handle, nil, &resultp)
-	err = processAsyncHcsResult(container, err, resultp, hcsNotificationSystemResumeCompleted, &defaultTimeout)
+	err = processAsyncHcsResult(err, resultp, container.callbackNumber, hcsNotificationSystemResumeCompleted, &defaultTimeout)
 	if err != nil {
 		err := &ContainerError{Container: container, Operation: operation, Err: err}
 		logrus.Error(err)
@@ -379,6 +386,14 @@ func (container *container) CreateProcess(c *ProcessConfig) (Process, error) {
 		},
 	}
 
+	if hcsCallbacksSupported {
+		if err := process.registerCallback(); err != nil {
+			err = &ContainerError{Container: container, Operation: operation, Err: err}
+			logrus.Error(err)
+			return nil, err
+		}
+	}
+
 	logrus.Debugf(title+" succeeded id=%s processid=%s", container.id, process.processID)
 	runtime.SetFinalizer(process, closeProcess)
 	return process, nil
@@ -408,6 +423,12 @@ func (container *container) OpenProcess(pid int) (Process, error) {
 		container: container,
 	}
 
+	if err := process.registerCallback(); err != nil {
+		err = &ContainerError{Container: container, Operation: operation, Err: err}
+		logrus.Error(err)
+		return nil, err
+	}
+
 	logrus.Debugf(title+" succeeded id=%s processid=%s", container.id, process.processID)
 	runtime.SetFinalizer(process, closeProcess)
 	return process, nil
@@ -424,6 +445,14 @@ func (container *container) Close() error {
 		return nil
 	}
 
+	if hcsCallbacksSupported {
+		if err := container.unregisterCallback(); err != nil {
+			err = &ContainerError{Container: container, Operation: operation, Err: err}
+			logrus.Error(err)
+			return err
+		}
+	}
+
 	if err := hcsCloseComputeSystem(container.handle); err != nil {
 		err = &ContainerError{Container: container, Operation: operation, Err: err}
 		logrus.Error(err)
@@ -441,30 +470,32 @@ func closeContainer(container *container) {
 	container.Close()
 }
 
-func (container *container) registerCallback(expectedNotification hcsNotification) (uintptr, error) {
+func (container *container) registerCallback() error {
 	callbackMapLock.Lock()
 	defer callbackMapLock.Unlock()
 
 	callbackNumber := nextCallback
 	nextCallback++
 
 	context := &notifcationWatcherContext{
-		expectedNotification: expectedNotification,
-		channel:              make(chan error, 1),
+		channels: newChannels(),
 	}
 	callbackMap[callbackNumber] = context
 
 	var callbackHandle hcsCallback
 	err := hcsRegisterComputeSystemCallback(container.handle, notificationWatcherCallback, callbackNumber, &callbackHandle)
 	if err != nil {
-		return 0, err
+		return err
 	}
 	context.handle = callbackHandle
+	container.callbackNumber = callbackNumber
 
-	return callbackNumber, nil
+	return nil
 }
 
-func (container *container) unregisterCallback(callbackNumber uintptr) error {
+func (container *container) unregisterCallback() error {
+	callbackNumber := container.callbackNumber
+
 	callbackMapLock.Lock()
 	defer callbackMapLock.Unlock()
 

process.go

@@ -22,11 +22,11 @@ type ProcessError struct {
 }
 
 type process struct {
-	handle             hcsProcess
-	processID          int
-	container          *container
-	cachedPipes        *cachedPipes
-	killCallbackNumber uintptr
+	handle         hcsProcess
+	processID      int
+	container      *container
+	cachedPipes    *cachedPipes
+	callbackNumber uintptr
 }
 
 type cachedPipes struct {
@@ -101,7 +101,7 @@ func (process *process) Wait() error {
 	logrus.Debugf(title+" processid=%d", process.processID)
 
 	if hcsCallbacksSupported {
-		err := registerAndWaitForCallback(process, hcsNotificationProcessExited)
+		err := waitForNotification(process.callbackNumber, hcsNotificationProcessExited, nil)
 		if err != nil {
 			err := &ProcessError{Operation: operation, Process: process, Err: err}
 			logrus.Error(err)
@@ -128,7 +128,7 @@ func (process *process) WaitTimeout(timeout time.Duration) error {
 	logrus.Debugf(title+" processid=%d", process.processID)
 
 	if hcsCallbacksSupported {
-		err := registerAndWaitForCallbackTimeout(process, hcsNotificationProcessExited, timeout)
+		err := waitForNotification(process.callbackNumber, hcsNotificationProcessExited, &timeout)
 		if err == ErrTimeout {
 			return ErrTimeout
 		} else if err != nil {
@@ -344,6 +344,14 @@ func (process *process) Close() error {
 		return nil
 	}
 
+	if hcsCallbacksSupported {
+		if err := process.unregisterCallback(); err != nil {
+			err = &ProcessError{Operation: operation, Process: process, Err: err}
+			logrus.Error(err)
+			return err
+		}
+	}
+
 	if err := hcsCloseProcess(process.handle); err != nil {
 		err = &ProcessError{Operation: operation, Process: process, Err: err}
 		logrus.Error(err)
@@ -361,30 +369,32 @@ func closeProcess(process *process) {
 	process.Close()
 }
 
-func (process *process) registerCallback(expectedNotification hcsNotification) (uintptr, error) {
+func (process *process) registerCallback() error {
 	callbackMapLock.Lock()
 	defer callbackMapLock.Unlock()
 
 	callbackNumber := nextCallback
 	nextCallback++
 
 	context := &notifcationWatcherContext{
-		expectedNotification: expectedNotification,
-		channel:              make(chan error, 1),
+		channels: newChannels(),
 	}
 	callbackMap[callbackNumber] = context
 
 	var callbackHandle hcsCallback
 	err := hcsRegisterProcessCallback(process.handle, notificationWatcherCallback, callbackNumber, &callbackHandle)
 	if err != nil {
-		return 0, err
+		return err
 	}
 	context.handle = callbackHandle
+	process.callbackNumber = callbackNumber
 
-	return callbackNumber, nil
+	return nil
 }
 
-func (process *process) unregisterCallback(callbackNumber uintptr) error {
+func (process *process) unregisterCallback() error {
+	callbackNumber := process.callbackNumber
+
 	callbackMapLock.Lock()
 	defer callbackMapLock.Unlock()
 	handle := callbackMap[callbackNumber].handle