Add restore support to runsc shim.

pkg/shim/extension/BUILD

@@ -9,5 +9,8 @@ go_library(
     name = "extension",
     srcs = ["extension.go"],
     visibility = ["//visibility:public"],
-    deps = ["@com_github_containerd_containerd//runtime/v2/task:go_default_library"],
+    deps = [
+        "@com_github_containerd_containerd//pkg/process:go_default_library",
+        "@com_github_containerd_containerd//runtime/v2/task:go_default_library",
+    ],
 )

pkg/shim/extension/extension.go

@@ -18,16 +18,37 @@ package extension
 import (
 	"context"
 
+	"github.com/containerd/containerd/pkg/process"
 	"github.com/containerd/containerd/runtime/v2/task"
 )
 
 // NewExtension registers an extension constructor. It may return nil, nil to indicate that the
 // extension should not handle this task request. Returning an error will fail the task request.
 var NewExtension func(ctx context.Context, next TaskServiceExt, req *task.CreateTaskRequest) (TaskServiceExt, error)
 
+// RestoreRequest is a request to restore a container. It extends
+// task.StartRequest with restore functionality.
+type RestoreRequest struct {
+	Start task.StartRequest
+	Conf  RestoreConfig
+}
+
+// Process extends process.Process with extra restore functionality.
+type Process interface {
+	process.Process
+	// Restore restores the container from a snapshot.
+	Restore(context.Context, *RestoreConfig) error
+}
+
+// RestoreConfig is the configuration for a restore request.
+type RestoreConfig struct {
+	ImagePath string
+	Direct    bool
+}
+
 // TaskServiceExt extends TaskRequest with extra functionality required by the shim.
 type TaskServiceExt interface {
 	task.TaskService
 	Cleanup(ctx context.Context) (*task.DeleteResponse, error)
-	Restore(ctx context.Context, req *task.StartRequest) (*task.StartResponse, error)
+	Restore(ctx context.Context, req *RestoreRequest) (*task.StartResponse, error)
 }

pkg/shim/proc/BUILD

