1
2
3 package download
4
5 import (
6 "fmt"
7 "log"
8 "strconv"
9 "time"
10
11 "github.com/vbauerster/mpb/v5"
12 "github.com/vbauerster/mpb/v5/decor"
13 "gitlab.com/tslocum/gophast/pkg/config"
14 . "gitlab.com/tslocum/gophast/pkg/utils"
15 )
16
17 const (
18 RefreshRate = 120 * time.Millisecond
19 SpeedAverageWindow = 45.0
20 BarStyle = "|-> |"
21 )
22
23 var (
24 p *mpb.Progress
25 bars []*mpb.Bar
26 )
27
28 func InitializeProgress(width int) {
29 if p != nil {
30 return
31 }
32
33 if config.C.ProgressLevel == config.ProgressDynamic {
34 p = mpb.New(mpb.WithWidth(width), mpb.WithRefreshRate(RefreshRate), mpb.WithWaitGroup(config.WG))
35 }
36 }
37
38 func AddLocalProgressBar(index int, pieceCount int, r *ByteRange) {
39 if config.C.ProgressLevel != config.ProgressDynamic {
40 return
41 }
42 if p == nil {
43 log.Fatal("progress was not initialized")
44 }
45
46 if pieceCount <= 1 {
47 return
48 }
49
50 task := fmt.Sprintf(config.C.TaskFormat, index+1)
51 if config.C.TaskWidth > len(task) {
52 task = fmt.Sprintf("%"+strconv.Itoa(config.C.TaskWidth)+"s", task)
53 }
54 bar := p.AddBar(r.End-r.Start+1,
55 mpb.BarStyle(BarStyle),
56 mpb.PrependDecorators(
57 decor.Name(task, decor.WC{W: config.C.TaskWidth, C: decor.DidentRight})),
58 mpb.AppendDecorators(
59 decor.OnComplete(decor.Percentage(decor.WC{W: 5}), "100 % 00:00:00"),
60 decor.OnComplete(decor.EwmaETA(decor.ET_STYLE_HHMMSS, SpeedAverageWindow, decor.WCSyncWidth, decor.WC{W: 9}), "")),
61 )
62
63 bars = append(bars, bar)
64 go handleLocalProgressBar(bar, r)
65 }
66
67 func handleLocalProgressBar(bar *mpb.Bar, r *ByteRange) {
68 var (
69 ticker = time.NewTicker(RefreshRate)
70
71 lastProgress int64
72 lastUpdate time.Time
73 )
74
75 lastUpdate = time.Now()
76 for {
77 select {
78 case <-ticker.C:
79 }
80
81 if r == nil {
82 return
83 } else if r.Progress != lastProgress {
84 bar.IncrInt64(r.Progress - lastProgress)
85 bar.DecoratorEwmaUpdate(time.Now().Sub(lastUpdate))
86
87 lastProgress = r.Progress
88 lastUpdate = time.Now()
89 }
90
91 }
92 }
93
94 func AddGlobalProgressBar(d *Download) {
95 if config.C.ProgressLevel != config.ProgressDynamic {
96 return
97 }
98 if p == nil {
99 log.Fatal("progress was not initialized")
100 }
101
102 if len(d.Ranges) > 1 {
103
104 p.AddBar(0, mpb.BarStyle(" ")).SetTotal(0, true)
105 }
106
107 bar := p.AddBar(d.Size,
108 mpb.BarStyle(BarStyle),
109 mpb.PrependDecorators(
110 decor.Name(Ellipsize(d.Name), decor.WC{W: config.C.TaskWidth, C: decor.DidentRight})),
111 mpb.AppendDecorators(
112 decor.OnComplete(decor.Percentage(decor.WC{W: 5}), "100 % 00:00:00"),
113 decor.OnComplete(decor.EwmaETA(decor.ET_STYLE_HHMMSS, SpeedAverageWindow, decor.WCSyncWidth, decor.WC{W: 9}), "")),
114 )
115
116 bars = append(bars, bar)
117
118 go handleGlobalProgressBar(bar, d)
119 }
120
121 func handleGlobalProgressBar(bar *mpb.Bar, d *Download) {
122 var (
123 ticker = time.NewTicker(RefreshRate)
124
125 progress int64
126 lastProgress int64
127 lastUpdate time.Time
128 )
129
130 lastUpdate = time.Now()
131 for {
132 select {
133 case <-ticker.C:
134 }
135
136 if d == nil || d.Ranges == nil {
137 return
138 }
139
140 progress = d.Size - d.Remaining
141 for _, r := range d.Ranges {
142 progress += r.Wrote
143 }
144
145 if progress != lastProgress {
146 bar.IncrInt64(progress - lastProgress)
147 bar.DecoratorEwmaUpdate(time.Now().Sub(lastUpdate))
148
149 lastProgress = progress
150 lastUpdate = time.Now()
151 }
152 }
153 }
154
155 func ShutdownProgress(success bool) {
156 if p == nil {
157 return
158 }
159
160 for _, b := range bars {
161 if success {
162 b.SetTotal(1, true)
163 } else {
164 b.Abort(false)
165 }
166 }
167
168 p.Wait()
169 }
170
View as plain text