*: estimate copy to cloud speed and remaining time

Implement logic to track copy progress and calculate average speed.Estimate remaining time based on current throughput to improve observability.

Signed-off-by: huanghaoyuanhhy <haoyuan.huang@zilliz.com>

Author: huanghaoyuanhhy

core/meta/taskmgr/migrate.go

@@ -2,6 +2,7 @@ package taskmgr
 
 import (
 	"sync"
+	"time"
 
 	"github.com/vbauerster/mpb/v8"
 	"github.com/vbauerster/mpb/v8/decor"
@@ -20,15 +21,15 @@ func SetMigrateJobID(jobID string) MigrateTaskOpt {
 	}
 }
 
-func IncMigrateCopiedSize(size int64) MigrateTaskOpt {
+func IncMigrateCopiedSize(size int64, cost time.Duration) MigrateTaskOpt {
 	return func(task *MigrateTask) {
 		task.mu.Lock()
 		defer task.mu.Unlock()
 
 		task.copiedSize += size
 
 		if task.bar != nil {
-			task.bar.IncrInt64(size)
+			task.bar.EwmaIncrInt64(size, cost)
 		}
 	}
 }
@@ -38,8 +39,12 @@ func SetMigrateCopyStart() MigrateTaskOpt {
 		opts := []mpb.BarOption{
 			mpb.PrependDecorators(
 				decor.Name("Uploading", decor.WC{C: decor.DindentRight | decor.DextraSpace}),
-				decor.Counters(decor.SizeB1024(0), "% .1f / % .1f")),
-			mpb.AppendDecorators(decor.Percentage()),
+				decor.Counters(decor.SizeB1024(0), "% .2f / % .2f")),
+			mpb.AppendDecorators(
+				decor.EwmaETA(decor.ET_STYLE_GO, 30),
+				decor.Name(" ] "),
+				decor.EwmaSpeed(decor.SizeB1024(0), "% .2f", 30),
+			),
 			mpb.BarRemoveOnComplete(),
 		}
 
@@ -49,7 +54,7 @@ func SetMigrateCopyStart() MigrateTaskOpt {
 		// The size in backupinfo does not include all partition L0 (partition id == -1)
 		// and meta, so the total cannot be set directly.
 		// Need to use dyn total.
-		task.bar = progressbar.Progress().AddBar(0, opts...)
+		task.bar = progressbar.Progress().New(0, mpb.BarStyle().Rbound("|"), opts...)
 		task.bar.SetTotal(task.totalSize, false)
 	}
 }

core/migrate/task.go

