From 4f728689f51d80616de5661f10ab9fda6dce12e4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Leszczy=C5=84ski?= <2000michal@wp.pl>
Date: Wed, 23 Oct 2024 13:42:29 +0200
Subject: [PATCH] feat(managerclient): display bandwidth in sctool progress

Fixes #4042
---
 v3/pkg/managerclient/model.go | 68 +++++++++++++++++++++++++++++++++--
 v3/pkg/managerclient/utils.go |  6 ++--
 2 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/v3/pkg/managerclient/model.go b/v3/pkg/managerclient/model.go
index 68566846ad..bc594eb077 100644
--- a/v3/pkg/managerclient/model.go
+++ b/v3/pkg/managerclient/model.go
@@ -1021,6 +1021,9 @@ Duration:	{{ FormatDuration .StartTime .EndTime }}
 {{ end -}}
 {{ with .Progress }}Progress:	{{ if ne .Size 0 }}{{ FormatRestoreProgress .Size .Restored .Downloaded .Failed }}{{else}}-{{ end }}
 Snapshot Tag:	{{ .SnapshotTag }}
+Bandwidth:
+  - Download:    {{ avgDownload .Hosts }}
+  - Load&stream: {{ avgStream .Hosts }}
 {{ else }}Progress:	0%
 {{ end }}
 {{- if .Errors -}}
@@ -1038,6 +1041,8 @@ func (rp RestoreProgress) addHeader(w io.Writer) error {
 		"FormatError":           FormatError,
 		"FormatRestoreProgress": FormatRestoreProgress,
 		"status":                rp.status,
+		"avgDownload":           avgDownload,
+		"avgStream":             avgStream,
 	}).Parse(restoreProgressTemplate))
 	return temp.Execute(w, rp)
 }
@@ -1055,6 +1060,37 @@ func (rp RestoreProgress) status() string {
 	return s
 }
 
+func avgDownload(hosts []*models.RestoreHostProgress) string {
+	var bytes, milliseconds, shards int64
+	for _, hp := range hosts {
+		bytes += hp.DownloadedBytes
+		milliseconds += hp.DownloadDuration
+		shards += hp.ShardCnt
+	}
+	return formatBandwidth(bytes, milliseconds, shards)
+}
+
+func avgStream(hosts []*models.RestoreHostProgress) string {
+	var bytes, milliseconds, shards int64
+	for _, hp := range hosts {
+		bytes += hp.StreamedBytes
+		milliseconds += hp.StreamDuration
+		shards += hp.ShardCnt
+	}
+	return formatBandwidth(bytes, milliseconds, shards)
+}
+
+func formatBandwidth(bytes, milliseconds, shards int64) string {
+	if milliseconds <= 0 {
+		return "unknown"
+	}
+	bs := bytes * 1000 / milliseconds
+	if shards <= 0 {
+		return FormatSizeSuffix(bs) + "/s"
+	}
+	return FormatSizeSuffix(bs/shards) + "/s/shard"
+}
+
 func (rp RestoreProgress) hideKeyspace(keyspace string) bool {
 	if rp.KeyspaceFilter.Size() > 0 {
 		return rp.KeyspaceFilter.FirstMatch(keyspace) == -1
@@ -1082,10 +1118,15 @@ func (rp RestoreProgress) Render(w io.Writer) error {
 		}
 	}
 
-	if rp.Detailed && rp.Progress.Size > 0 {
-		if err := rp.addTableProgress(w); err != nil {
+	if rp.Detailed {
+		if err := rp.addHostProgress(w); err != nil {
 			return err
 		}
+		if rp.Progress.Size > 0 {
+			if err := rp.addTableProgress(w); err != nil {
+				return err
+			}
+		}
 	}
 
 	// Check if there is repair progress to display
@@ -1211,6 +1252,29 @@ func (rp RestoreProgress) addTableProgress(w io.Writer) error {
 	return nil
 }
 
+func (rp RestoreProgress) addHostProgress(w io.Writer) error {
+	_, _ = fmt.Fprintf(w, "\nHosts info\n")
+	t := table.New("Host", "Shards", "Download bandwidth", "Download duration", "Load&stream bandwidth", "Load&stream duration")
+	for i, hp := range rp.Progress.Hosts {
+		if i > 0 {
+			t.AddSeparator()
+		}
+		t.AddRow(
+			hp.Host,
+			hp.ShardCnt,
+			formatBandwidth(hp.DownloadedBytes, hp.DownloadDuration, hp.ShardCnt),
+			FormatMsDuration(hp.DownloadDuration),
+			formatBandwidth(hp.StreamedBytes, hp.StreamDuration, hp.ShardCnt),
+			FormatMsDuration(hp.StreamDuration),
+		)
+	}
+	t.SetColumnAlignment(termtables.AlignRight, 1, 2, 3, 4, 5)
+	if _, err := w.Write([]byte(t.String())); err != nil {
+		return err
+	}
+	return nil
+}
+
 func (rp RestoreProgress) addViewProgress(w io.Writer) {
 	if len(rp.Progress.Views) == 0 {
 		return
diff --git a/v3/pkg/managerclient/utils.go b/v3/pkg/managerclient/utils.go
index 4a5fe69f22..850c079d70 100644
--- a/v3/pkg/managerclient/utils.go
+++ b/v3/pkg/managerclient/utils.go
@@ -191,10 +191,10 @@ func FormatDuration(t0, t1 strfmt.DateTime) string {
 	return d.Truncate(time.Second).String()
 }
 
-// FormatMsDuration returns string representation of duration as number of
-// milliseconds.
+// FormatMsDuration returns string representation of given amount of milliseconds
+// as duration rounded up to a full second.
 func FormatMsDuration(d int64) string {
-	return (time.Duration(d) * time.Millisecond).Truncate(time.Second).String()
+	return (time.Duration(d)*time.Millisecond + time.Second - 1).Truncate(time.Second).String()
 }
 
 func isZero(t strfmt.DateTime) bool {