shim: handle connection-closed errors during kill after live migration

Protobuild.toml

@@ -3,11 +3,9 @@ generators = ["go", "go-grpc"]
 
 # Control protoc include paths.
 [includes]
-  before = ["./protobuf"]
-
   # defaults are "/usr/local/include" and "/usr/include", which don't exist on Windows.
   # override defaults to supress errors about non-existant directories.
-  after = []
+  # after = []
 
 # This section maps protobuf imports to Go packages.
 [packages]
@@ -21,5 +19,20 @@ prefixes = [
   "github.com/Microsoft/hcsshim/internal/computeagent",
   "github.com/Microsoft/hcsshim/internal/ncproxyttrpc",
   "github.com/Microsoft/hcsshim/internal/vmservice",
+  "github.com/Microsoft/hcsshim/internal/save",
+]
+generators = ["go", "go-ttrpc"]
+
+[[overrides]]
+prefixes = [
+  "github.com/Microsoft/hcsshim/internal/lm/proto",
 ]
 generators = ["go", "go-ttrpc"]
+# Override the go package path for the generated code. This allows us to not need a go_package
+# directive in the proto file itself, which is nice so we can share it with other projects easily.
+parameters.go."Mgithub.com/Microsoft/hcsshim/internal/lm/proto/lm.proto" = '"github.com/Microsoft/hcsshim/internal/lm/proto;lmproto"'
+parameters.go-ttrpc."Mgithub.com/Microsoft/hcsshim/internal/lm/proto/lm.proto" = '"github.com/Microsoft/hcsshim/internal/lm/proto;lmproto"'
+
+[[overrides]]
+prefixes = ["github.com/Microsoft/hcsshim/internal/state"]
+generators = ["go"]

cmd/containerd-shim-runhcs-v1/main.go

