1package progress_bar
 2
 3import (
 4	"fmt"
 5	"os"
 6	"strings"
 7	"sync"
 8	"sync/atomic"
 9	"time"
10)
11
12type ProgressBar struct {
13	label   string
14	total   int64
15	current int64
16	stop    chan struct{}
17	wg      sync.WaitGroup
18}
19
20func NewProgressBar(label string, total int) *ProgressBar {
21	p := &ProgressBar{label: label, total: int64(total), stop: make(chan struct{})}
22	p.wg.Add(1)
23	go func() {
24		defer p.wg.Done()
25		ticker := time.NewTicker(100 * time.Millisecond)
26		defer ticker.Stop()
27		// initial draw
28		p.draw(atomic.LoadInt64(&p.current))
29		for {
30			select {
31			case <-p.stop:
32				return
33			case <-ticker.C:
34				cur := atomic.LoadInt64(&p.current)
35				if cur > p.total {
36					cur = p.total
37				}
38				p.draw(cur)
39			}
40		}
41	}()
42	return p
43}
44
45func (p *ProgressBar) Inc() {
46	for {
47		cur := atomic.LoadInt64(&p.current)
48		if cur >= p.total {
49			return
50		}
51		if atomic.CompareAndSwapInt64(&p.current, cur, cur+1) {
52			return
53		}
54		// retry on race
55	}
56}
57
58func (p *ProgressBar) Done() {
59	atomic.StoreInt64(&p.current, p.total)
60	close(p.stop)
61	p.wg.Wait()
62	p.draw(p.total)
63	_, _ = fmt.Fprintln(os.Stderr)
64}
65
66func (p *ProgressBar) draw(current int64) {
67	if p.total <= 0 {
68		return
69	}
70	percent := 0
71	if p.total > 0 {
72		percent = int(current * 100 / p.total)
73	}
74	barLen := 24
75	filled := 0
76	if p.total > 0 {
77		filled = int(current * int64(barLen) / p.total)
78	}
79	if filled > barLen {
80		filled = barLen
81	}
82	bar := strings.Repeat("#", filled) + strings.Repeat(" ", barLen-filled)
83	_, _ = fmt.Fprintf(os.Stderr, "\r[%s] %4d/%-4d (%3d%%) %s", bar, current, p.total, percent, p.label)
84}