Add high level job object wrapper

internal/jobobject/iocp.go

@@ -0,0 +1,111 @@
+package jobobject
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"unsafe"
+
+	"github.com/Microsoft/hcsshim/internal/log"
+	"github.com/Microsoft/hcsshim/internal/queue"
+	"github.com/Microsoft/hcsshim/internal/winapi"
+	"github.com/sirupsen/logrus"
+	"golang.org/x/sys/windows"
+)
+
+var (
+	ioInitOnce sync.Once
+	initIOErr  error
+	// Global iocp handle that will be re-used for every job object
+	ioCompletionPort windows.Handle
+	// Mapping of job handle to queue to place notifications in.
+	jobMap sync.Map
+)
+
+// MsgAllProcessesExited is a type representing a message that every process in a job has exited.
+type MsgAllProcessesExited struct{}
+
+// MsgUnimplemented represents a message that we are aware of, but that isn't implemented currently.
+// This should not be treated as an error.
+type MsgUnimplemented struct{}
+
+// pollIOCP polls the io completion port forever.
+func pollIOCP(ctx context.Context, iocpHandle windows.Handle) {
+	var (
+		overlapped uintptr
+		code       uint32
+		key        uintptr
+	)
+
+	for {
+		err := winapi.GetQueuedCompletionStatus(iocpHandle, &code, &key, (**windows.Overlapped)(unsafe.Pointer(&overlapped)), windows.INFINITE)
+		if err != nil {
+			log.G(ctx).WithError(err).Error("failed to poll for job object message")
+			continue
+		}
+		if val, ok := jobMap.Load(key); ok {
+			msq, ok := val.(*queue.MessageQueue)
+			if !ok {
+				log.G(ctx).WithField("value", msq).Warn("encountered non queue type in job map")
+				continue
+			}
+			notification, err := parseMessage(code, overlapped)
+			if err != nil {
+				log.G(ctx).WithFields(logrus.Fields{
+					"code":       code,
+					"overlapped": overlapped,
+				}).Warn("failed to parse job object message")
+				continue
+			}
+			if err := msq.Write(notification); err == queue.ErrQueueClosed {
+				// Write will only return an error when the queue is closed.
+				// The only time a queue would ever be closed is when we call `Close` on
+				// the job it belongs to which also removes it from the jobMap, so something
+				// went wrong here. We can't return as this is reading messages for all jobs
+				// so just log it and move on.
+				log.G(ctx).WithFields(logrus.Fields{
+					"code":       code,
+					"overlapped": overlapped,
+				}).Warn("tried to write to a closed queue")
+				continue
+			}
+		} else {
+			log.G(ctx).Warn("received a message for a job not present in the mapping")
+		}
+	}
+}
+
+func parseMessage(code uint32, overlapped uintptr) (interface{}, error) {
+	// Check code and parse out relevant information related to that notification
+	// that we care about. For now all we handle is the message that all processes
+	// in the job have exited.
+	switch code {
+	case winapi.JOB_OBJECT_MSG_ACTIVE_PROCESS_ZERO:
+		return MsgAllProcessesExited{}, nil
+	// Other messages for completeness and a check to make sure that if we fall
+	// into the default case that this is a code we don't know how to handle.
+	case winapi.JOB_OBJECT_MSG_END_OF_JOB_TIME:
+	case winapi.JOB_OBJECT_MSG_END_OF_PROCESS_TIME:
+	case winapi.JOB_OBJECT_MSG_ACTIVE_PROCESS_LIMIT:
+	case winapi.JOB_OBJECT_MSG_NEW_PROCESS:
+	case winapi.JOB_OBJECT_MSG_EXIT_PROCESS:
+	case winapi.JOB_OBJECT_MSG_ABNORMAL_EXIT_PROCESS:
+	case winapi.JOB_OBJECT_MSG_PROCESS_MEMORY_LIMIT:
+	case winapi.JOB_OBJECT_MSG_JOB_MEMORY_LIMIT:
+	case winapi.JOB_OBJECT_MSG_NOTIFICATION_LIMIT:
+	default:
+		return nil, fmt.Errorf("unknown job notification type: %d", code)
+	}
+	return MsgUnimplemented{}, nil
+}
+
+// Assigns an IO completion port to get notified of events for the registered job
+// object.
+func attachIOCP(job windows.Handle, iocp windows.Handle) error {
+	info := winapi.JOBOBJECT_ASSOCIATE_COMPLETION_PORT{
+		CompletionKey:  job,
+		CompletionPort: iocp,
+	}
+	_, err := windows.SetInformationJobObject(job, windows.JobObjectAssociateCompletionPortInformation, uintptr(unsafe.Pointer(&info)), uint32(unsafe.Sizeof(info)))
+	return err
+}