@@ -80,7 +80,7 @@ func (s *stage) IsExpired() bool {
 }
 
 func (s *stage) apply(ctx context.Context) error {
-	s.logger.Info("apply stage")
+	s.logger.Info("apply stage, it will take about 10 seconds")
 	// use taskID as dir
 	start := time.Now()
 	resp, err := s.cloudCli.ApplyStage(ctx, s.clusterID, s.taskID)
@@ -192,9 +192,11 @@ func (t *Task) copyToCloud(ctx context.Context) error {
 		SrcPrefix:  t.backupDir,
 		DestPrefix: t.stage.resp.UploadPath,
 		Sem:        t.copySem,
-		OnSuccess: func(copyAttr storage.CopyAttr) {
-			t.taskMgr.UpdateMigrateTask(t.taskID, taskmgr.IncMigrateCopiedSize(copyAttr.Src.Length))
+
+		TraceFn: func(size int64, cost time.Duration) {
+			t.taskMgr.UpdateMigrateTask(t.taskID, taskmgr.IncMigrateCopiedSize(size, cost))
 		},
+
 		CopyByServer: true,
 	}
 

core/storage/copier.go

@@ -0,0 +1,111 @@
+package storage
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"time"
+
+	"go.uber.org/zap"
+
+	"github.com/zilliztech/milvus-backup/internal/log"
+	"github.com/zilliztech/milvus-backup/internal/util/retry"
+)
+
+type CopyAttr struct {
+	Src     ObjectAttr
+	DestKey string
+}
+
+type copierOpt struct {
+	traceFn TraceFn
+}
+
+type TraceFn func(size int64, cost time.Duration)
+
+type trackReader struct {
+	inner io.Reader
+
+	lastRead time.Time
+	traceFn  TraceFn
+}
+
+func newTrackReader(inner io.Reader, traceFn TraceFn) *trackReader {
+	return &trackReader{inner: inner, traceFn: traceFn, lastRead: time.Now()}
+}
+
+func (t *trackReader) Read(p []byte) (int, error) {
+	n, err := t.inner.Read(p)
+	if n > 0 {
+		t.traceFn(int64(n), time.Since(t.lastRead))
+		t.lastRead = time.Now()
+	}
+	return n, err
+}
+
+type copier interface {
+	copy(ctx context.Context, copyAttr CopyAttr) error
+}
+
+// remoteCopier copy data from src to dest by calling dest.CopyObject
+type remoteCopier struct {
+	src  Client
+	dest Client
+
+	logger *zap.Logger
+}
+
+func (rp *remoteCopier) copy(ctx context.Context, copyAttr CopyAttr) error {
+	rp.logger.Debug("copy object", zap.String("src", copyAttr.Src.Key), zap.String("dest", copyAttr.DestKey))
+	i := CopyObjectInput{SrcCli: rp.src, SrcKey: copyAttr.Src.Key, DestKey: copyAttr.DestKey}
+	if err := rp.dest.CopyObject(ctx, i); err != nil {
+		return fmt.Errorf("storage: remote copier copy object %w", err)
+	}
+
+	return nil
+}
+
+// serverCopier copy data from src to dest by backup server
+type serverCopier struct {
+	src  Client
+	dest Client
+
+	opt copierOpt
+
+	logger *zap.Logger
+}
+
+func (sc *serverCopier) copy(ctx context.Context, copyAttr CopyAttr) error {
+	sc.logger.Debug("copy object", zap.String("src_key", copyAttr.Src.Key), zap.String("dest_key", copyAttr.DestKey))
+
+	return retry.Do(ctx, func() error {
+		obj, err := sc.src.GetObject(ctx, copyAttr.Src.Key)
+		if err != nil {
+			return fmt.Errorf("storage: server copier get object %w", err)
+		}
+		defer obj.Body.Close()
+
+		body := io.Reader(obj.Body)
+		if sc.opt.traceFn != nil {
+			body = newTrackReader(obj.Body, sc.opt.traceFn)
+		}
+
+		i := UploadObjectInput{Body: body, Key: copyAttr.DestKey, Size: copyAttr.Src.Length}
+		if err := sc.dest.UploadObject(ctx, i); err != nil {
+			return fmt.Errorf("storage: copier upload object %w", err)
+		}
+
+		return nil
+	})
+}
+
+func newCopier(src, dest Client, copyByServer bool, opt copierOpt) copier {
+	logger := log.L().With(
+		zap.String("src", src.Config().Bucket),
+		zap.String("dest", dest.Config().Bucket),
+	)
+	if copyByServer {
+		return &serverCopier{src: src, dest: dest, opt: opt, logger: logger.With(zap.String("copier", "server"))}
+	}
+	return &remoteCopier{src: src, dest: dest, logger: logger.With(zap.String("copier", "remote"))}
+}

core/storage/copy_task_test.go core/storage/copier_test.gocore/storage/copy_task_test.go renamed to core/storage/copier_test.go

@@ -5,13 +5,39 @@ import (
 	"context"
 	"io"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/zap"
 )
 
-func TestRemoteCopier_copy(t *testing.T) {
+func TestTrackReader_Read(t *testing.T) {
+	t.Run("ReadGTZero", func(t *testing.T) {
+		tr := newTrackReader(bytes.NewReader([]byte("hello")), func(size int64, cost time.Duration) {
+			assert.Equal(t, 5, size)
+			assert.True(t, cost > 0)
+		})
+
+		buf := make([]byte, 10)
+		n, err := tr.Read(buf)
+		assert.NoError(t, err)
+		assert.Equal(t, 5, n)
+	})
+
+	t.Run("ReadZero", func(t *testing.T) {
+		tr := newTrackReader(bytes.NewReader([]byte("hello")), func(size int64, cost time.Duration) {
+			assert.Fail(t, "should not be called")
+		})
+
+		buf := make([]byte, 0)
+		n, err := tr.Read(buf)
+		assert.NoError(t, err)
+		assert.Equal(t, 0, n)
+	})
+}
+
+func TestRemoteCopier_Copy(t *testing.T) {
 	dest := NewMockClient(t)
 	src := NewMockClient(t)
 	rp := &remoteCopier{src: src, dest: dest, logger: zap.NewNop()}

core/storage/copy_task.go

@@ -10,76 +10,8 @@ import (
 	"golang.org/x/sync/semaphore"
 
 	"github.com/zilliztech/milvus-backup/internal/log"
-	"github.com/zilliztech/milvus-backup/internal/util/retry"
 )
 
-type OnSuccessFn func(copyAttr CopyAttr)
-
-type CopyAttr struct {
-	Src     ObjectAttr
-	DestKey string
-}
-
-type copier interface {
-	copy(ctx context.Context, copyAttr CopyAttr) error
-}
-
-// remoteCopier copy data from src to dest by calling dest.CopyObject
-type remoteCopier struct {
-	src  Client
-	dest Client
-
-	logger *zap.Logger
-}
-
-func (rp *remoteCopier) copy(ctx context.Context, copyAttr CopyAttr) error {
-	rp.logger.Debug("copy object", zap.String("src", copyAttr.Src.Key), zap.String("dest", copyAttr.DestKey))
-	i := CopyObjectInput{SrcCli: rp.src, SrcKey: copyAttr.Src.Key, DestKey: copyAttr.DestKey}
-	if err := rp.dest.CopyObject(ctx, i); err != nil {
-		return fmt.Errorf("storage: remote copier copy object %w", err)
-	}
-
-	return nil
-}
-
-// serverCopier copy data from src to dest by backup server
-type serverCopier struct {
-	src  Client
-	dest Client
-
-	logger *zap.Logger
-}
-
-func (sc *serverCopier) copy(ctx context.Context, copyAttr CopyAttr) error {
-	sc.logger.Debug("copy object", zap.String("src_key", copyAttr.Src.Key), zap.String("dest_key", copyAttr.DestKey))
-
-	return retry.Do(ctx, func() error {
-		obj, err := sc.src.GetObject(ctx, copyAttr.Src.Key)
-		if err != nil {
-			return fmt.Errorf("storage: server copier get object %w", err)
-		}
-		defer obj.Body.Close()
-
-		i := UploadObjectInput{Body: obj.Body, Key: copyAttr.DestKey, Size: copyAttr.Src.Length}
-		if err := sc.dest.UploadObject(ctx, i); err != nil {
-			return fmt.Errorf("storage: copier upload object %w", err)
-		}
-
-		return nil
-	})
-}
-
-func newCopier(src, dest Client, copyByServer bool) copier {
-	logger := log.L().With(
-		zap.String("src", src.Config().Bucket),
-		zap.String("dest", dest.Config().Bucket),
-	)
-	if copyByServer {
-		return &serverCopier{src: src, dest: dest, logger: logger.With(zap.String("copier", "server"))}
-	}
-	return &remoteCopier{src: src, dest: dest, logger: logger.With(zap.String("copier", "remote"))}
-}
-
 type CopyPrefixOpt struct {
 	Src  Client
 	Dest Client
@@ -89,7 +21,7 @@ type CopyPrefixOpt struct {
 
 	Sem *semaphore.Weighted
 
-	OnSuccess OnSuccessFn
+	TraceFn TraceFn
 
 	CopyByServer bool
 }
@@ -106,7 +38,7 @@ func NewCopyPrefixTask(opt CopyPrefixOpt) *CopyPrefixTask {
 	return &CopyPrefixTask{
 		opt: opt,
 
-		copier: newCopier(opt.Src, opt.Dest, opt.CopyByServer),
+		copier: newCopier(opt.Src, opt.Dest, opt.CopyByServer, copierOpt{traceFn: opt.TraceFn}),
 
 		logger: log.L().With(zap.String("src", opt.SrcPrefix), zap.String("dest", opt.DestPrefix)),
 	}
@@ -120,10 +52,6 @@ func (c *CopyPrefixTask) copy(ctx context.Context, src ObjectAttr) error {
 		return fmt.Errorf("storage: copy prefix %w", err)
 	}
 
-	if c.opt.OnSuccess != nil {
-		c.opt.OnSuccess(attr)
-	}
-
 	return nil
 }
 
@@ -186,7 +114,7 @@ func NewCopyObjectsTask(opt CopyObjectsOpt) *CopyObjectsTask {
 	return &CopyObjectsTask{
 		opt: opt,
 
-		copier: newCopier(opt.Src, opt.Dest, opt.CopyByServer),
+		copier: newCopier(opt.Src, opt.Dest, opt.CopyByServer, copierOpt{}),
 	}
 }