Return restored container pid from Create

kevpar · Nov 18, 2024 · a9b2ce6 · a9b2ce6
1 parent bff6351
commit a9b2ce6
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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,