From d59e8a84048b7c815050c8a72374fa230e4eb6bb Mon Sep 17 00:00:00 2001
From: Samuel Karp <samuelkarp@google.com>
Date: Wed, 7 Aug 2024 16:14:48 -0700
Subject: [PATCH 1/2] ctr: shim state for secondary tasks

The v2 shim interface supports grouping, so a single shim can manage
multiple tasks.  Prior to this change, the `shim state` command could
only query the state of the primary task (task that shares the same ID
as the shim).

Signed-off-by: Samuel Karp <samuelkarp@google.com>
---
 cmd/ctr/commands/shim/shim.go | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/cmd/ctr/commands/shim/shim.go b/cmd/ctr/commands/shim/shim.go
index 7ea60199a..72de3d9b5 100644
--- a/cmd/ctr/commands/shim/shim.go
+++ b/cmd/ctr/commands/shim/shim.go
@@ -67,7 +67,7 @@ var Command = &cli.Command{
 	Flags: []cli.Flag{
 		&cli.StringFlag{
 			Name:  "id",
-			Usage: "Container id",
+			Usage: "shim ID",
 		},
 	},
 	Subcommands: []*cli.Command{
@@ -115,14 +115,26 @@ var deleteCommand = &cli.Command{
 
 var stateCommand = &cli.Command{
 	Name:  "state",
-	Usage: "Get the state of all the processes of the task",
+	Usage: "Get the state of all main process of the task",
+	Flags: []cli.Flag{
+		&cli.StringFlag{
+			Name:    "task-id",
+			Aliases: []string{"t"},
+			Usage:   "task ID",
+		},
+	},
 	Action: func(cliContext *cli.Context) error {
 		service, err := getTaskService(cliContext)
 		if err != nil {
 			return err
 		}
+		id := cliContext.String("task-id")
+		if id == "" {
+			id = cliContext.String("id")
+		}
+
 		r, err := service.State(context.Background(), &task.StateRequest{
-			ID: cliContext.String("id"),
+			ID: id,
 		})
 		if err != nil {
 			return err

From 7d4da0cb2831a58bb2a98e99145cce2b3a627f2b Mon Sep 17 00:00:00 2001
From: Samuel Karp <samuelkarp@google.com>
Date: Wed, 7 Aug 2024 16:42:59 -0700
Subject: [PATCH 2/2] ctr: shim state query for old shims

Old shims do not implement containerd.task.v3.Task, but it can be
useful to use a new ctr with an older shim especially during upgrade
scenarios.

Signed-off-by: Samuel Karp <samuelkarp@google.com>
---
 cmd/ctr/commands/shim/shim.go | 64 +++++++++++++++++++++++++++++------
 1 file changed, 53 insertions(+), 11 deletions(-)

diff --git a/cmd/ctr/commands/shim/shim.go b/cmd/ctr/commands/shim/shim.go
index 72de3d9b5..ab869a110 100644
--- a/cmd/ctr/commands/shim/shim.go
+++ b/cmd/ctr/commands/shim/shim.go
@@ -28,7 +28,8 @@ import (
 	"strings"
 
 	"github.com/containerd/console"
-	"github.com/containerd/containerd/api/runtime/task/v3"
+	taskv2 "github.com/containerd/containerd/api/runtime/task/v2"
+	task "github.com/containerd/containerd/api/runtime/task/v3"
 	"github.com/containerd/containerd/v2/cmd/ctr/commands"
 	"github.com/containerd/containerd/v2/pkg/namespaces"
 	ptypes "github.com/containerd/containerd/v2/pkg/protobuf/types"
@@ -122,22 +123,48 @@ var stateCommand = &cli.Command{
 			Aliases: []string{"t"},
 			Usage:   "task ID",
 		},
+		&cli.IntFlag{
+			Name:  "api-version",
+			Usage: "shim API version {2,3}",
+			Value: 3,
+			Action: func(c *cli.Context, v int) error {
+				if v != 2 && v != 3 {
+					return fmt.Errorf("api-version must be 2 or 3")
+				}
+				return nil
+			},
+		},
 	},
 	Action: func(cliContext *cli.Context) error {
-		service, err := getTaskService(cliContext)
-		if err != nil {
-			return err
-		}
 		id := cliContext.String("task-id")
 		if id == "" {
 			id = cliContext.String("id")
 		}
 
-		r, err := service.State(context.Background(), &task.StateRequest{
-			ID: id,
-		})
-		if err != nil {
-			return err
+		var r any
+		switch cliContext.Int("api-version") {
+		case 2:
+			service, err := getTaskServiceV2(cliContext)
+			if err != nil {
+				return err
+			}
+			r, err = service.State(context.Background(), &taskv2.StateRequest{
+				ID: id,
+			})
+			if err != nil {
+				return err
+			}
+		default:
+			service, err := getTaskService(cliContext)
+			if err != nil {
+				return err
+			}
+			r, err = service.State(context.Background(), &task.StateRequest{
+				ID: id,
+			})
+			if err != nil {
+				return err
+			}
 		}
 		commands.PrintAsJSON(r)
 		return nil
@@ -246,6 +273,21 @@ var execCommand = &cli.Command{
 }
 
 func getTaskService(cliContext *cli.Context) (task.TTRPCTaskService, error) {
+	client, err := getTTRPCClient(cliContext)
+	if err != nil {
+		return nil, err
+	}
+	return task.NewTTRPCTaskClient(client), nil
+}
+func getTaskServiceV2(cliContext *cli.Context) (taskv2.TaskService, error) {
+	client, err := getTTRPCClient(cliContext)
+	if err != nil {
+		return nil, err
+	}
+	return taskv2.NewTaskClient(client), nil
+}
+
+func getTTRPCClient(cliContext *cli.Context) (*ttrpc.Client, error) {
 	id := cliContext.String("id")
 	if id == "" {
 		return nil, fmt.Errorf("container id must be specified")
@@ -268,7 +310,7 @@ func getTaskService(cliContext *cli.Context) (task.TTRPCTaskService, error) {
 			// TODO(stevvooe): This actually leaks the connection. We were leaking it
 			// before, so may not be a huge deal.
 
-			return task.NewTTRPCTaskClient(client), nil
+			return client, nil
 		}
 	}