@@ -3,7 +3,6 @@
 package main
 
 import (
-	"context"
 	"errors"
 	"fmt"
 	"os"
@@ -20,7 +19,6 @@ import (
 
 	"github.com/Microsoft/hcsshim/internal/log"
 	"github.com/Microsoft/hcsshim/internal/oc"
-	"github.com/Microsoft/hcsshim/internal/shimdiag"
 	hcsversion "github.com/Microsoft/hcsshim/internal/version"
 
 	// register common types spec with typeurl
@@ -55,17 +53,17 @@ var (
 )
 
 func etwCallback(sourceID guid.GUID, state etw.ProviderState, level etw.Level, matchAnyKeyword uint64, matchAllKeyword uint64, filterData uintptr) {
-	if state == etw.ProviderStateCaptureState {
-		resp, err := svc.DiagStacks(context.Background(), &shimdiag.StacksRequest{})
-		if err != nil {
-			return
-		}
-		log := logrus.WithField("tid", svc.tid)
-		log.WithField("stack", resp.Stacks).Info("goroutine stack dump")
-		if resp.GuestStacks != "" {
-			log.WithField("stack", resp.GuestStacks).Info("guest stack dump")
-		}
-	}
+	// if state == etw.ProviderStateCaptureState {
+	// 	resp, err := svc.DiagStacks(context.Background(), &shimdiag.StacksRequest{})
+	// 	if err != nil {
+	// 		return
+	// 	}
+	// 	log := logrus.WithField("tid", svc.tid)
+	// 	log.WithField("stack", resp.Stacks).Info("goroutine stack dump")
+	// 	if resp.GuestStacks != "" {
+	// 		log.WithField("stack", resp.GuestStacks).Info("guest stack dump")
+	// 	}
+	// }
 }
 
 func main() {

cmd/containerd-shim-runhcs-v1/options/runhcs.proto

@@ -16,6 +16,7 @@ message Options {
 		NPIPE = 0;
 		FILE = 1;
 		ETW = 2;
+		STDIO = 3;
 	}
 
 	// debug tracing output type
@@ -107,6 +108,15 @@ message Options {
 
 	// scrub_logs enables removing environment variables and other potentially sensitive information from logs
 	bool scrub_logs = 20;
+
+	BundleType bundle_type = 22;
+}
+
+enum BundleType {
+	BUNDLE_OCI = 0;
+	BUNDLE_POD_RESTORE = 1;
+	BUNDLE_CONTAINER_RESTORE = 2;
+	BUNDLE_POD_LM = 3;
 }
 
 // ProcessDetails contains additional information about a process. This is the additional

cmd/containerd-shim-runhcs-v1/serve.go

@@ -24,13 +24,16 @@ import (
 	"google.golang.org/protobuf/types/known/anypb"
 
 	runhcsopts "github.com/Microsoft/hcsshim/cmd/containerd-shim-runhcs-v1/options"
-	"github.com/Microsoft/hcsshim/internal/extendedtask"
+	"github.com/Microsoft/hcsshim/internal/ctrdpub"
+	lmproto "github.com/Microsoft/hcsshim/internal/lm/proto"
 	hcslog "github.com/Microsoft/hcsshim/internal/log"
 	"github.com/Microsoft/hcsshim/internal/shimdiag"
+	"github.com/Microsoft/hcsshim/internal/taskserver"
+	"github.com/Microsoft/hcsshim/internal/winapi"
 	"github.com/Microsoft/hcsshim/pkg/octtrpc"
 )
 
-var svc *service
+var svc task.TaskService
 
 var serveCommand = cli.Command{
 	Name:           "serve",
@@ -47,6 +50,17 @@ var serveCommand = cli.Command{
 		},
 	},
 	Action: func(ctx *cli.Context) error {
+
+		_, err := os.Stat(`c:\debugwait`)
+		if err == nil {
+			for {
+				if winapi.IsDebuggerPresent() {
+					break
+				}
+				time.Sleep(1 * time.Second)
+			}
+		}
+
 		// On Windows the serve command is internally used to actually create
 		// the process that hosts the containerd/ttrpc entrypoint to the Runtime
 		// V2 API's. The model requires this 2nd invocation of the shim process
@@ -77,17 +91,15 @@ var serveCommand = cli.Command{
 		}
 
 		// containerd passes the shim options protobuf via stdin.
+		logrus.Info("reading options")
 		newShimOpts, err := readOptions(os.Stdin)
 		if err != nil {
 			return errors.Wrap(err, "failed to read shim options from stdin")
 		} else if newShimOpts != nil {
 			// We received a valid shim options struct.
 			shimOpts = newShimOpts
 		}
-
-		if shimOpts.Debug && shimOpts.LogLevel != "" {
-			logrus.Warning("Both Debug and LogLevel specified, Debug will be overridden")
-		}
+		logrus.Info("done reading options")
 
 		// For now keep supporting the debug option, this used to be the only way to specify a different logging
 		// level for the shim.
@@ -156,6 +168,8 @@ var serveCommand = cli.Command{
 		case runhcsopts.Options_ETW:
 			logrus.SetFormatter(hcslog.NopFormatter{})
 			logrus.SetOutput(io.Discard)
+		case runhcsopts.Options_STDIO:
+			logrus.Info("sending logs to stdio")
 		}
 
 		os.Stdin.Close()
@@ -175,20 +189,27 @@ var serveCommand = cli.Command{
 		}
 
 		ttrpcAddress := os.Getenv(ttrpcAddressEnv)
-		ttrpcEventPublisher, err := newEventPublisher(ttrpcAddress, namespaceFlag)
+		// ttrpcEventPublisher, err := newEventPublisher(ttrpcAddress, namespaceFlag)
+		// if err != nil {
+		// 	return err
+		// }
+		// defer func() {
+		// 	if err != nil {
+		// 		ttrpcEventPublisher.close()
+		// 	}
+		// }()
+		publisher, err := ctrdpub.NewPublisher(ttrpcAddress, namespaceFlag)
 		if err != nil {
 			return err
 		}
-		defer func() {
-			if err != nil {
-				ttrpcEventPublisher.close()
-			}
-		}()
 
 		// Setup the ttrpc server
-		svc, err = NewService(WithEventPublisher(ttrpcEventPublisher),
-			WithTID(idFlag),
-			WithIsSandbox(ctx.Bool("is-sandbox")))
+		closeCh := make(chan struct{})
+		coreSvc := taskserver.NewService(idFlag, closeCh, publisher)
+		taskSvc := taskserver.NewInstrumentedService(coreSvc)
+		// svc, err = NewService(WithEventPublisher(ttrpcEventPublisher),
+		// 	WithTID(idFlag),
+		// 	WithIsSandbox(ctx.Bool("is-sandbox")))
 		if err != nil {
 			return fmt.Errorf("failed to create new service: %w", err)
 		}
@@ -198,16 +219,19 @@ var serveCommand = cli.Command{
 			return err
 		}
 		defer s.Close()
-		task.RegisterTaskService(s, svc)
-		shimdiag.RegisterShimDiagService(s, svc)
-		extendedtask.RegisterExtendedTaskService(s, svc)
+		task.RegisterTaskService(s, taskSvc)
+		shimdiag.RegisterShimDiagService(s, coreSvc)
+		lmproto.RegisterMigrationService(s, coreSvc)
+		// extendedtask.RegisterExtendedTaskService(s, svc)
+		// save.RegisterSaveService(s, svc)
 
 		sl, err := winio.ListenPipe(socket, nil)
 		if err != nil {
 			return err
 		}
 		defer sl.Close()
 
+		logrus.Info("serving")
 		serrs := make(chan error, 1)
 		defer close(serrs)
 		go func() {
@@ -246,12 +270,12 @@ var serveCommand = cli.Command{
 		select {
 		case err = <-serrs:
 			// the ttrpc server shutdown without processing a shutdown request
-		case <-svc.Done():
-			if !svc.gracefulShutdown {
-				// Return immediately, but still close ttrpc server, pipes, and spans
-				// Shouldn't need to os.Exit without clean up (ie, deferred `.Close()`s)
-				return nil
-			}
+		case <-closeCh:
+			// if !svc.gracefulShutdown {
+			// 	// Return immediately, but still close ttrpc server, pipes, and spans
+			// 	// Shouldn't need to os.Exit without clean up (ie, deferred `.Close()`s)
+			// 	return nil
+			// }
 			// currently the ttrpc shutdown is the only clean up to wait on
 			sctx, cancel := context.WithTimeout(context.Background(), gracefulShutdownTimeout)
 			defer cancel()

cmd/containerd-shim-runhcs-v1/start.go

@@ -3,6 +3,7 @@
 package main
 
 import (
+	"bytes"
 	gocontext "context"
 	"encoding/json"
 	"fmt"
@@ -12,6 +13,8 @@ import (
 	"path/filepath"
 
 	"github.com/Microsoft/go-winio"
+	"github.com/Microsoft/hcsshim/cmd/containerd-shim-runhcs-v1/options"
+	lmproto "github.com/Microsoft/hcsshim/internal/lm/proto"
 	"github.com/Microsoft/hcsshim/internal/oci"
 	"github.com/Microsoft/hcsshim/pkg/annotations"
 	task "github.com/containerd/containerd/api/runtime/task/v2"
@@ -20,6 +23,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli"
+	"google.golang.org/protobuf/proto"
 )
 
 var startCommand = cli.Command{
@@ -66,7 +70,15 @@ The start command can either start a new shim or return an address to an existin
 			return err
 		}
 
-		a, err := getSpecAnnotations(cwd)
+		stdin, err := io.ReadAll(os.Stdin)
+		if err != nil {
+			return fmt.Errorf("read stdin: %w", err)
+		}
+		opts, err := readOptions(bytes.NewBuffer(stdin))
+		if err != nil {
+			return fmt.Errorf("parse options: %w", err)
+		}
+		a, err := getSpecAnnotations(cwd, opts.BundleType)
 		if err != nil {
 			return err
 		}
@@ -145,7 +157,7 @@ The start command can either start a new shim or return an address to an existin
 				Args:   args,
 				Env:    os.Environ(),
 				Dir:    cwd,
-				Stdin:  os.Stdin,
+				Stdin:  bytes.NewBuffer(stdin),
 				Stdout: w,
 				Stderr: f,
 			}
@@ -183,21 +195,33 @@ The start command can either start a new shim or return an address to an existin
 	},
 }
 
-func getSpecAnnotations(bundlePath string) (map[string]string, error) {
-	// specAnnotations is a minimal representation for oci.Spec that we need
-	// to serve a shim.
-	type specAnnotations struct {
-		// Annotations contains arbitrary metadata for the container.
-		Annotations map[string]string `json:"annotations,omitempty"`
-	}
-	f, err := os.OpenFile(filepath.Join(bundlePath, "config.json"), os.O_RDONLY, 0)
+func getSpecAnnotations(bundlePath string, bundleType options.BundleType) (map[string]string, error) {
+	raw, err := os.ReadFile(filepath.Join(bundlePath, "config.json"))
 	if err != nil {
 		return nil, err
 	}
-	defer f.Close()
-	var spec specAnnotations
-	if err := json.NewDecoder(f).Decode(&spec); err != nil {
-		return nil, errors.Wrap(err, "failed to deserialize valid OCI spec")
+	switch bundleType {
+	case options.BundleType_BUNDLE_OCI:
+		var spec struct {
+			Annotations map[string]string `json:"annotations,omitempty"`
+		}
+		if err := json.Unmarshal(raw, &spec); err != nil {
+			return nil, fmt.Errorf("deserialize OCI spec: %w", err)
+		}
+		return spec.Annotations, nil
+	case options.BundleType_BUNDLE_POD_LM:
+		var spec lmproto.SandboxLMSpec
+		if err := (proto.UnmarshalOptions{}).Unmarshal(raw, &spec); err != nil {
+			return nil, err
+		}
+		return spec.Annotations, nil
+	case options.BundleType_BUNDLE_CONTAINER_RESTORE:
+		var spec lmproto.ContainerRestoreSpec
+		if err := (proto.UnmarshalOptions{}).Unmarshal(raw, &spec); err != nil {
+			return nil, err
+		}
+		return spec.Annotations, nil
+	default:
+		return nil, fmt.Errorf("unrecognized bundle type: %v", bundleType)
 	}
-	return spec.Annotations, nil
 }

cmd/containerd-shim-runhcs-v1/task_hcs.go

@@ -211,7 +211,7 @@ func newHcsTask(
 	// Default to an infinite timeout (zero value)
 	var ioRetryTimeout time.Duration
 	if shimOpts != nil {
-		ioRetryTimeout = time.Duration(shimOpts.IoRetryTimeoutInSec) * time.Second
+		ioRetryTimeout = time.Duration(shimOpts.IORetryTimeoutInSec) * time.Second
 	}
 	io, err := cmd.NewUpstreamIO(ctx, req.ID, req.Stdout, req.Stderr, req.Stdin, req.Terminal, ioRetryTimeout)
 	if err != nil {