From 6df46463a9bcb850be5bb5f8ba6546e8a7105f80 Mon Sep 17 00:00:00 2001
From: allencloud <allen.sun@daocloud.io>
Date: Sun, 4 Sep 2016 15:38:50 +0800
Subject: [PATCH] support docker node ps multiNodes

Signed-off-by: allencloud <allen.sun@daocloud.io>
---
 command/node/ps.go | 66 +++++++++++++++++++++++++++++++---------------
 1 file changed, 45 insertions(+), 21 deletions(-)

diff --git a/command/node/ps.go b/command/node/ps.go
index 84d4b375ac..607488f35e 100644
--- a/command/node/ps.go
+++ b/command/node/ps.go
@@ -1,7 +1,11 @@
 package node
 
 import (
+	"fmt"
+	"strings"
+
 	"github.com/docker/docker/api/types"
+	"github.com/docker/docker/api/types/swarm"
 	"github.com/docker/docker/cli"
 	"github.com/docker/docker/cli/command"
 	"github.com/docker/docker/cli/command/idresolver"
@@ -12,7 +16,7 @@ import (
 )
 
 type psOptions struct {
-	nodeID    string
+	nodeIDs   []string
 	noResolve bool
 	noTrunc   bool
 	filter    opts.FilterOpt
@@ -22,14 +26,14 @@ func newPsCommand(dockerCli *command.DockerCli) *cobra.Command {
 	opts := psOptions{filter: opts.NewFilterOpt()}
 
 	cmd := &cobra.Command{
-		Use:   "ps [OPTIONS] [NODE]",
-		Short: "List tasks running on a node, defaults to current node",
-		Args:  cli.RequiresRangeArgs(0, 1),
+		Use:   "ps [OPTIONS] [NODE...]",
+		Short: "List tasks running on one or more nodes, defaults to current node",
+		Args:  cli.RequiresMinArgs(0),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			opts.nodeID = "self"
+			opts.nodeIDs = []string{"self"}
 
 			if len(args) != 0 {
-				opts.nodeID = args[0]
+				opts.nodeIDs = args
 			}
 
 			return runPs(dockerCli, opts)
@@ -47,23 +51,43 @@ func runPs(dockerCli *command.DockerCli, opts psOptions) error {
 	client := dockerCli.Client()
 	ctx := context.Background()
 
-	nodeRef, err := Reference(ctx, client, opts.nodeID)
-	if err != nil {
-		return nil
-	}
-	node, _, err := client.NodeInspectWithRaw(ctx, nodeRef)
-	if err != nil {
-		return err
+	var (
+		errs  []string
+		tasks []swarm.Task
+	)
+
+	for _, nodeID := range opts.nodeIDs {
+		nodeRef, err := Reference(ctx, client, nodeID)
+		if err != nil {
+			errs = append(errs, err.Error())
+			continue
+		}
+
+		node, _, err := client.NodeInspectWithRaw(ctx, nodeRef)
+		if err != nil {
+			errs = append(errs, err.Error())
+			continue
+		}
+
+		filter := opts.filter.Value()
+		filter.Add("node", node.ID)
+
+		nodeTasks, err := client.TaskList(ctx, types.TaskListOptions{Filter: filter})
+		if err != nil {
+			errs = append(errs, err.Error())
+			continue
+		}
+
+		tasks = append(tasks, nodeTasks...)
 	}
 
-	filter := opts.filter.Value()
-	filter.Add("node", node.ID)
-	tasks, err := client.TaskList(
-		ctx,
-		types.TaskListOptions{Filter: filter})
-	if err != nil {
-		return err
+	if err := task.Print(dockerCli, ctx, tasks, idresolver.New(client, opts.noResolve), opts.noTrunc); err != nil {
+		errs = append(errs, err.Error())
 	}
 
-	return task.Print(dockerCli, ctx, tasks, idresolver.New(client, opts.noResolve), opts.noTrunc)
+	if len(errs) > 0 {
+		return fmt.Errorf("%s", strings.Join(errs, "\n"))
+	}
+
+	return nil
 }