github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/utils/progress.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package utils
     4  
     5  import (
     6  	"context"
     7  	"encoding/json"
     8  	"io"
     9  	"sync/atomic"
    10  	"time"
    11  
    12  	"github.com/cheggaaa/pb/v3"
    13  	"github.com/pingcap/errors"
    14  	"github.com/pingcap/log"
    15  	"go.uber.org/zap"
    16  )
    17  
    18  type logFunc func(msg string, fields ...zap.Field)
    19  
    20  // ProgressPrinter prints a progress bar.
    21  type ProgressPrinter struct {
    22  	name        string
    23  	total       int64
    24  	redirectLog bool
    25  	progress    int64
    26  
    27  	cancel context.CancelFunc
    28  }
    29  
    30  // NewProgressPrinter returns a new progress printer.
    31  func NewProgressPrinter(
    32  	name string,
    33  	total int64,
    34  	redirectLog bool,
    35  ) *ProgressPrinter {
    36  	return &ProgressPrinter{
    37  		name:        name,
    38  		total:       total,
    39  		redirectLog: redirectLog,
    40  		cancel: func() {
    41  			log.Warn("canceling non-started progress printer")
    42  		},
    43  	}
    44  }
    45  
    46  // Inc increases the current progress bar.
    47  func (pp *ProgressPrinter) Inc() {
    48  	atomic.AddInt64(&pp.progress, 1)
    49  }
    50  
    51  // Close closes the current progress bar.
    52  func (pp *ProgressPrinter) Close() {
    53  	pp.cancel()
    54  }
    55  
    56  // goPrintProgress starts a gorouinte and prints progress.
    57  func (pp *ProgressPrinter) goPrintProgress(
    58  	ctx context.Context,
    59  	logFuncImpl logFunc,
    60  	testWriter io.Writer, // Only for tests
    61  ) {
    62  	cctx, cancel := context.WithCancel(ctx)
    63  	pp.cancel = cancel
    64  	bar := pb.New64(pp.total)
    65  	if pp.redirectLog || testWriter != nil {
    66  		tmpl := `{"P":"{{percent .}}","C":"{{counters . }}","E":"{{etime .}}","R":"{{rtime .}}","S":"{{speed .}}"}`
    67  		bar.SetTemplateString(tmpl)
    68  		bar.SetRefreshRate(2 * time.Minute)
    69  		bar.Set(pb.Static, false)       // Do not update automatically
    70  		bar.Set(pb.ReturnSymbol, false) // Do not append '\r'
    71  		bar.Set(pb.Terminal, false)     // Do not use terminal width
    72  		// Hack! set Color to avoid separate progress string
    73  		bar.Set(pb.Color, true)
    74  		if logFuncImpl == nil {
    75  			logFuncImpl = log.Info
    76  		}
    77  		bar.SetWriter(&wrappedWriter{name: pp.name, log: logFuncImpl})
    78  	} else {
    79  		tmpl := `{{string . "barName" | green}} {{ bar . "<" "-" (cycle . "-" "\\" "|" "/" ) "." ">"}} {{percent .}}`
    80  		bar.SetTemplateString(tmpl)
    81  		bar.Set("barName", pp.name)
    82  	}
    83  	if testWriter != nil {
    84  		bar.SetWriter(testWriter)
    85  		bar.SetRefreshRate(2 * time.Second)
    86  	}
    87  	bar.Start()
    88  
    89  	go func() {
    90  		t := time.NewTicker(time.Second)
    91  		defer t.Stop()
    92  		defer bar.Finish()
    93  
    94  		for {
    95  			select {
    96  			case <-cctx.Done():
    97  				// a hacky way to adapt the old behavior:
    98  				// when canceled by the outer context, leave the progress unchanged.
    99  				// when canceled by Close method (the 'internal' way), push the progress to 100%.
   100  				if ctx.Err() != nil {
   101  					return
   102  				}
   103  				bar.SetCurrent(pp.total)
   104  				return
   105  			case <-t.C:
   106  			}
   107  
   108  			currentProgress := atomic.LoadInt64(&pp.progress)
   109  			if currentProgress <= pp.total {
   110  				bar.SetCurrent(currentProgress)
   111  			} else {
   112  				bar.SetCurrent(pp.total)
   113  			}
   114  		}
   115  	}()
   116  }
   117  
   118  type wrappedWriter struct {
   119  	name string
   120  	log  logFunc
   121  }
   122  
   123  func (ww *wrappedWriter) Write(p []byte) (int, error) {
   124  	var info struct {
   125  		P string
   126  		C string
   127  		E string
   128  		R string
   129  		S string
   130  	}
   131  	if err := json.Unmarshal(p, &info); err != nil {
   132  		return 0, errors.Trace(err)
   133  	}
   134  	ww.log("progress",
   135  		zap.String("step", ww.name),
   136  		zap.String("progress", info.P),
   137  		zap.String("count", info.C),
   138  		zap.String("speed", info.S),
   139  		zap.String("elapsed", info.E),
   140  		zap.String("remaining", info.R))
   141  	return len(p), nil
   142  }
   143  
   144  // StartProgress starts progress bar.
   145  func StartProgress(
   146  	ctx context.Context,
   147  	name string,
   148  	total int64,
   149  	redirectLog bool,
   150  	log logFunc,
   151  ) *ProgressPrinter {
   152  	progress := NewProgressPrinter(name, total, redirectLog)
   153  	progress.goPrintProgress(ctx, log, nil)
   154  	return progress
   155  }