From c26b7469bdb3a8cedf5da9e936c9678b43503447 Mon Sep 17 00:00:00 2001
From: Ken Hibino <ken.hibino7@gmail.com>
Date: Mon, 7 Sep 2020 12:55:32 -0700
Subject: [PATCH] Display cluster info in stats command when --cluster flag is
 passed

---
 internal/rdb/inspect.go  | 15 ++++++++++++++-
 tools/asynq/cmd/stats.go | 30 +++++++++++++++++++++++++++---
 2 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/internal/rdb/inspect.go b/internal/rdb/inspect.go
index a17535a..bec0752 100644
--- a/internal/rdb/inspect.go
+++ b/internal/rdb/inspect.go
@@ -222,8 +222,21 @@ func (r *RDB) RedisInfo() (map[string]string, error) {
 	if err != nil {
 		return nil, err
 	}
+	return parseInfo(res)
+}
+
+// RedisClusterInfo returns a map of redis cluster info.
+func (r *RDB) RedisClusterInfo() (map[string]string, error) {
+	res, err := r.client.ClusterInfo().Result()
+	if err != nil {
+		return nil, err
+	}
+	return parseInfo(res)
+}
+
+func parseInfo(infoStr string) (map[string]string, error) {
 	info := make(map[string]string)
-	lines := strings.Split(res, "\r\n")
+	lines := strings.Split(infoStr, "\r\n")
 	for _, l := range lines {
 		kv := strings.Split(l, ":")
 		if len(kv) == 2 {
diff --git a/tools/asynq/cmd/stats.go b/tools/asynq/cmd/stats.go
index 1a9cb73..00bf059 100644
--- a/tools/asynq/cmd/stats.go
+++ b/tools/asynq/cmd/stats.go
@@ -6,6 +6,7 @@ package cmd
 
 import (
 	"fmt"
+	"io"
 	"os"
 	"strconv"
 	"strings"
@@ -88,7 +89,12 @@ func stats(cmd *cobra.Command, args []string) {
 		aggStats.Timestamp = s.Timestamp
 		stats = append(stats, s)
 	}
-	info, err := r.RedisInfo()
+	var info map[string]string
+	if useRedisCluster {
+		info, err = r.RedisClusterInfo()
+	} else {
+		info, err = r.RedisInfo()
+	}
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(1)
@@ -105,8 +111,13 @@ func stats(cmd *cobra.Command, args []string) {
 	printSuccessFailureStats(&aggStats)
 	fmt.Println()
 
-	fmt.Println("REDIS INFO")
-	printInfo(info)
+	if useRedisCluster {
+		fmt.Println("REDIS CLUSTER INFO")
+		printClusterInfo(info)
+	} else {
+		fmt.Println("REDIS INFO")
+		printInfo(info)
+	}
 	fmt.Println()
 }
 
@@ -174,6 +185,19 @@ func printInfo(info map[string]string) {
 	tw.Flush()
 }
 
+func printClusterInfo(info map[string]string) {
+	printTable(
+		[]string{"State", "Known Nodes", "Cluster Size"},
+		func(w io.Writer, tmpl string) {
+			fmt.Fprintf(w, tmpl,
+				strings.ToUpper(info["cluster_state"]),
+				info["cluster_known_nodes"],
+				info["cluster_size"],
+			)
+		},
+	)
+}
+
 func toInterfaceSlice(strs []string) []interface{} {
 	var res []interface{}
 	for _, s := range strs {