feat: add uc ps command to list all containers in the cluster (#165)

* feat: Add `uc ps` command to list all containers in the cluster

* Fix indentation

* Improvements to the container state highlights code (use enum, better var naming)

* Add group/sort option by "health"

* Rename `uc ps --group-by` flag to `--sort`. Also, remove `--context` flag.

* Use huh.spinner with standard style/type

* Remove pointless TrimPrefix on container name

* Move machine lookup out of services loop

* Use more efficient `ListServiceContainers`

* Use consistent header other

* Use api.ProxyMachinesContext helper for context to query all machines for service containers

* Make errored MachineServiceContainer entries a warning, not return an error immediately, in `uc ps`

* Use lipgloss table instead of tabwriter for "uc ps" command

* Handle nil metadata in `uc ps` on container when cluster only has one machine

cmd/uncloud/main.go

@@ -112,6 +112,7 @@ func main() {
 		NewDocsCommand(),
 		NewBuildCommand(),
 		NewImagesCommand(),
+		NewPsCommand(),
 		caddy.NewRootCommand(),
 		cmdcontext.NewRootCommand(),
 		dns.NewRootCommand(),

cmd/uncloud/ps.go

@@ -0,0 +1,266 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"sort"
+
+	"github.com/charmbracelet/huh/spinner"
+	"github.com/charmbracelet/lipgloss"
+	"github.com/charmbracelet/lipgloss/table"
+	"github.com/docker/docker/api/types/container"
+	"github.com/spf13/cobra"
+
+	"github.com/psviderski/uncloud/internal/cli"
+	"github.com/psviderski/uncloud/pkg/api"
+	"github.com/psviderski/uncloud/pkg/client"
+)
+
+const (
+	sortByService = "service"
+	sortByMachine = "machine"
+	sortByHealth  = "health"
+)
+
+type containerHighlight int
+
+const (
+	highlightDanger containerHighlight = iota
+	highlightWarning
+	highlightSuccess
+	highlightNormal
+)
+
+type psOptions struct {
+	sortBy string
+}
+
+func NewPsCommand() *cobra.Command {
+	opts := psOptions{}
+	cmd := &cobra.Command{
+		Use:   "ps",
+		Short: "List all service containers in the cluster",
+		Long: `List all service containers across all machines in the cluster.
+
+This command provides a comprehensive overview of all running containers that are part of a service,
+making it easy to see the distribution and status of containers across the cluster.`,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			if opts.sortBy != sortByService && opts.sortBy != sortByMachine && opts.sortBy != sortByHealth {
+				return fmt.Errorf("invalid value for --sort: %q, must be one of '%s', '%s' or '%s'", opts.sortBy, sortByService, sortByMachine, sortByHealth)
+			}
+			return runPs(cmd, opts)
+		},
+	}
+	cmd.Flags().StringVarP(&opts.sortBy, "sort", "s", sortByService, "Sort containers by 'service', 'machine' or 'health'")
+	return cmd
+}
+
+type containerInfo struct {
+	serviceName string
+	machineName string
+	id          string
+	name        string
+	image       string
+	status      string
+	highlight   containerHighlight
+}
+
+func runPs(cmd *cobra.Command, opts psOptions) error {
+	uncli := cmd.Context().Value("cli").(*cli.CLI)
+	client, err := uncli.ConnectCluster(cmd.Context())
+	if err != nil {
+		return fmt.Errorf("connect to cluster: %w", err)
+	}
+	defer client.Close()
+
+	var containers []containerInfo
+	err = spinner.New().
+		Title(" Collecting container info...").
+		Type(spinner.MiniDot).
+		Style(lipgloss.NewStyle().Foreground(lipgloss.Color("3"))).
+		ActionWithErr(func(ctx context.Context) error {
+			containers, err = collectContainers(ctx, client)
+			return err
+		}).
+		Run()
+	if err != nil {
+		return fmt.Errorf("collect containers: %w", err)
+	}
+
+	// Sort the containers based on the sorting option
+	sort.SliceStable(containers, func(i, j int) bool {
+		a, b := containers[i], containers[j]
+		switch opts.sortBy {
+		case sortByHealth:
+			if a.highlight != b.highlight {
+				return a.highlight < b.highlight
+			}
+			if a.serviceName != b.serviceName {
+				return a.serviceName < b.serviceName
+			}
+			if a.machineName != b.machineName {
+				return a.machineName < b.machineName
+			}
+		case sortByMachine:
+			if a.machineName != b.machineName {
+				return a.machineName < b.machineName
+			}
+			if a.serviceName != b.serviceName {
+				return a.serviceName < b.serviceName
+			}
+		default: // sortByService
+			if a.serviceName != b.serviceName {
+				return a.serviceName < b.serviceName
+			}
+			if a.machineName != b.machineName {
+				return a.machineName < b.machineName
+			}
+		}
+		// Final tie-breaker
+		return a.name < b.name
+	})
+
+	return printContainers(containers)
+}
+
+func printContainers(containers []containerInfo) error {
+	t := table.New().
+		// Remove the default border.
+		Border(lipgloss.Border{}).
+		BorderTop(false).
+		BorderBottom(false).
+		BorderLeft(false).
+		BorderRight(false).
+		BorderHeader(false).
+		BorderColumn(false).
+		StyleFunc(func(row, col int) lipgloss.Style {
+			if row == table.HeaderRow {
+				return lipgloss.NewStyle().Bold(true).PaddingRight(3)
+			}
+			// Regular style for data rows with padding.
+			return lipgloss.NewStyle().PaddingRight(3)
+		})
+
+	t.Headers("SERVICE", "CONTAINER ID", "NAME", "IMAGE", "STATUS", "MACHINE")
+
+	for _, ctr := range containers {
+		id := ctr.id
+		if len(id) > 12 {
+			id = id[:12]
+		}
+
+		var statusStyle lipgloss.Style
+		switch ctr.highlight {
+		case highlightSuccess:
+			statusStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("2")) // Green
+		case highlightDanger:
+			statusStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("1")) // Red
+		case highlightWarning:
+			statusStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("3")) // Yellow
+		default:
+			statusStyle = lipgloss.NewStyle() // Default
+		}
+
+		t.Row(
+			ctr.serviceName,
+			id,
+			ctr.name,
+			ctr.image,
+			statusStyle.Render(ctr.status),
+			ctr.machineName,
+		)
+	}
+
+	fmt.Println(t)
+	return nil
+}
+
+func collectContainers(ctx context.Context, cli *client.Client) ([]containerInfo, error) {
+	listCtx, machines, err := api.ProxyMachinesContext(ctx, cli, nil)
+	if err != nil {
+		return nil, fmt.Errorf("proxy machines context: %w", err)
+	}
+
+	// Create a map of IP to machine name for resolving response metadata
+	machinesNamesByIP := make(map[string]string)
+	for _, m := range machines {
+		if addr, err := m.Machine.Network.ManagementIp.ToAddr(); err == nil {
+			machinesNamesByIP[addr.String()] = m.Machine.Name
+		}
+	}
+
+	// List all service containers across all machines in the cluster.
+	machineContainers, err := cli.Docker.ListServiceContainers(
+		listCtx, "", container.ListOptions{All: true},
+	)
+	if err != nil {
+		return nil, fmt.Errorf("list service containers: %w", err)
+	}
+
+	var containers []containerInfo
+	for _, msc := range machineContainers {
+		machineName := "unknown"
+
+		// Metadata can be nil if the request was broadcasted to only one machine.
+		if msc.Metadata == nil && len(machineContainers) > 1 {
+			return nil, fmt.Errorf("something went wrong with gRPC proxy: metadata is missing for a machine response")
+		}
+		if msc.Metadata != nil && msc.Metadata.Error != "" {
+			client.PrintWarning(fmt.Sprintf("failed to list containers on machine %s: %s", machineName, msc.Metadata.Error))
+			continue
+		}
+
+		if msc.Metadata != nil {
+			var ok bool
+			machineName, ok = machinesNamesByIP[msc.Metadata.Machine]
+			if !ok {
+				// Fallback to machine's IP as name.
+				machineName = msc.Metadata.Machine
+			}
+		} else {
+			// Fallback to the first available machine name.
+			if len(machines) > 0 {
+				machineName = machines[0].Machine.Name
+			}
+		}
+
+		for _, ctr := range msc.Containers {
+			if ctr.Container.State == nil || ctr.Container.Config == nil {
+				continue
+			}
+
+			status, err := ctr.Container.HumanState()
+			if err != nil {
+				return nil, fmt.Errorf("get human state for container %s: %w", ctr.Container.ID, err)
+			}
+
+			var highlight containerHighlight
+			healthStatus := ""
+			if ctr.Container.State.Health != nil {
+				healthStatus = ctr.Container.State.Health.Status
+			}
+
+			if healthStatus == container.Unhealthy || ctr.Container.State.Status == "dead" || ctr.Container.State.OOMKilled || ctr.Container.State.Dead {
+				highlight = highlightDanger
+			} else if healthStatus == container.Healthy {
+				highlight = highlightSuccess
+			} else if ctr.Container.State.Status == "running" {
+				highlight = highlightNormal
+			} else { // Other non-critical but noteworthy states
+				highlight = highlightWarning
+			}
+
+			info := containerInfo{
+				serviceName: ctr.ServiceName(),
+				machineName: machineName,
+				id:          ctr.Container.ID,
+				name:        ctr.Container.Name,
+				image:       ctr.Container.Config.Image,
+				status:      status,
+				highlight:   highlight,
+			}
+			containers = append(containers, info)
+		}
+	}
+	return containers, nil
+}

