summaryrefslogtreecommitdiff
path: root/src/mongo/gotools/common/progress/manager.go
blob: e1c5f0db5f7370ae0fd840309d1d03f5c0575d6e (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
package progress

import (
	"fmt"
	"github.com/mongodb/mongo-tools/common/text"
	"io"
	"sync"
	"time"
)

const GridPadding = 2

// Manager handles thread-safe synchronized progress bar writing, so that all
// given progress bars are written in a group at a given interval.
// The current implementation maintains insert order when printing,
// such that new bars appear at the bottom of the group.
type Manager struct {
	waitTime time.Duration
	writer   io.Writer
	bars     []*Bar
	barsLock *sync.Mutex
	stopChan chan struct{}
}

// NewProgressBarManager returns an initialized Manager with the given
// time.Duration to wait between writes
func NewProgressBarManager(w io.Writer, waitTime time.Duration) *Manager {
	return &Manager{
		waitTime: waitTime,
		writer:   w,
		barsLock: &sync.Mutex{},
		stopChan: make(chan struct{}),
	}
}

// Attach registers the given progress bar with the manager. Should be used as
//  myManager.Attach(myBar)
//  defer myManager.Detach(myBar)
func (manager *Manager) Attach(pb *Bar) {
	// first some quick error checks
	if pb.Name == "" {
		panic("cannot attach a nameless bar to a progress bar manager")
	}
	pb.validate()

	manager.barsLock.Lock()
	defer manager.barsLock.Unlock()

	// make sure we are not adding the same bar again
	for _, bar := range manager.bars {
		if bar.Name == pb.Name {
			panic(fmt.Sprintf("progress bar with name '%v' already exists in manager", pb.Name))
		}
	}

	manager.bars = append(manager.bars, pb)
}

// Detach removes the given progress bar from the manager.
// Insert order is maintained for consistent ordering of the printed bars.
//  Note: the manager removes progress bars by "Name" not by memory location
func (manager *Manager) Detach(pb *Bar) {
	if pb.Name == "" {
		panic("cannot detach a nameless bar from a progress bar manager")
	}

	manager.barsLock.Lock()
	defer manager.barsLock.Unlock()

	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 *Manager) renderAllBars() {
	manager.barsLock.Lock()
	defer manager.barsLock.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 *Manager) Start() {
	if manager.writer == nil {
		panic("Cannot use a progress.Manager with an unset Writer")
	}
	go manager.start()
}

func (manager *Manager) 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 *Manager) Stop() {
	manager.stopChan <- struct{}{}
}