forked from mongodb/mongo-tools
-
Notifications
You must be signed in to change notification settings - Fork 0
/
manager.go
166 lines (144 loc) · 4.27 KB
/
manager.go
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
// Copyright (C) MongoDB, Inc. 2014-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
package progress
import (
"fmt"
"io"
"sync"
"time"
"github.com/huimingz/mongo-tools/common/text"
)
// Manager is an interface which tools can use to registers progressors which
// track the progress of any arbitrary operation.
type Manager interface {
// Attach registers the progressor with the manager under the given name.
// Any call to Attach must have a matching call to Detach.
Attach(name string, progressor Progressor)
// Detach removes the progressor with the given name from the manager
Detach(name string)
}
const GridPadding = 2
// BarWriter implements Manager. It periodically prints the status of all of its
// progressors in the form of pretty progress bars. It handles thread-safe
// synchronized progress bar writing, so that its progressors are written in a
// group at a given interval. It maintains insertion order when printing, such
// that new bars appear at the bottom of the group.
type BarWriter struct {
sync.Mutex
waitTime time.Duration
writer io.Writer
bars []*Bar
stopChan chan struct{}
barLength int
isBytes bool
}
// NewBarWriter returns an initialized BarWriter with the given bar length and
// byte-formatting toggle, waiting the given duration between writes
func NewBarWriter(w io.Writer, waitTime time.Duration, barLength int, isBytes bool) *BarWriter {
return &BarWriter{
waitTime: waitTime,
writer: w,
stopChan: make(chan struct{}),
barLength: barLength,
isBytes: isBytes,
}
}
// Attach registers the given progressor with the manager
func (manager *BarWriter) Attach(name string, progressor Progressor) {
pb := &Bar{
Name: name,
Watching: progressor,
BarLength: manager.barLength,
IsBytes: manager.isBytes,
}
pb.validate()
manager.Lock()
defer manager.Unlock()
// make sure we are not adding the same bar again
for _, bar := range manager.bars {
if bar.Name == name {
panic(fmt.Sprintf("progress bar with name '%s' already exists in manager", name))
}
}
manager.bars = append(manager.bars, pb)
}
// Detach removes the progressor with the given name from the manager. Insert
// order is maintained for consistent ordering of the printed bars.
func (manager *BarWriter) Detach(name string) {
manager.Lock()
defer manager.Unlock()
var pb *Bar
for _, bar := range manager.bars {
if bar.Name == name {
pb = bar
break
}
}
if pb == nil {
panic("could not find progressor")
}
grid := &text.GridWriter{
ColumnPadding: GridPadding,
}
if pb.hasRendered {
// if we've rendered this bar at least once, render it one last time
pb.renderToGridRow(grid)
}
grid.FlushRows(manager.writer)
updatedBars := make([]*Bar, 0, len(manager.bars)-1)
for _, bar := range manager.bars {
// move all bars to the updated list except for the bar we want to detach
if bar.Name != pb.Name {
updatedBars = append(updatedBars, bar)
}
}
manager.bars = updatedBars
}
// helper to render all bars in order
func (manager *BarWriter) renderAllBars() {
manager.Lock()
defer manager.Unlock()
grid := &text.GridWriter{
ColumnPadding: GridPadding,
}
for _, bar := range manager.bars {
bar.renderToGridRow(grid)
}
grid.FlushRows(manager.writer)
// add padding of one row if we have more than one active bar
if len(manager.bars) > 1 {
// we just write an empty array here, since a write call of any
// length to our log.Writer will trigger a new logline.
manager.writer.Write([]byte{})
}
}
// Start kicks of the timed batch writing of progress bars.
func (manager *BarWriter) Start() {
if manager.writer == nil {
panic("Cannot use a progress.BarWriter with an unset Writer")
}
go manager.start()
}
func (manager *BarWriter) start() {
if manager.waitTime <= 0 {
manager.waitTime = DefaultWaitTime
}
ticker := time.NewTicker(manager.waitTime)
defer ticker.Stop()
for {
select {
case <-manager.stopChan:
return
case <-ticker.C:
manager.renderAllBars()
}
}
}
// Stop ends the main manager goroutine, stopping the manager's bars
// from being rendered.
func (manager *BarWriter) Stop() {
manager.stopChan <- struct{}{}
}