cmd/uncloud/ps_test.go

@@ -0,0 +1,424 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"net/netip"
+	"testing"
+
+	"github.com/psviderski/uncloud/internal/machine/api/pb"
+	"github.com/psviderski/uncloud/internal/machine/docker"
+	"github.com/psviderski/uncloud/pkg/client"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"google.golang.org/grpc"
+	"google.golang.org/protobuf/types/known/emptypb"
+)
+
+// mockDockerClient implements pb.DockerClient
+type mockDockerClient struct {
+	pb.DockerClient // Embed to avoid implementing all methods
+	listResp        *pb.ListServiceContainersResponse
+	listErr         error
+}
+
+func (m *mockDockerClient) ListServiceContainers(ctx context.Context, in *pb.ListServiceContainersRequest, opts ...grpc.CallOption) (*pb.ListServiceContainersResponse, error) {
+	return m.listResp, m.listErr
+}
+
+// mockClusterClient implements pb.ClusterClient
+type mockClusterClient struct {
+	pb.ClusterClient // Embed to avoid implementing all methods
+	machinesResp     *pb.ListMachinesResponse
+	machinesErr      error
+}
+
+func (m *mockClusterClient) ListMachines(ctx context.Context, in *emptypb.Empty, opts ...grpc.CallOption) (*pb.ListMachinesResponse, error) {
+	return m.machinesResp, m.machinesErr
+}
+
+func TestCollectContainers_NilMetadata(t *testing.T) {
+	// Setup container data
+	containerData := map[string]interface{}{
+		"Id":   "container1",
+		"Name": "test-container",
+		"Config": map[string]interface{}{
+			"Image": "test-image",
+		},
+		"State": map[string]interface{}{
+			"Status":     "running",
+			"StartedAt":  "2023-01-01T12:00:00Z",
+			"FinishedAt": "0001-01-01T00:00:00Z",
+		},
+	}
+	containerJSON, err := json.Marshal(containerData)
+	require.NoError(t, err)
+
+	serviceSpecJSON, err := json.Marshal(map[string]interface{}{})
+	require.NoError(t, err)
+
+	// Setup mocks
+	mockDocker := &mockDockerClient{
+		listResp: &pb.ListServiceContainersResponse{
+			Messages: []*pb.MachineServiceContainers{
+				{
+					Metadata: nil, // Simulating the issue: nil metadata
+					Containers: []*pb.ServiceContainer{
+						{
+							Container:   containerJSON,
+							ServiceSpec: serviceSpecJSON,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	machineIP := "10.0.0.1"
+	mockCluster := &mockClusterClient{
+		machinesResp: &pb.ListMachinesResponse{
+			Machines: []*pb.MachineMember{
+				{
+					Machine: &pb.MachineInfo{
+						Name: "machine-1",
+						Network: &pb.NetworkConfig{
+							ManagementIp: pb.NewIP(netip.MustParseAddr(machineIP)),
+						},
+					},
+					State: pb.MachineMember_UP,
+				},
+			},
+		},
+	}
+
+	// Construct client with mocks
+	cli := &client.Client{
+		Docker: &docker.Client{
+			GRPCClient: mockDocker,
+		},
+		ClusterClient: mockCluster,
+	}
+
+	// Execute
+	containers, err := collectContainers(context.Background(), cli)
+	require.NoError(t, err)
+
+	// Verify
+	assert.Len(t, containers, 1)
+	if len(containers) > 0 {
+		c := containers[0]
+		assert.Equal(t, "container1", c.id)
+		assert.Equal(t, "test-container", c.name)
+		assert.Equal(t, "machine-1", c.machineName, "Should fall back to the single machine name when metadata is nil")
+	}
+}
+
+func TestCollectContainers_NilMetadata_MultipleMachines_Error(t *testing.T) {
+	// If we have multiple machines but receive nil metadata, it should return an error as it is ambiguous
+
+	// Setup container data
+	containerData1 := map[string]interface{}{
+		"Id": "container1",
+	}
+	containerJSON1, _ := json.Marshal(containerData1)
+
+	containerData2 := map[string]interface{}{
+		"Id": "container2",
+		"Config": map[string]interface{}{
+			"Image": "test-image",
+		},
+		"State": map[string]interface{}{
+			"Status":     "running",
+			"StartedAt":  "2023-01-01T12:00:00Z",
+			"FinishedAt": "0001-01-01T00:00:00Z",
+		},
+	}
+	containerJSON2, _ := json.Marshal(containerData2)
+
+	serviceSpecJSON, _ := json.Marshal(map[string]interface{}{})
+
+	// Setup mocks
+	mockDocker := &mockDockerClient{
+		listResp: &pb.ListServiceContainersResponse{
+			Messages: []*pb.MachineServiceContainers{
+				{
+					Metadata: nil, // Nil metadata
+					Containers: []*pb.ServiceContainer{
+						{
+							Container:   containerJSON1,
+							ServiceSpec: serviceSpecJSON,
+						},
+					},
+				},
+				{
+					Metadata: &pb.Metadata{Machine: "10.0.0.2"},
+					Containers: []*pb.ServiceContainer{
+						{
+							Container:   containerJSON2,
+							ServiceSpec: serviceSpecJSON,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	mockCluster := &mockClusterClient{
+		machinesResp: &pb.ListMachinesResponse{
+			Machines: []*pb.MachineMember{
+				{
+					Machine: &pb.MachineInfo{
+						Name: "machine-1",
+						Network: &pb.NetworkConfig{
+							ManagementIp: pb.NewIP(netip.MustParseAddr("10.0.0.1")),
+						},
+					},
+				},
+				{
+					Machine: &pb.MachineInfo{
+						Name: "machine-2",
+						Network: &pb.NetworkConfig{
+							ManagementIp: pb.NewIP(netip.MustParseAddr("10.0.0.2")),
+						},
+					},
+				},
+			},
+		},
+	}
+
+	// Construct client with mocks
+	cli := &client.Client{
+		Docker: &docker.Client{
+			GRPCClient: mockDocker,
+		},
+		ClusterClient: mockCluster,
+	}
+
+	// Execute
+	_, err := collectContainers(context.Background(), cli)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "metadata is missing for a machine response")
+}
+
+func TestCollectContainers_MetadataPresent_MultipleMachines(t *testing.T) {
+	// Verify correct mapping of containers to machines when metadata is present
+
+	// Setup container data
+	containerData1 := map[string]interface{}{
+		"Id":   "container1",
+		"Name": "container-1",
+		"Config": map[string]interface{}{
+			"Image": "image-1",
+		},
+		"State": map[string]interface{}{
+			"Status":     "running",
+			"StartedAt":  "2023-01-01T12:00:00Z",
+			"FinishedAt": "0001-01-01T00:00:00Z",
+		},
+	}
+	containerJSON1, _ := json.Marshal(containerData1)
+
+	containerData2 := map[string]interface{}{
+		"Id":   "container2",
+		"Name": "container-2",
+		"Config": map[string]interface{}{
+			"Image": "image-2",
+		},
+		"State": map[string]interface{}{
+			"Status":     "running",
+			"StartedAt":  "2023-01-01T12:00:00Z",
+			"FinishedAt": "0001-01-01T00:00:00Z",
+		},
+	}
+	containerJSON2, _ := json.Marshal(containerData2)
+
+	serviceSpecJSON, _ := json.Marshal(map[string]interface{}{})
+
+	// Setup mocks
+	mockDocker := &mockDockerClient{
+		listResp: &pb.ListServiceContainersResponse{
+			Messages: []*pb.MachineServiceContainers{
+				{
+					Metadata: &pb.Metadata{Machine: "10.0.0.1"},
+					Containers: []*pb.ServiceContainer{
+						{
+							Container:   containerJSON1,
+							ServiceSpec: serviceSpecJSON,
+						},
+					},
+				},
+				{
+					Metadata: &pb.Metadata{Machine: "10.0.0.2"},
+					Containers: []*pb.ServiceContainer{
+						{
+							Container:   containerJSON2,
+							ServiceSpec: serviceSpecJSON,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	mockCluster := &mockClusterClient{
+		machinesResp: &pb.ListMachinesResponse{
+			Machines: []*pb.MachineMember{
+				{
+					Machine: &pb.MachineInfo{
+						Name: "machine-1",
+						Network: &pb.NetworkConfig{
+							ManagementIp: pb.NewIP(netip.MustParseAddr("10.0.0.1")),
+						},
+					},
+				},
+				{
+					Machine: &pb.MachineInfo{
+						Name: "machine-2",
+						Network: &pb.NetworkConfig{
+							ManagementIp: pb.NewIP(netip.MustParseAddr("10.0.0.2")),
+						},
+					},
+				},
+			},
+		},
+	}
+
+	cli := &client.Client{
+		Docker:        &docker.Client{GRPCClient: mockDocker},
+		ClusterClient: mockCluster,
+	}
+
+	containers, err := collectContainers(context.Background(), cli)
+	require.NoError(t, err)
+
+	assert.Len(t, containers, 2)
+	// Order is not guaranteed by the map iteration in logic or parallel fetch (though here it's mocked sequential),
+	// but collectContainers just appends.
+	// We'll find them by ID.
+	for _, c := range containers {
+		if c.id == "container1" {
+			assert.Equal(t, "machine-1", c.machineName)
+		} else if c.id == "container2" {
+			assert.Equal(t, "machine-2", c.machineName)
+		} else {
+			t.Errorf("unexpected container id: %s", c.id)
+		}
+	}
+}
+
+func TestCollectContainers_NilMetadata_NoMachines(t *testing.T) {
+	// Case: 1 msc with nil metadata but no machines at all
+
+	containerData := map[string]interface{}{
+		"Id": "container1",
+		"Config": map[string]interface{}{
+			"Image": "test-image",
+		},
+		"State": map[string]interface{}{
+			"Status":     "running",
+			"StartedAt":  "2023-01-01T12:00:00Z",
+			"FinishedAt": "0001-01-01T00:00:00Z",
+		},
+	}
+	containerJSON, _ := json.Marshal(containerData)
+	serviceSpecJSON, _ := json.Marshal(map[string]interface{}{})
+
+	mockDocker := &mockDockerClient{
+		listResp: &pb.ListServiceContainersResponse{
+			Messages: []*pb.MachineServiceContainers{
+				{
+					Metadata: nil,
+					Containers: []*pb.ServiceContainer{
+						{
+							Container:   containerJSON,
+							ServiceSpec: serviceSpecJSON,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	// No machines in cluster response
+	mockCluster := &mockClusterClient{
+		machinesResp: &pb.ListMachinesResponse{
+			Machines: []*pb.MachineMember{},
+		},
+	}
+
+	cli := &client.Client{
+		Docker:        &docker.Client{GRPCClient: mockDocker},
+		ClusterClient: mockCluster,
+	}
+
+	containers, err := collectContainers(context.Background(), cli)
+	require.NoError(t, err)
+
+	assert.Len(t, containers, 1)
+	if len(containers) > 0 {
+		assert.Equal(t, "unknown", containers[0].machineName)
+	}
+}
+
+func TestCollectContainers_MetadataPresent_NotInMapping(t *testing.T) {
+	// Case: msc with metadata that is not in the IP-to-name mapping
+
+	containerData := map[string]interface{}{
+		"Id": "container1",
+		"Config": map[string]interface{}{
+			"Image": "test-image",
+		},
+		"State": map[string]interface{}{
+			"Status":     "running",
+			"StartedAt":  "2023-01-01T12:00:00Z",
+			"FinishedAt": "0001-01-01T00:00:00Z",
+		},
+	}
+	containerJSON, _ := json.Marshal(containerData)
+	serviceSpecJSON, _ := json.Marshal(map[string]interface{}{})
+
+	mockDocker := &mockDockerClient{
+		listResp: &pb.ListServiceContainersResponse{
+			Messages: []*pb.MachineServiceContainers{
+				{
+					Metadata: &pb.Metadata{Machine: "10.0.0.99"}, // Unknown IP
+					Containers: []*pb.ServiceContainer{
+						{
+							Container:   containerJSON,
+							ServiceSpec: serviceSpecJSON,
+						},
+					},
+				},
+			},
+		},
+	}
+
+	mockCluster := &mockClusterClient{
+		machinesResp: &pb.ListMachinesResponse{
+			Machines: []*pb.MachineMember{
+				{
+					Machine: &pb.MachineInfo{
+						Name: "machine-1",
+						Network: &pb.NetworkConfig{
+							ManagementIp: pb.NewIP(netip.MustParseAddr("10.0.0.1")),
+						},
+					},
+				},
+			},
+		},
+	}
+
+	cli := &client.Client{
+		Docker:        &docker.Client{GRPCClient: mockDocker},
+		ClusterClient: mockCluster,
+	}
+
+	containers, err := collectContainers(context.Background(), cli)
+	require.NoError(t, err)
+
+	assert.Len(t, containers, 1)
+	if len(containers) > 0 {
+		// Should fallback to the IP/string in metadata
+		assert.Equal(t, "10.0.0.99", containers[0].machineName)
+	}
+}