summaryrefslogtreecommitdiff
path: root/vendor/github.com/vbauerster/mpb/v4/decor/speed.go
blob: 795a5536f240e863fcefedcddce3852492067f47 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
package decor

import (
	"fmt"
	"io"
	"math"
	"time"

	"github.com/VividCortex/ewma"
)

// SpeedFormatter is wrapper for SizeB1024 and SizeB1000 to format value as speed/s.
type SpeedFormatter struct {
	fmt.Formatter
}

func (self *SpeedFormatter) Format(st fmt.State, verb rune) {
	self.Formatter.Format(st, verb)
	io.WriteString(st, "/s")
}

// EwmaSpeed exponential-weighted-moving-average based speed decorator.
// Note that it's necessary to supply bar.Incr* methods with incremental
// work duration as second argument, in order for this decorator to
// work correctly. This decorator is a wrapper of MovingAverageSpeed.
func EwmaSpeed(unit int, format string, age float64, wcc ...WC) Decorator {
	var average MovingAverage
	if age == 0 {
		average = ewma.NewMovingAverage()
	} else {
		average = ewma.NewMovingAverage(age)
	}
	return MovingAverageSpeed(unit, format, average, wcc...)
}

// MovingAverageSpeed decorator relies on MovingAverage implementation
// to calculate its average.
//
//	`unit` one of [0|UnitKiB|UnitKB] zero for no unit
//
//	`format` printf compatible verb for value, like "%f" or "%d"
//
//	`average` MovingAverage implementation
//
//	`wcc` optional WC config
//
// format examples:
//
//	unit=UnitKiB, format="%.1f"  output: "1.0MiB/s"
//	unit=UnitKiB, format="% .1f" output: "1.0 MiB/s"
//	unit=UnitKB,  format="%.1f"  output: "1.0MB/s"
//	unit=UnitKB,  format="% .1f" output: "1.0 MB/s"
//
func MovingAverageSpeed(unit int, format string, average MovingAverage, wcc ...WC) Decorator {
	var wc WC
	for _, widthConf := range wcc {
		wc = widthConf
	}
	if format == "" {
		format = "%.0f"
	}
	d := &movingAverageSpeed{
		WC:       wc.Init(),
		average:  average,
		producer: chooseSpeedProducer(unit, format),
	}
	return d
}

type movingAverageSpeed struct {
	WC
	producer func(float64) string
	average  ewma.MovingAverage
	msg      string
}

func (d *movingAverageSpeed) Decor(st *Statistics) string {
	if !st.Completed {
		var speed float64
		if v := d.average.Value(); v > 0 {
			speed = 1 / v
		}
		d.msg = d.producer(speed * 1e9)
	}
	return d.FormatMsg(d.msg)
}

func (d *movingAverageSpeed) NextAmount(n int64, wdd ...time.Duration) {
	var workDuration time.Duration
	for _, wd := range wdd {
		workDuration = wd
	}
	durPerByte := float64(workDuration) / float64(n)
	if math.IsInf(durPerByte, 0) || math.IsNaN(durPerByte) {
		return
	}
	d.average.Add(durPerByte)
}

// AverageSpeed decorator with dynamic unit measure adjustment. It's
// a wrapper of NewAverageSpeed.
func AverageSpeed(unit int, format string, wcc ...WC) Decorator {
	return NewAverageSpeed(unit, format, time.Now(), wcc...)
}

// NewAverageSpeed decorator with dynamic unit measure adjustment and
// user provided start time.
//
//	`unit` one of [0|UnitKiB|UnitKB] zero for no unit
//
//	`format` printf compatible verb for value, like "%f" or "%d"
//
//	`startTime` start time
//
//	`wcc` optional WC config
//
// format examples:
//
//	unit=UnitKiB, format="%.1f"  output: "1.0MiB/s"
//	unit=UnitKiB, format="% .1f" output: "1.0 MiB/s"
//	unit=UnitKB,  format="%.1f"  output: "1.0MB/s"
//	unit=UnitKB,  format="% .1f" output: "1.0 MB/s"
//
func NewAverageSpeed(unit int, format string, startTime time.Time, wcc ...WC) Decorator {
	var wc WC
	for _, widthConf := range wcc {
		wc = widthConf
	}
	if format == "" {
		format = "%.0f"
	}
	d := &averageSpeed{
		WC:        wc.Init(),
		startTime: startTime,
		producer:  chooseSpeedProducer(unit, format),
	}
	return d
}

type averageSpeed struct {
	WC
	startTime time.Time
	producer  func(float64) string
	msg       string
}

func (d *averageSpeed) Decor(st *Statistics) string {
	if !st.Completed {
		speed := float64(st.Current) / float64(time.Since(d.startTime))
		d.msg = d.producer(speed * 1e9)
	}

	return d.FormatMsg(d.msg)
}

func (d *averageSpeed) AverageAdjust(startTime time.Time) {
	d.startTime = startTime
}

func chooseSpeedProducer(unit int, format string) func(float64) string {
	switch unit {
	case UnitKiB:
		return func(speed float64) string {
			return fmt.Sprintf(format, &SpeedFormatter{SizeB1024(math.Round(speed))})
		}
	case UnitKB:
		return func(speed float64) string {
			return fmt.Sprintf(format, &SpeedFormatter{SizeB1000(math.Round(speed))})
		}
	default:
		return func(speed float64) string {
			return fmt.Sprintf(format, speed)
		}
	}
}