-
Notifications
You must be signed in to change notification settings - Fork 61
/
progress_logger.go
119 lines (104 loc) · 2.89 KB
/
progress_logger.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
// Copyright 2021 VMware, Inc.
// SPDX-License-Identifier: Apache-2.0
package util
import (
"context"
"fmt"
"os"
pb "github.com/cheggaaa/pb/v3"
regv1 "github.com/google/go-containerregistry/pkg/v1"
"github.com/mattn/go-isatty"
)
// ProgressLogger Progress bar
type ProgressLogger interface {
Start(ctx context.Context, progress <-chan regv1.Update)
End()
}
// NewProgressBar constructor to build a ProgressLogger responsible for printing out a progress bar using updates when
// writing to a registry via ggcr
func NewProgressBar(logger LoggerWithLevels, finalMessage, errorMessagePrefix string) ProgressLogger {
if isatty.IsTerminal(os.Stdout.Fd()) {
return &ProgressBarLogger{logger: logger, finalMessage: finalMessage, errorMessagePrefix: errorMessagePrefix}
}
return &ProgressBarNoTTYLogger{logger: logger, finalMessage: finalMessage}
}
// NewNoopProgressBar constructs a Noop Progress bar that will not display anything
func NewNoopProgressBar() ProgressLogger {
return &ProgressBarNoTTYLogger{}
}
// ProgressBarLogger display progress bar on output
type ProgressBarLogger struct {
cancelFunc context.CancelFunc
bar *pb.ProgressBar
logger LoggerWithLevels
finalMessage string
errorMessagePrefix string
}
// Start the display of the Progress Bar
func (l *ProgressBarLogger) Start(ctx context.Context, progressChan <-chan regv1.Update) {
ctx, cancelFunc := context.WithCancel(ctx)
l.cancelFunc = cancelFunc
// Add a new empty line to separate the progress bar from prior output
fmt.Println()
l.bar = pb.New64(0)
l.bar.Set(pb.Bytes, true)
go func() {
for {
select {
case <-ctx.Done():
return
case update := <-progressChan:
if update.Error != nil {
l.logger.Errorf("%s: %s\n", l.errorMessagePrefix, update.Error)
continue
}
if update.Total == 0 {
return
}
if !l.bar.IsStarted() {
l.bar.SetTotal(update.Total)
l.bar.Start()
}
l.bar.SetCurrent(update.Complete)
l.bar.Write()
}
}
}()
}
// End stops the progress bar and writes the final message
func (l *ProgressBarLogger) End() {
if l.cancelFunc != nil {
l.cancelFunc()
}
l.bar.Finish()
l.logger.Logf("\n%s", l.finalMessage)
}
// ProgressBarNoTTYLogger does not display the progress bar
type ProgressBarNoTTYLogger struct {
cancelFunc context.CancelFunc
logger Logger
finalMessage string
}
// Start consuming the progress channel but does not display anything
func (l *ProgressBarNoTTYLogger) Start(ctx context.Context, progressChan <-chan regv1.Update) {
ctx, cancelFunc := context.WithCancel(ctx)
l.cancelFunc = cancelFunc
go func() {
for {
select {
case <-ctx.Done():
return
case <-progressChan:
}
}
}()
}
// End Write the final message
func (l *ProgressBarNoTTYLogger) End() {
if l.cancelFunc != nil {
l.cancelFunc()
}
if l.logger != nil && l.finalMessage != "" {
l.logger.Logf(l.finalMessage)
}
}