package observe import ( "context" "fmt" "io" "os" "strings" "sync" "github.com/alcionai/clues" "github.com/dustin/go-humanize" "github.com/spf13/cobra" "github.com/spf13/pflag" "github.com/vbauerster/mpb/v8" "github.com/vbauerster/mpb/v8/decor" "github.com/alcionai/corso/src/internal/common/color" "github.com/alcionai/corso/src/pkg/logger" ) const ( hideProgressBarsFN = "hide-progress" retainProgressBarsFN = "retain-progress" progressBarWidth = 40 ) func init() { makeSpinFrames(progressBarWidth) } // adds the persistent boolean flag --hide-progress to the provided command. // This is a hack for help displays. Due to seeding the context, we also // need to parse the configuration before we execute the command. func AddProgressBarFlags(cmd *cobra.Command) { fs := cmd.PersistentFlags() fs.Bool(hideProgressBarsFN, false, "turn off the progress bar displays") fs.Bool(retainProgressBarsFN, false, "retain the progress bar displays after completion") } // Due to races between the lazy evaluation of flags in cobra and the need to init observer // behavior in a ctx, these options get pre-processed manually here using pflags. The canonical // AddProgressBarFlag() ensures the flags are displayed as part of the help/usage output. func PreloadFlags() config { fs := pflag.NewFlagSet("seed-observer", pflag.ContinueOnError) fs.ParseErrorsWhitelist.UnknownFlags = true fs.Bool(hideProgressBarsFN, false, "turn off the progress bar displays") fs.Bool(retainProgressBarsFN, false, "retain the progress bar displays after completion") // prevents overriding the corso/cobra help processor fs.BoolP("help", "h", false, "") // parse the os args list to find the observer display flags if err := fs.Parse(os.Args[1:]); err != nil { return config{} } // retrieve the user's preferred display // automatically defaults to "info" shouldHide, err := fs.GetBool(hideProgressBarsFN) if err != nil { return config{} } // retrieve the user's preferred display // automatically defaults to "info" shouldAlwaysShow, err := fs.GetBool(retainProgressBarsFN) if err != nil { return config{} } return config{ doNotDisplay: shouldHide, displayIsTerminal: true, keepBarsAfterComplete: shouldAlwaysShow, } } // --------------------------------------------------------------------------- // configuration // --------------------------------------------------------------------------- // config handles observer configuration type config struct { // under certain conditions (ex: testing) we aren't outputting // to a terminal. When this happens the observe bars need to be // given a specific optional value or they'll never flush the // writer. displayIsTerminal bool doNotDisplay bool keepBarsAfterComplete bool } type observerKey string const ctxKey observerKey = "corsoObserver" type observer struct { cfg config mp *mpb.Progress w io.Writer wg *sync.WaitGroup } func (o observer) hidden() bool { return o.cfg.doNotDisplay || o.w == nil } func (o *observer) resetWriter(ctx context.Context) { opts := []mpb.ContainerOption{ mpb.WithWidth(progressBarWidth), mpb.WithWaitGroup(o.wg), mpb.WithOutput(o.w), } // needed for tests to pass with mpb versions > 8.1.6 // if !o.cfg.displayIsTerminal { // opts = append(opts, mpb.WithAutoRefresh()) // } o.mp = mpb.NewWithContext(ctx, opts...) } // SeedObserver adds an observer to the context. Any calls to observe // funcs will retrieve the observer from the context. If no observer // is found in the context, the call no-ops. func SeedObserver(ctx context.Context, w io.Writer, cfg config) context.Context { obs := &observer{ w: w, cfg: cfg, wg: &sync.WaitGroup{}, } obs.resetWriter(ctx) return setObserver(ctx, obs) } func setObserver(ctx context.Context, obs *observer) context.Context { return context.WithValue(ctx, ctxKey, obs) } func getObserver(ctx context.Context) *observer { o := ctx.Value(ctxKey) if o == nil { return &observer{cfg: config{doNotDisplay: true}} } return o.(*observer) } // Flush blocks until the progress finishes writing out all data. // Afterwards, the progress instance is reset. func Flush(ctx context.Context) { obs := getObserver(ctx) if obs.mp != nil { obs.mp.Wait() } obs.resetWriter(ctx) } const ( ItemBackupMsg = "Backing up item" ItemRestoreMsg = "Restoring item" ItemExportMsg = "Exporting item" ItemQueueMsg = "Queuing items" ) // --------------------------------------------------------------------------- // Progress Updates // --------------------------------------------------------------------------- type ProgressCfg struct { NewSection bool SectionIdentifier any Indent int CompletionMessage func() string } func DefaultCfg() ProgressCfg { return ProgressCfg{} } // Message is used to display a progress message func Message(ctx context.Context, cfg ProgressCfg, msgs ...any) { var ( obs = getObserver(ctx) plainSl = make([]string, 0, len(msgs)) loggableSl = make([]string, 0, len(msgs)) ) for _, m := range msgs { plainSl = append(plainSl, plainString(m)) loggableSl = append(loggableSl, fmt.Sprintf("%v", m)) } plain := strings.Join(plainSl, " ") loggable := strings.Join(loggableSl, " ") if cfg.SectionIdentifier != nil { plain = fmt.Sprintf( "%s %s %s", plain, Bullet, color.MagentaOutput(plainString(cfg.SectionIdentifier))) loggable = fmt.Sprintf("%s - %v", loggable, cfg.SectionIdentifier) } logger.Ctx(ctx).Info(loggable) if obs.hidden() { return } if cfg.NewSection { // Empty bar to separate out section obs.wg.Add(1) empty := obs.mp.New(-1, mpb.NopStyle()) empty.SetTotal(-1, true) waitAndCloseBar(ctx, empty, obs.wg, func() {})() } obs.wg.Add(1) bar := obs.mp.New( -1, mpb.NopStyle(), mpb.PrependDecorators( decor.Name("", decor.WC{W: cfg.Indent * 2}), decor.Name( plain, decor.WC{ W: len(plain) + 1, C: decor.DidentRight, }))) // Complete the bar immediately bar.SetTotal(-1, true) waitAndCloseBar(ctx, bar, obs.wg, func() {})() } // MessageWithCompletion is used to display progress with a spinner // that switches to "done" when the completion channel is signalled func MessageWithCompletion( ctx context.Context, cfg ProgressCfg, msgs ...any, ) chan<- struct{} { var ( obs = getObserver(ctx) log = logger.Ctx(ctx) plainSl = make([]string, 0, len(msgs)) loggableSl = make([]string, 0, len(msgs)) ch = make(chan struct{}, 1) ) for _, m := range msgs { plainSl = append(plainSl, plainString(m)) loggableSl = append(loggableSl, fmt.Sprintf("%v", m)) } plain := strings.Join(plainSl, " ") loggable := strings.Join(loggableSl, " ") if cfg.SectionIdentifier != nil { plain = fmt.Sprintf( "%s %s %s", plain, Bullet, color.MagentaOutput(plainString(cfg.SectionIdentifier))) loggable = fmt.Sprintf("%s - %v", loggable, cfg.SectionIdentifier) } if cfg.Indent > 0 { plain = color.CyanOutput(plain) } log.Info(loggable) if obs.hidden() { defer log.Info("done - " + loggable) return ch } if cfg.NewSection { // Empty bar to separate out section obs.wg.Add(1) empty := obs.mp.New(-1, mpb.NopStyle()) empty.SetTotal(-1, true) waitAndCloseBar(ctx, empty, obs.wg, func() {})() } obs.wg.Add(1) frames := []string{"∙∙∙", "●∙∙", "∙●∙", "∙∙●", "∙∙∙"} // https://github.com/vbauerster/mpb/issues/71 bfoc := mpb.BarFillerOnComplete(color.GreenOutput("done")) if cfg.CompletionMessage != nil { bfoc = mpb.BarFillerMiddleware(func(base mpb.BarFiller) mpb.BarFiller { filler := func(w io.Writer, st decor.Statistics) error { if st.Completed { msg := fmt.Sprintf("%s %s", color.GreenOutput("done"), color.GreyOutput(cfg.CompletionMessage())) _, err := io.WriteString(w, msg) return err } return base.Fill(w, st) } return mpb.BarFillerFunc(filler) }) } bar := obs.mp.New( -1, mpb.SpinnerStyle(frames...).PositionLeft(), mpb.PrependDecorators( decor.Name("", decor.WC{W: cfg.Indent * 2}), decor.Name(plain, decor.WC{W: progressBarWidth - cfg.Indent*2, C: decor.DidentRight}), decor.Elapsed(decor.ET_STYLE_GO, decor.WC{W: 8})), bfoc) go listen( ctx, ch, func() { bar.SetTotal(-1, true) bar.Abort(true) }, // callers should close the channel func() {}) go waitAndCloseBar(ctx, bar, obs.wg, func() { log.Info("done - " + loggable) })() return ch } // --------------------------------------------------------------------------- // Progress for Known Quantities // --------------------------------------------------------------------------- type autoCloser struct { rc io.ReadCloser close func() closed bool } func (ac *autoCloser) Read(p []byte) (n int, err error) { return ac.rc.Read(p) } func (ac *autoCloser) Close() error { if ac.closed { return nil } ac.closed = true ac.close() return ac.rc.Close() } // ItemProgress tracks the display of an item in a folder by counting the bytes // read through the provided readcloser, up until the byte count matches // the totalBytes. // // The progress bar will close automatically when the reader closes. If an early // close is needed due to abort or other issue, the reader can be closed manually. func ItemProgress( ctx context.Context, rc io.ReadCloser, header string, iname any, totalBytes int64, ) io.ReadCloser { var ( obs = getObserver(ctx) plain = plainString(iname) log = logger.Ctx(ctx).With( "item", iname, "size", humanize.Bytes(uint64(totalBytes))) ) log.Debug(header) if obs.hidden() || rc == nil { defer log.Debug("done - " + header) return rc } obs.wg.Add(1) barOpts := []mpb.BarOption{ mpb.PrependDecorators( decor.Name(header, decor.WCSyncSpaceR), decor.Name(plain, decor.WCSyncSpaceR), decor.CountersKibiByte(" %.1f/%.1f ", decor.WC{W: 8}), decor.NewPercentage("%d ", decor.WC{W: 4})), } if !obs.cfg.keepBarsAfterComplete { barOpts = append(barOpts, mpb.BarRemoveOnComplete()) } bar := obs.mp.New(totalBytes, mpb.NopStyle(), barOpts...) go waitAndCloseBar(ctx, bar, obs.wg, func() { // might be overly chatty, we can remove if needed. log.Debug("done - " + header) })() closer := &autoCloser{rc: bar.ProxyReader(rc)} closer.close = func() { bar.SetTotal(-1, true) bar.Abort(true) } return closer } // ItemSpinner is similar to ItemProgress, but for use in cases where // we don't know the file size but want to show progress. // // The progress bar will close automatically when the reader closes. If an early // close is needed due to abort or other issue, the reader can be closed manually. func ItemSpinner( ctx context.Context, rc io.ReadCloser, header string, iname any, ) io.ReadCloser { var ( obs = getObserver(ctx) plain = plainString(iname) log = logger.Ctx(ctx).With("item", iname) ) log.Debug(header) if obs.hidden() || rc == nil { defer log.Debug("done - " + header) return rc } obs.wg.Add(1) barOpts := []mpb.BarOption{ mpb.PrependDecorators( decor.Name(header, decor.WCSyncSpaceR), decor.Name(plain, decor.WCSyncSpaceR), decor.CurrentKibiByte(" %.1f", decor.WC{W: 8})), } if !obs.cfg.keepBarsAfterComplete { barOpts = append(barOpts, mpb.BarRemoveOnComplete()) } bar := obs.mp.New(-1, mpb.NopStyle(), barOpts...) go waitAndCloseBar(ctx, bar, obs.wg, func() { // might be overly chatty, we can remove if needed. log.Debug("done - " + header) })() closer := &autoCloser{rc: bar.ProxyReader(rc)} closer.close = func() { bar.SetTotal(-1, true) bar.Abort(true) } return closer } // ProgressWithCount tracks the display of a bar that tracks the completion // of the specified count. // Each write to the provided channel counts as a single increment. // The caller is expected to close the channel. func ProgressWithCount( ctx context.Context, header string, msg any, count int64, ) chan<- struct{} { var ( obs = getObserver(ctx) plain = plainString(msg) loggable = fmt.Sprintf("%s %v - %d", header, msg, count) log = logger.Ctx(ctx) ch = make(chan struct{}) ) log.Info(loggable) if obs.hidden() { go listen(ctx, ch, nop, nop) defer log.Info("done - " + loggable) return ch } obs.wg.Add(1) barOpts := []mpb.BarOption{ mpb.PrependDecorators( decor.Name(header, decor.WCSyncSpaceR), decor.Name(plain), decor.Counters(0, " %d/%d ")), } if !obs.cfg.keepBarsAfterComplete { barOpts = append(barOpts, mpb.BarRemoveOnComplete()) } bar := obs.mp.New(count, mpb.NopStyle(), barOpts...) go listen( ctx, ch, func() { bar.Abort(true) }, bar.Increment) go waitAndCloseBar(ctx, bar, obs.wg, func() { log.Info("done - " + loggable) })() return ch } // --------------------------------------------------------------------------- // Progress for Unknown Quantities // --------------------------------------------------------------------------- var spinFrames []string // The bar width is set to a static 32 characters. The default spinner is only // one char wide, which puts a lot of white space between it and the useful text. // This builds a custom spinner animation to fill up that whitespace for a cleaner // display. func makeSpinFrames(barWidth int) { s, l := rune('∙'), rune('●') line := []rune{} for i := 0; i < barWidth; i++ { line = append(line, s) } sl := make([]string, 0, barWidth+1) sl = append(sl, string(line)) for i := 1; i < barWidth; i++ { l2 := make([]rune, len(line)) copy(l2, line) l2[i] = l sl = append(sl, string(l2)) } spinFrames = sl } // CollectionProgress tracks the display a spinner that idles while the collection // incrementing the count of items handled. Each write to the provided channel // counts as a single increment. The caller is expected to close the channel. func CollectionProgress( ctx context.Context, category string, dirName any, ) chan<- struct{} { var ( obs = getObserver(ctx) counted int plain = plainString(dirName) ch = make(chan struct{}) log = logger.Ctx(ctx).With( "category", category, "dir", dirName) message = "Collecting Directory" ) log.Info(message) incCount := func() { counted++ // Log every 1000 items that are processed if counted%1000 == 0 { log.Infow("uploading", "count_items", counted) } } if obs.hidden() || len(plain) == 0 { go listen(ctx, ch, nop, incCount) defer log.Infow("done - "+message, "count_items", counted) return ch } obs.wg.Add(1) barOpts := []mpb.BarOption{ mpb.PrependDecorators( decor.Name("", decor.WC{W: 2}), decor.Name(color.CyanOutput(string(category)))), mpb.AppendDecorators( decor.CurrentNoUnit("%d - ", decor.WCSyncSpace), decor.Name(plain)), mpb.BarFillerOnComplete(spinFrames[0]), } if !obs.cfg.keepBarsAfterComplete { barOpts = append(barOpts, mpb.BarRemoveOnComplete()) } bar := obs.mp.New( -1, // -1 to indicate an unbounded count mpb.SpinnerStyle(spinFrames...), barOpts...) go listen( ctx, ch, func() { bar.SetTotal(-1, true) }, func() { incCount() bar.Increment() }) go waitAndCloseBar(ctx, bar, obs.wg, func() { log.Infow("done - "+message, "count_items", counted) })() return ch } func waitAndCloseBar(ctx context.Context, bar *mpb.Bar, wg *sync.WaitGroup, log func()) func() { return func() { bar.Wait() wg.Done() if !bar.Aborted() { log() } } } // --------------------------------------------------------------------------- // other funcs // --------------------------------------------------------------------------- var nop = func() {} // listen handles reading, and exiting, from a channel. It assumes the // caller will run it inside a goroutine (ex: go listen(...)). // On context timeout or channel close, the loop exits. // onEnd() is called on both ctx.Done() and channel close. onInc is // called on every channel read except when closing. func listen(ctx context.Context, ch <-chan struct{}, onEnd, onInc func()) { for { select { case <-ctx.Done(): onEnd() return case _, ok := <-ch: if !ok { onEnd() return } onInc() } } } // --------------------------------------------------------------------------- // Styling // --------------------------------------------------------------------------- const Bullet = "∙" type bulletf struct { tmpl string vs []any } func Bulletf(template string, vs ...any) bulletf { return bulletf{template, vs} } func (b bulletf) PlainString() string { ps := make([]any, 0, len(b.vs)) for _, v := range b.vs { ps = append(ps, plainString(v)) } return fmt.Sprintf("∙ "+b.tmpl, ps...) } func (b bulletf) String() string { return fmt.Sprintf("∙ "+b.tmpl, b.vs...) } // plainString attempts to cast v to a PlainStringer // interface, and retrieve the un-altered value. If // v is not compliant with PlainStringer, returns the // %v fmt of v. // // This should only be used to display the value in the // observe progress bar. Logged values should only use // the fmt %v to ensure Concealers hide PII. func plainString(v any) string { if c, ok := v.(clues.Concealer); ok { return c.PlainString() } return fmt.Sprintf("%v", v) }