Register TaskState after LMFinalize. New drop for Apurv

kevpar · Nov 15, 2024 · e940475 · e940475
1 parent 4eec37e
commit e940475
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 2 deletions.
diff --git a/cmd/task/core.go b/cmd/task/core.go
@@ -322,7 +322,7 @@ var execCommand = cli.Command{
 			Usage: "Named pipe path",
 		},
 		cli.BoolFlag{
-			Name:  "terminal",
+			Name:  "tty",
 			Usage: "Enable terminal mode for task IO",
 		},
 	},
@@ -443,6 +443,33 @@ var waitCommand = cli.Command{
 	},
 }
 
+var connectCommand = cli.Command{
+	Name:           "connect",
+	ArgsUsage:      "<pipe> <task id>",
+	SkipArgReorder: true,
+	Before:         appargs.Validate(appargs.String, appargs.String),
+	Action: func(clictx *cli.Context) error {
+		args := clictx.Args()
+		address := args[0]
+		id := args[1]
+
+		conn, err := winio.DialPipe(address, nil)
+		if err != nil {
+			return fmt.Errorf("dial %s: %w", address, err)
+		}
+
+		client := ttrpc.NewClient(conn)
+		svc := task.NewTaskClient(client)
+
+		ctx := context.Background()
+
+		if _, err := svc.Connect(ctx, &task.ConnectRequest{ID: id}); err != nil {
+			return err
+		}
+		return nil
+	},
+}
+
 var shutdownCommand = cli.Command{
 	Name:           "shutdown",
 	ArgsUsage:      "<pipe> <task id>",

diff --git a/cmd/task/main.go b/cmd/task/main.go
@@ -23,6 +23,7 @@ func main() {
 		execCommand,
 		closeIOCommand,
 		waitCommand,
+		connectCommand,
 		shutdownCommand,
 		// Extra
 		pipeCommand,

diff --git a/internal/taskserver/migration.go b/internal/taskserver/migration.go
@@ -28,6 +28,7 @@ type migrationState struct {
 	c         windows.Handle
 	sandbox   core.Migratable
 	taskState map[string]*statepkg.TaskState
+	newID     string
 }
 
 var _ (lmproto.MigrationService) = (*service)(nil)
@@ -86,6 +87,7 @@ func (s *service) newSandboxLM(ctx context.Context, shimOpts *runhcsopts.Options
 		return err
 	}
 	s.migState = &migrationState{
+		newID:     req.ID,
 		sandbox:   sandbox,
 		taskState: config.Tasks,
 	}
@@ -205,6 +207,9 @@ func (s *service) FinalizeSandbox(ctx context.Context, req *lmproto.FinalizeSand
 			return nil, err
 		}
 		s.sandbox = &Sandbox{
+			State: &State{
+				TaskID: s.migState.newID,
+			},
 			Sandbox: s.migState.sandbox.(core.Sandbox),
 			Tasks:   make(map[string]*Task),
 		}

diff --git a/internal/vm2/vm.go b/internal/vm2/vm.go
@@ -244,7 +244,8 @@ func convertConfig(config *Config) (*hcsschema.ComputeSystem, error) {
 				LinuxKernelDirect: &hcsschema.LinuxKernelDirect{
 					KernelFilePath: filepath.Join(bootFilesDir, "vmlinux"),
 					InitRdPath:     filepath.Join(bootFilesDir, "initrd.img"),
-					KernelCmdLine:  fmt.Sprintf("init=/init 8250_core.nr_uarts=1 8250_core.skip_txen_test=1 console=ttyS0,115200 panic=-1 debug pci=off nr_cpus=%d brd.rd_nr=0 pmtmr=0 printk.devkmsg=on -- -e 1 /bin/gcs -v4 /log-format json -loglevel debug", config.ProcessorCount),
+					// KernelCmdLine:  fmt.Sprintf("init=/init 8250_core.nr_uarts=1 8250_core.skip_txen_test=1 console=ttyS0,115200 panic=-1 debug pci=off nr_cpus=%d brd.rd_nr=0 pmtmr=0 printk.devkmsg=on -- -e 1 sh -c \"/bin/gcs -v4 /log-format json -loglevel debug & exec sh\"", config.ProcessorCount),
+					KernelCmdLine: fmt.Sprintf("init=/init 8250_core.nr_uarts=1 8250_core.skip_txen_test=1 console=ttyS0,115200 panic=-1 debug pci=off nr_cpus=%d brd.rd_nr=0 pmtmr=0 printk.devkmsg=on -- -e 1 /bin/gcs -v4 /log-format json -loglevel debug", config.ProcessorCount),
 				},
 			},
 			ComputeTopology: &hcsschema.Topology{