summary/: export new log collector (#602)

* export log exporter * make sure dumpling can redirect log in progress.go
pingcap · Nov 19, 2020 · bb4294d · bb4294d
1 parent e3faf51
commit bb4294d
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 11 deletions.
diff --git a/pkg/gluetikv/glue.go b/pkg/gluetikv/glue.go
@@ -48,7 +48,7 @@ func (Glue) OwnsStorage() bool {
 
 // StartProgress implements glue.Glue.
 func (Glue) StartProgress(ctx context.Context, cmdName string, total int64, redirectLog bool) glue.Progress {
-	return utils.StartProgress(ctx, cmdName, total, redirectLog)
+	return utils.StartProgress(ctx, cmdName, total, redirectLog, nil)
 }
 
 // Record implements glue.Glue.

diff --git a/pkg/summary/collector.go b/pkg/summary/collector.go
@@ -44,7 +44,7 @@ type LogCollector interface {
 
 type logFunc func(msg string, fields ...zap.Field)
 
-var collector LogCollector = newLogCollector(log.Info)
+var collector LogCollector = NewLogCollector(log.Info)
 
 // InitCollector initilize global collector instance.
 func InitCollector( // revive:disable-line:flag-parameter
@@ -62,7 +62,7 @@ func InitCollector( // revive:disable-line:flag-parameter
 			}
 		}
 	}
-	collector = newLogCollector(logF)
+	collector = NewLogCollector(logF)
 }
 
 type logCollector struct {
@@ -82,7 +82,8 @@ type logCollector struct {
 	log logFunc
 }
 
-func newLogCollector(log logFunc) LogCollector {
+// NewLogCollector returns a new LogCollector.
+func NewLogCollector(log logFunc) LogCollector {
 	return &logCollector{
 		successUnitCount: 0,
 		failureUnitCount: 0,

diff --git a/pkg/summary/collector_test.go b/pkg/summary/collector_test.go
@@ -24,7 +24,7 @@ func (suit *testCollectorSuite) TestSumDurationInt(c *C) {
 	logger := func(msg string, fs ...zap.Field) {
 		fields = append(fields, fs...)
 	}
-	col := newLogCollector(logger)
+	col := NewLogCollector(logger)
 	col.CollectDuration("a", time.Second)
 	col.CollectDuration("b", time.Second)
 	col.CollectDuration("b", time.Second)

diff --git a/pkg/utils/progress.go b/pkg/utils/progress.go
@@ -14,6 +14,8 @@ import (
 	"go.uber.org/zap"
 )
 
+type logFunc func(msg string, fields ...zap.Field)
+
 // ProgressPrinter prints a progress bar.
 type ProgressPrinter struct {
 	name        string
@@ -53,6 +55,7 @@ func (pp *ProgressPrinter) Close() {
 // goPrintProgress starts a gorouinte and prints progress.
 func (pp *ProgressPrinter) goPrintProgress(
 	ctx context.Context,
+	logFuncImpl logFunc,
 	testWriter io.Writer, // Only for tests
 ) {
 	cctx, cancel := context.WithCancel(ctx)
@@ -67,7 +70,10 @@ func (pp *ProgressPrinter) goPrintProgress(
 		bar.Set(pb.Terminal, false)     // Do not use terminal width
 		// Hack! set Color to avoid separate progress string
 		bar.Set(pb.Color, true)
-		bar.SetWriter(&wrappedWriter{name: pp.name})
+		if logFuncImpl == nil {
+			logFuncImpl = log.Info
+		}
+		bar.SetWriter(&wrappedWriter{name: pp.name, log: logFuncImpl})
 	} else {
 		tmpl := `{{string . "barName" | green}} {{ bar . "<" "-" (cycle . "-" "\\" "|" "/" ) "." ">"}} {{percent .}}`
 		bar.SetTemplateString(tmpl)
@@ -110,6 +116,7 @@ func (pp *ProgressPrinter) goPrintProgress(
 
 type wrappedWriter struct {
 	name string
+	log  logFunc
 }
 
 func (ww *wrappedWriter) Write(p []byte) (int, error) {
@@ -123,7 +130,7 @@ func (ww *wrappedWriter) Write(p []byte) (int, error) {
 	if err := json.Unmarshal(p, &info); err != nil {
 		return 0, err
 	}
-	log.Info("progress",
+	ww.log("progress",
 		zap.String("step", ww.name),
 		zap.String("progress", info.P),
 		zap.String("count", info.C),
@@ -139,8 +146,9 @@ func StartProgress(
 	name string,
 	total int64,
 	redirectLog bool,
+	log logFunc,
 ) *ProgressPrinter {
 	progress := NewProgressPrinter(name, total, redirectLog)
-	progress.goPrintProgress(ctx, nil)
+	progress.goPrintProgress(ctx, log, nil)
 	return progress
 }
diff --git a/pkg/utils/progress_test.go b/pkg/utils/progress_test.go
@@ -28,7 +28,7 @@ func (r *testProgressSuite) TestProgress(c *C) {
 	var p string
 	pCh2 := make(chan string, 2)
 	progress2 := NewProgressPrinter("test", 2, false)
-	progress2.goPrintProgress(ctx, &testWriter{
+	progress2.goPrintProgress(ctx, nil, &testWriter{
 		fn: func(p string) { pCh2 <- p },
 	})
 	progress2.Inc()
@@ -46,7 +46,7 @@ func (r *testProgressSuite) TestProgress(c *C) {
 
 	pCh4 := make(chan string, 4)
 	progress4 := NewProgressPrinter("test", 4, false)
-	progress4.goPrintProgress(ctx, &testWriter{
+	progress4.goPrintProgress(ctx, nil, &testWriter{
 		fn: func(p string) { pCh4 <- p },
 	})
 	progress4.Inc()
@@ -61,7 +61,7 @@ func (r *testProgressSuite) TestProgress(c *C) {
 
 	pCh8 := make(chan string, 8)
 	progress8 := NewProgressPrinter("test", 8, false)
-	progress8.goPrintProgress(ctx, &testWriter{
+	progress8.goPrintProgress(ctx, nil, &testWriter{
 		fn: func(p string) { pCh8 <- p },
 	})
 	progress8.Inc()