@@ -25,13 +25,13 @@ go_library(
     deps = [
         "//pkg/atomicbitops",
         "//pkg/cleanup",
+        "//pkg/shim/extension",
         "//pkg/shim/runsccmd",
         "//pkg/shim/utils",
         "@com_github_containerd_console//:go_default_library",
         "@com_github_containerd_containerd//errdefs:go_default_library",
         "@com_github_containerd_containerd//log:go_default_library",
         "@com_github_containerd_containerd//mount:go_default_library",
-        "@com_github_containerd_containerd//pkg/process:go_default_library",
         "@com_github_containerd_containerd//pkg/stdio:go_default_library",
         "@com_github_containerd_fifo//:go_default_library",
         "@com_github_containerd_go_runc//:go_default_library",

pkg/shim/proc/deleted_state.go

@@ -21,8 +21,8 @@ import (
 
 	"github.com/containerd/console"
 	"github.com/containerd/containerd/errdefs"
-	"github.com/containerd/containerd/pkg/process"
 	runc "github.com/containerd/go-runc"
+	"gvisor.dev/gvisor/pkg/shim/extension"
 )
 
 type deletedState struct{}
@@ -31,7 +31,7 @@ func (*deletedState) Resize(console.WinSize) error {
 	return fmt.Errorf("cannot resize a deleted container/process")
 }
 
-func (*deletedState) Start(context.Context) error {
+func (*deletedState) Start(context.Context, *extension.RestoreConfig) error {
 	return fmt.Errorf("cannot start a deleted container/process")
 }
 
@@ -45,7 +45,7 @@ func (*deletedState) Kill(_ context.Context, signal uint32, _ bool) error {
 
 func (*deletedState) SetExited(int) {}
 
-func (*deletedState) Exec(context.Context, string, *ExecConfig) (process.Process, error) {
+func (*deletedState) Exec(context.Context, string, *ExecConfig) (extension.Process, error) {
 	return nil, fmt.Errorf("cannot exec in a deleted state")
 }
 

pkg/shim/proc/exec.go

@@ -33,6 +33,7 @@ import (
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"golang.org/x/sys/unix"
 	"gvisor.dev/gvisor/pkg/cleanup"
+	"gvisor.dev/gvisor/pkg/shim/extension"
 	"gvisor.dev/gvisor/pkg/shim/runsccmd"
 )
 
@@ -168,7 +169,7 @@ func (e *execProcess) Start(ctx context.Context) error {
 	e.mu.Lock()
 	defer e.mu.Unlock()
 
-	return e.execState.Start(ctx)
+	return e.execState.Start(ctx, nil /* restoreConf */)
 }
 
 func (e *execProcess) start(ctx context.Context) error {
@@ -275,6 +276,10 @@ func (e *execProcess) start(ctx context.Context) error {
 	return nil
 }
 
+func (e *execProcess) Restore(context.Context, *extension.RestoreConfig) error {
+	return fmt.Errorf("cannot restore an exec'd process")
+}
+
 func (e *execProcess) Status(context.Context) (string, error) {
 	e.mu.Lock()
 	defer e.mu.Unlock()

pkg/shim/proc/exec_state.go

@@ -20,11 +20,12 @@ import (
 	"fmt"
 
 	"github.com/containerd/console"
+	"gvisor.dev/gvisor/pkg/shim/extension"
 )
 
 type execState interface {
 	Resize(console.WinSize) error
-	Start(context.Context) error
+	Start(context.Context, *extension.RestoreConfig) error
 	Delete(context.Context) error
 	Kill(context.Context, uint32, bool) error
 	SetExited(int)
@@ -55,7 +56,10 @@ func (s *execCreatedState) Resize(ws console.WinSize) error {
 	return s.p.resize(ws)
 }
 
-func (s *execCreatedState) Start(ctx context.Context) error {
+func (s *execCreatedState) Start(ctx context.Context, restoreConf *extension.RestoreConfig) error {
+	if restoreConf != nil {
+		return fmt.Errorf("cannot restore an exec'd process")
+	}
 	if err := s.p.start(ctx); err != nil {
 		return err
 	}
@@ -99,7 +103,7 @@ func (s *execRunningState) Resize(ws console.WinSize) error {
 	return s.p.resize(ws)
 }
 
-func (s *execRunningState) Start(context.Context) error {
+func (s *execRunningState) Start(context.Context, *extension.RestoreConfig) error {
 	return fmt.Errorf("cannot start a running process")
 }
 
@@ -137,7 +141,7 @@ func (s *execStoppedState) Resize(console.WinSize) error {
 	return fmt.Errorf("cannot resize a stopped container")
 }
 
-func (s *execStoppedState) Start(context.Context) error {
+func (s *execStoppedState) Start(context.Context, *extension.RestoreConfig) error {
 	return fmt.Errorf("cannot start a stopped process")
 }
 

pkg/shim/proc/init.go

@@ -30,13 +30,13 @@ import (
 	"github.com/containerd/containerd/errdefs"
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/mount"
-	"github.com/containerd/containerd/pkg/process"
 	"github.com/containerd/containerd/pkg/stdio"
 
 	"github.com/containerd/fifo"
 	runc "github.com/containerd/go-runc"
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"golang.org/x/sys/unix"
+	"gvisor.dev/gvisor/pkg/shim/extension"
 	"gvisor.dev/gvisor/pkg/shim/runsccmd"
 	"gvisor.dev/gvisor/pkg/shim/utils"
 )
@@ -226,16 +226,25 @@ func (p *Init) Start(ctx context.Context) error {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
-	return p.initState.Start(ctx)
+	return p.initState.Start(ctx, nil /* restoreConf */)
 }
 
-func (p *Init) start(ctx context.Context) error {
+func (p *Init) start(ctx context.Context, restoreConf *extension.RestoreConfig) error {
 	var cio runc.IO
 	if !p.Sandbox {
 		cio = p.io
 	}
-	if err := p.runtime.Start(ctx, p.id, cio); err != nil {
-		return p.runtimeError(err, "OCI runtime start failed")
+	if restoreConf == nil {
+		if err := p.runtime.Start(ctx, p.id, cio); err != nil {
+			return p.runtimeError(err, "OCI runtime start failed")
+		}
+	} else {
+		if err := p.runtime.Restore(ctx, p.id, cio, &runsccmd.RestoreOpts{
+			ImagePath: restoreConf.ImagePath,
+			Direct:    restoreConf.Direct,
+		}); err != nil {
+			return p.runtimeError(err, "OCI runtime restore failed")
+		}
 	}
 	go func() {
 		status, err := p.runtime.Wait(context.Background(), p.id)
@@ -253,7 +262,15 @@ func (p *Init) start(ctx context.Context) error {
 	return nil
 }
 
-// SetExited set the exit stauts of the init process.
+// Restore restores the container from a snapshot.
+func (p *Init) Restore(ctx context.Context, conf *extension.RestoreConfig) error {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	return p.initState.Start(ctx, conf)
+}
+
+// SetExited set the exit status of the init process.
 func (p *Init) SetExited(status int) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
@@ -392,15 +409,15 @@ func (p *Init) Runtime() *runsccmd.Runsc {
 }
 
 // Exec returns a new child process.
-func (p *Init) Exec(ctx context.Context, path string, r *ExecConfig) (process.Process, error) {
+func (p *Init) Exec(ctx context.Context, path string, r *ExecConfig) (extension.Process, error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
 
 	return p.initState.Exec(ctx, path, r)
 }
 
 // exec returns a new exec'd process.
-func (p *Init) exec(path string, r *ExecConfig) (process.Process, error) {
+func (p *Init) exec(path string, r *ExecConfig) (extension.Process, error) {
 	var spec specs.Process
 	if err := json.Unmarshal(r.Spec.Value, &spec); err != nil {
 		return nil, err

pkg/shim/proc/init_state.go

@@ -20,9 +20,10 @@ import (
 	"fmt"
 
 	"github.com/containerd/containerd/errdefs"
-	"github.com/containerd/containerd/pkg/process"
 	runc "github.com/containerd/go-runc"
 	"golang.org/x/sys/unix"
+
+	"gvisor.dev/gvisor/pkg/shim/extension"
 )
 
 type stateTransition int
@@ -47,9 +48,11 @@ func (s stateTransition) String() string {
 }
 
 type initState interface {
-	Start(context.Context) error
+	// Start starts the process. If RestoreConfig is provided, the process is
+	// restored using the checkpoint image provided in the config.
+	Start(context.Context, *extension.RestoreConfig) error
 	Delete(context.Context) error
-	Exec(context.Context, string, *ExecConfig) (process.Process, error)
+	Exec(context.Context, string, *ExecConfig) (extension.Process, error)
 	State(ctx context.Context) (string, error)
 	Stats(context.Context, string) (*runc.Stats, error)
 	Kill(context.Context, uint32, bool) error
@@ -77,14 +80,14 @@ func (s *createdState) transition(transition stateTransition) {
 	}
 }
 
-func (s *createdState) Start(ctx context.Context) error {
-	if err := s.p.start(ctx); err != nil {
+func (s *createdState) Start(ctx context.Context, restoreConf *extension.RestoreConfig) error {
+	if err := s.p.start(ctx, restoreConf); err != nil {
 		// Containerd doesn't allow deleting container in created state.
-		// However, for gvisor, a non-root container in created state can
-		// only go to running state. If the container can't be started,
+		// However, for gVisor, a non-root container in created state can
+		// only go to running state. If the container can't be started/restored,
 		// it can only stay in created state, and never be deleted.
-		// To work around that, we treat non-root container in start failure
-		// state as stopped.
+		// To work around that, we treat non-root container in start/restore
+		// failure state as stopped.
 		if !s.p.Sandbox {
 			s.p.io.Close()
 			s.p.setExited(internalErrorCode)
@@ -113,7 +116,7 @@ func (s *createdState) SetExited(status int) {
 	s.transition(stopped)
 }
 
-func (s *createdState) Exec(ctx context.Context, path string, r *ExecConfig) (process.Process, error) {
+func (s *createdState) Exec(ctx context.Context, path string, r *ExecConfig) (extension.Process, error) {
 	return s.p.exec(path, r)
 }
 
@@ -146,7 +149,7 @@ func (s *runningState) transition(transition stateTransition) {
 	}
 }
 
-func (s *runningState) Start(ctx context.Context) error {
+func (s *runningState) Start(context.Context, *extension.RestoreConfig) error {
 	return fmt.Errorf("cannot start a running container")
 }
 
@@ -163,7 +166,7 @@ func (s *runningState) SetExited(status int) {
 	s.transition(stopped)
 }
 
-func (s *runningState) Exec(_ context.Context, path string, r *ExecConfig) (process.Process, error) {
+func (s *runningState) Exec(_ context.Context, path string, r *ExecConfig) (extension.Process, error) {
 	return s.p.exec(path, r)
 }
 
@@ -196,7 +199,7 @@ func (s *stoppedState) transition(transition stateTransition) {
 	}
 }
 
-func (s *stoppedState) Start(context.Context) error {
+func (s *stoppedState) Start(context.Context, *extension.RestoreConfig) error {
 	return fmt.Errorf("cannot start a stopped container")
 }
 
@@ -216,7 +219,7 @@ func (s *stoppedState) SetExited(status int) {
 	s.process.setExited(status)
 }
 
-func (s *stoppedState) Exec(context.Context, string, *ExecConfig) (process.Process, error) {
+func (s *stoppedState) Exec(context.Context, string, *ExecConfig) (extension.Process, error) {
 	return nil, fmt.Errorf("cannot exec in a stopped state")
 }
 

pkg/shim/runsc/service.go

@@ -111,7 +111,7 @@ type runscService struct {
 	task *proc.Init
 
 	// processes maps ExecId to processes running through exec.
-	processes map[string]process.Process
+	processes map[string]extension.Process
 
 	events chan any
 
@@ -148,7 +148,7 @@ func New(ctx context.Context, id string, publisher shim.Publisher) (extension.Ta
 	go ep.run(ctx)
 	s := &runscService{
 		id:        id,
-		processes: make(map[string]process.Process),
+		processes: make(map[string]extension.Process),
 		events:    make(chan any, 128),
 		ec:        proc.ExitCh,
 		oomPoller: ep,
@@ -611,9 +611,19 @@ func (s *runscService) Checkpoint(ctx context.Context, r *taskAPI.CheckpointTask
 }
 
 // Restore restores the container.
-func (s *runscService) Restore(ctx context.Context, r *taskAPI.StartRequest) (*taskAPI.StartResponse, error) {
-	log.L.Debugf("Restore, id: %s", r.ID)
-	return nil, errdefs.ErrNotImplemented
+func (s *runscService) Restore(ctx context.Context, r *extension.RestoreRequest) (*taskAPI.StartResponse, error) {
+	p, err := s.getProcess(r.Start.ExecID)
+	if err != nil {
+		return nil, err
+	}
+	if err := p.Restore(ctx, &r.Conf); err != nil {
+		return nil, err
+	}
+	// TODO: Set the cgroup and oom notifications on restore.
+	// https://github.com/google/gvisor-containerd-shim/issues/58
+	return &taskAPI.StartResponse{
+		Pid: uint32(p.Pid()),
+	}, nil
 }
 
 // Connect returns shim information such as the shim's pid.
@@ -798,7 +808,7 @@ func (s *runscService) forward(ctx context.Context, publisher shim.Publisher) {
 	}
 }
 
-func (s *runscService) getProcess(execID string) (process.Process, error) {
+func (s *runscService) getProcess(execID string) (extension.Process, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()