From 52987e7b59022c7a55eb9c3059cf188f095e14b0 Mon Sep 17 00:00:00 2001 From: Shaunak Kashyap Date: Wed, 18 Oct 2023 22:17:13 -0700 Subject: [PATCH] Add units to download rate --- .../download/http/progress_observer.go | 10 ++-- .../application/upgrade/details/details.go | 49 ++++++++++++++----- .../upgrade/details/details_test.go | 4 +- 3 files changed, 43 insertions(+), 20 deletions(-) diff --git a/internal/pkg/agent/application/upgrade/artifact/download/http/progress_observer.go b/internal/pkg/agent/application/upgrade/artifact/download/http/progress_observer.go index 08977cccc85..ca024c53c88 100644 --- a/internal/pkg/agent/application/upgrade/artifact/download/http/progress_observer.go +++ b/internal/pkg/agent/application/upgrade/artifact/download/http/progress_observer.go @@ -110,21 +110,21 @@ func newDetailsProgressObserver(upgradeDetails *details.Details) *detailsProgres } } -func (dpObs *detailsProgressObserver) Report(sourceURI string, timePast time.Duration, downloadedBytes, totalBytes, percentComplete, downloadRate float64) { +func (dpObs *detailsProgressObserver) Report(sourceURI string, timePast time.Duration, downloadedBytes, totalBytes, percentComplete, downloadRateBytesPerSecond float64) { dpObs.mu.Lock() defer dpObs.mu.Unlock() - dpObs.upgradeDetails.SetDownloadProgress(percentComplete, downloadRate) + dpObs.upgradeDetails.SetDownloadProgress(percentComplete, downloadRateBytesPerSecond) } -func (dpObs *detailsProgressObserver) ReportCompleted(sourceURI string, timePast time.Duration, downloadRate float64) { +func (dpObs *detailsProgressObserver) ReportCompleted(sourceURI string, timePast time.Duration, downloadRateBytesPerSecond float64) { dpObs.mu.Lock() defer dpObs.mu.Unlock() - dpObs.upgradeDetails.SetDownloadProgress(1, downloadRate) + dpObs.upgradeDetails.SetDownloadProgress(1, downloadRateBytesPerSecond) } -func (dpObs *detailsProgressObserver) ReportFailed(sourceURI string, timePast time.Duration, downloadedBytes, totalBytes, percentComplete, downloadRate float64, err error) { +func (dpObs *detailsProgressObserver) ReportFailed(sourceURI string, timePast time.Duration, downloadedBytes, totalBytes, percentComplete, downloadRateBytesPerSecond float64, err error) { dpObs.mu.Lock() defer dpObs.mu.Unlock() diff --git a/internal/pkg/agent/application/upgrade/details/details.go b/internal/pkg/agent/application/upgrade/details/details.go index e8b29f6d2da..028990aafcd 100644 --- a/internal/pkg/agent/application/upgrade/details/details.go +++ b/internal/pkg/agent/application/upgrade/details/details.go @@ -6,13 +6,17 @@ package details import ( "encoding/json" + "fmt" "math" + "strings" "sync" "time" + + "github.com/docker/go-units" ) // downloadRate is a float64 that can be safely marshalled to JSON -// when the value is Infinity. +// when the value is Infinity. The rate is always in bytes/second units. type downloadRate float64 // Observer is a function that will be called with upgrade details @@ -68,12 +72,12 @@ func (d *Details) SetState(s State) { // SetDownloadProgress is a convenience method to set the download percent // when the upgrade is in UPG_DOWNLOADING state. -func (d *Details) SetDownloadProgress(percent, rate float64) { +func (d *Details) SetDownloadProgress(percent, rateBytesPerSecond float64) { d.mu.Lock() defer d.mu.Unlock() d.Metadata.DownloadPercent = percent - d.Metadata.DownloadRate = downloadRate(rate) + d.Metadata.DownloadRate = downloadRate(rateBytesPerSecond) d.notifyObservers() } @@ -128,16 +132,35 @@ func (d *Details) notifyObserver(observer Observer) { } } -func (dr downloadRate) MarshalJSON() ([]byte, error) { - // When the Agent artifact is downloaded really fast during an upgrade, the - // download rate gets set to +Inf. This value is set on the - // details.Metadata.DownloadRate field. Unfortunately, JSON does not support - // +/-Inf or NaN values; see https://www.rfc-editor.org/rfc/rfc8259. So we - // reset this field to a sentinel value of -1 before marshalling the object - // to JSON. - if math.IsInf(float64(dr), 0) { - return json.Marshal(-1.0) +func (dr *downloadRate) MarshalJSON() ([]byte, error) { + downloadRateBytesPerSecond := float64(*dr) + if math.IsInf(downloadRateBytesPerSecond, 0) { + return json.Marshal("+Inf bps") + } + + return json.Marshal( + fmt.Sprintf("%sps", units.HumanSizeWithPrecision(downloadRateBytesPerSecond, 2)), + ) +} + +func (dr *downloadRate) UnmarshalJSON(data []byte) error { + var downloadRateStr string + err := json.Unmarshal(data, &downloadRateStr) + if err != nil { + return err + } + + if downloadRateStr == "+Inf bps" { + *dr = downloadRate(math.Inf(1)) + return nil + } + + downloadRateStr = strings.TrimSuffix(downloadRateStr, "ps") + downloadRateBytesPerSecond, err := units.FromHumanSize(downloadRateStr) + if err != nil { + return err } - return json.Marshal(float64(dr)) + *dr = downloadRate(downloadRateBytesPerSecond) + return nil } diff --git a/internal/pkg/agent/application/upgrade/details/details_test.go b/internal/pkg/agent/application/upgrade/details/details_test.go index c1d3056e08e..88d239dfabf 100644 --- a/internal/pkg/agent/application/upgrade/details/details_test.go +++ b/internal/pkg/agent/application/upgrade/details/details_test.go @@ -74,7 +74,7 @@ func TestDetailsDownloadRateJSON(t *testing.T) { var unmarshalledDetails Details err = json.Unmarshal(data, &unmarshalledDetails) require.NoError(t, err) - require.Equal(t, 1794.7, float64(unmarshalledDetails.Metadata.DownloadRate)) + require.Equal(t, float64(1800), float64(unmarshalledDetails.Metadata.DownloadRate)) require.Equal(t, .8, unmarshalledDetails.Metadata.DownloadPercent) }) @@ -88,7 +88,7 @@ func TestDetailsDownloadRateJSON(t *testing.T) { var unmarshalledDetails Details err = json.Unmarshal(data, &unmarshalledDetails) require.NoError(t, err) - require.Equal(t, -1.0, float64(unmarshalledDetails.Metadata.DownloadRate)) + require.Equal(t, math.Inf(1), float64(unmarshalledDetails.Metadata.DownloadRate)) require.Equal(t, 0.99, unmarshalledDetails.Metadata.DownloadPercent) })