From a9b2ce6978ca26b3035f8a012003bf4d346fad70 Mon Sep 17 00:00:00 2001
From: Kevin Parsons <kevpar@microsoft.com>
Date: Mon, 18 Nov 2024 08:55:13 -0800
Subject: [PATCH] Return restored container pid from Create

---
 internal/lm/proto/lm.proto       | 2 +-
 internal/taskserver/internal.go  | 7 +++++++
 internal/taskserver/migration.go | 3 ++-
 internal/taskserver/service.go   | 9 +++++----
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/internal/lm/proto/lm.proto b/internal/lm/proto/lm.proto
index ab4cc135bb..480680739a 100644
--- a/internal/lm/proto/lm.proto
+++ b/internal/lm/proto/lm.proto
@@ -93,7 +93,7 @@ message TransferSandboxResponse {
         STATUS_UNSPECIFIED = 0;
         STATUS_BROWNOUT_IN_PROGRESS = 1;
         STATUS_BLACKOUT_IN_PROGRESS = 2;
-        STATUS_CONMPLETE = 3;
+        STATUS_COMPLETE = 3;
         STATUS_FAILED = 4;
         STATUS_CANCEL = 5;
     }
diff --git a/internal/taskserver/internal.go b/internal/taskserver/internal.go
index 0827969e61..53ba2c91c8 100644
--- a/internal/taskserver/internal.go
+++ b/internal/taskserver/internal.go
@@ -67,6 +67,13 @@ func newExecState(req *task.ExecProcessRequest, bundle string) *State {
 	return newState(req.ID, req.ExecID, req.Stdin, req.Stdout, req.Stderr, req.Terminal, bundle)
 }
 
+func restoredTaskState(req *task.CreateTaskRequest, pid uint32) *State {
+	// TODO: Don't hardcode RUNNING here.
+	state := newState(req.ID, "", req.Stdin, req.Stdout, req.Stderr, req.Terminal, req.Bundle)
+	state.setStarted(pid)
+	return state
+}
+
 func (s *State) setStarted(pid uint32) {
 	s.Pid = pid
 	s.Status = containerd_v1_types.Status_RUNNING
diff --git a/internal/taskserver/migration.go b/internal/taskserver/migration.go
index c7ed1ffb0b..06e174be53 100644
--- a/internal/taskserver/migration.go
+++ b/internal/taskserver/migration.go
@@ -375,12 +375,13 @@ func (s *service) newRestoreContainer(ctx context.Context, shimOpts *runhcsopts.
 		return err
 	}
 	t := &Task{
-		State: newTaskState(req),
+		State: restoredTaskState(req, taskState.Pid),
 		Ctr:   ctr,
 		Execs: make(map[string]*Exec),
 	}
 	s.sandbox.Tasks[req.ID] = t
 
+	// TODO: Don't assume it's already started here.
 	go waitContainer(s.sandbox.waitCtx, ctr, t.State, s.publisher)
 
 	return nil
diff --git a/internal/taskserver/service.go b/internal/taskserver/service.go
index 55b567b498..e4910cf923 100644
--- a/internal/taskserver/service.go
+++ b/internal/taskserver/service.go
@@ -84,6 +84,10 @@ func (s *service) Create(ctx context.Context, req *task.CreateTaskRequest) (*tas
 		return nil, errors.Wrap(errdefs.ErrFailedPrecondition, "if using terminal, stderr must be empty")
 	}
 
+	resp := &task.CreateTaskResponse{
+		Pid: 0,
+	}
+
 	if s.sandbox == nil {
 		switch shimOpts.BundleType {
 		case runhcsopts.BundleType_BUNDLE_OCI:
@@ -109,15 +113,12 @@ func (s *service) Create(ctx context.Context, req *task.CreateTaskRequest) (*tas
 			if err := s.newRestoreContainer(ctx, shimOpts, req); err != nil {
 				return nil, err
 			}
+			resp.Pid = s.sandbox.Tasks[req.ID].Pid
 		default:
 			return nil, fmt.Errorf("unsupported bundle type: %s", shimOpts.BundleType)
 		}
 	}
 
-	resp := &task.CreateTaskResponse{
-		Pid: 0,
-	}
-
 	if err := s.publisher.PublishEvent(ctx, runtime.TaskCreateEventTopic, &events.TaskCreate{
 		ContainerID: req.ID,
 		Bundle:      req.Bundle,