ebs-warmup: added test cases && better exit logic

Makefile

@@ -208,7 +208,7 @@ fault-trigger:
 # `-race` for race detector.
 # GO_COVER: Whether to run tests with code coverage. Set to 'y' to enable coverage collection.
 #
-test: TEST_PACKAGES = ./cmd/backup-manager/app ./pkg
+test: TEST_PACKAGES = ./cmd/backup-manager/app ./pkg ./cmd/ebs-warmup/internal/tests
 test: ## Run unit tests
  @echo "Run unit tests"
 ifeq ($(GO_COVER),y)

cmd/ebs-warmup/filereader/config.go

@@ -13,6 +13,11 @@
 
 package filereader
 
+import (
+ "github.com/pingcap/tidb-operator/cmd/ebs-warmup/worker"
+ "github.com/pingcap/tidb-operator/cmd/ebs-warmup/worker/tasks"
+)
+
 type Config struct {
  Files string
  Type string
@@ -21,4 +26,7 @@ type Config struct {
  Direct bool
  CheckpointEvery uint64
  CheckpointFile string
+
+ OnStep worker.OnStepHook
+ OnFireRequest func(*tasks.ReadFile)
 }

cmd/ebs-warmup/filereader/exec.go

@@ -38,6 +38,8 @@ const (
  defaultSegmentCount = 16
 )
 
+type sendFileHook func(tasks.ReadFile) (brk bool)
+
 type StatedFile struct {
  Info os.FileInfo
  Path string
@@ -68,36 +70,31 @@ func StatFilesByGlob(glob string) ([]StatedFile, error) {
  return stats, nil
 }
 
-func WarmUpFooters(glob string, sendToWorker func(tasks.ReadFile)) error {
+func WarmUpFooters(glob string, sendToWorker sendFileHook) error {
  files, err := StatFilesByGlob(glob)
  if err != nil {
  return errors.Annotatef(err, "failed to stat files with glob %s", glob)
  }
  for _, file := range files {
- sendToWorker(tasks.ReadFile{
+ shouldBrk := sendToWorker(tasks.ReadFile{
  Type: tasks.ReadLastNBytes(16 * 1024),
  File: file.Info,
  FilePath: file.Path,
  })
+ if shouldBrk {
+ return nil
+ }
  }
  return nil
 }
 
-func WarmUpWholeFile(glob string, sendToWorker func(tasks.ReadFile)) error {
- return warmUpWholeFileBy(glob, func(sf StatedFile) {
- sendFileWithSegmenting(sf, defaultSegmentCount, sendToWorker)
- })
-}
-
-func WarmUpWholeFileAfter(glob string, after time.Time, sendToWorker func(tasks.ReadFile)) error {
- return warmUpWholeFileBy(glob, func(sf StatedFile) {
- if sf.Info.ModTime().After(after) {
- sendFileWithSegmenting(sf, defaultSegmentCount, sendToWorker)
- }
+func WarmUpWholeFile(glob string, sendToWorker sendFileHook) error {
+ return warmUpWholeFileBy(glob, func(sf StatedFile) bool {
+ return sendFileWithSegmenting(sf, defaultSegmentCount, sendToWorker)
  })
 }
 
-func warmUpWholeFileBy(glob string, onFile func(StatedFile)) error {
+func warmUpWholeFileBy(glob string, onFile func(StatedFile) (brk bool)) error {
  files, err := StatFilesByGlob(glob)
  if err != nil {
  return errors.Annotatef(err, "failed to stat files with glob %s", glob)
@@ -107,12 +104,14 @@ func warmUpWholeFileBy(glob string, onFile func(StatedFile)) error {
  return files[i].Info.ModTime().After(files[j].Info.ModTime())
  })
  for _, file := range files {
- onFile(file)
+ if onFile(file) {
+ return context.Canceled
+ }
  }
  return nil
 }
 
-func sendFileWithSegmenting(file StatedFile, partitions int, sendToWorker func(tasks.ReadFile)) {
+func sendFileWithSegmenting(file StatedFile, partitions int, sendToWorker sendFileHook) (brk bool) {
  partitionSize := file.Info.Size() / int64(partitions)
  if partitionSize < minimalSegmentSize {
  partitionSize = minimalSegmentSize
@@ -123,7 +122,7 @@ func sendFileWithSegmenting(file StatedFile, partitions int, sendToWorker func(t
  if offset+partitionSize > file.Info.Size() {
  length = file.Info.Size() - offset
  }
- sendToWorker(
+ shouldBrk := sendToWorker(
  tasks.ReadFile{
  Type: tasks.ReadOffsetAndLength{
  Offset: offset,
@@ -133,13 +132,18 @@ func sendFileWithSegmenting(file StatedFile, partitions int, sendToWorker func(t
  FilePath: file.Path,
  },
  )
+ if shouldBrk {
+ return true
+ }
  offset += partitionSize
  }
+ return false
 }
 
 type WorkersOpt struct {
  ObserveTotalSize *uint64
  RateLimitInMiB float64
+ OnStep worker.OnStepHook
 }
 
 func CreateWorkers(ctx context.Context, n int, opt WorkersOpt) ([]chan<- tasks.ReadFile, *errgroup.Group) {
@@ -160,6 +164,9 @@ func CreateWorkers(ctx context.Context, n int, opt WorkersOpt) ([]chan<- tasks.R
  wr := worker.New(ch)
  wr.RateLimiter = limiter
  wr.OnStep = func(file os.FileInfo, readBytes int, take time.Duration) {
+ if opt.OnStep != nil {
+ opt.OnStep(file, readBytes, take)
+ }
  if opt.ObserveTotalSize != nil {
  new := atomic.AddUint64(opt.ObserveTotalSize, uint64(readBytes))
  if atomic.AddUint64(&loopCounter, 1)%2048 == 0 {
@@ -200,16 +207,21 @@ func RoundRobin[T any](ts []T) func() T {
  return choose
 }
 
-func TrySync(workers []chan<- tasks.ReadFile) {
+func TrySync(ctx context.Context, workers []chan<- tasks.ReadFile) error {
  chs := make([]chan struct{}, 0)
  for _, w := range workers {
  ch := make(chan struct{})
  w <- tasks.ReadFile{Type: tasks.Sync{C: ch}}
  chs = append(chs, ch)
  }
  for _, ch := range chs {
- <-ch
+ select {
+ case <-ch:
+ case <-ctx.Done():
+ return ctx.Err()
+ }
  }
+ return nil
 }
 
 type ExecContext struct {
@@ -222,6 +234,8 @@ type ExecContext struct {
  chooser func() chan<- tasks.ReadFile
  lastSent uint64
  start time.Time
+
+ cancel context.CancelFunc
 }
 
 func (execCtx *ExecContext) perhapsCheckpoint() (uint64, error) {
@@ -241,7 +255,8 @@ func (execCtx *ExecContext) checkpoint() uint64 {
  ckp, err := execCtx.perhapsCheckpoint()
  if err != nil {
  klog.InfoS("Failed to read checkpoint file. Will use time.Now() as checkpoint.", "err", err)
- return uint64(time.Now().UnixMilli())
+ // add a 1s offset for newly created files. (usually in test)
+ return uint64(time.Now().UnixMilli() + 1000)
  }
  return ckp
 }
@@ -250,34 +265,38 @@ func (execCtx *ExecContext) saveCheckpoint(ckp uint64) error {
  return os.WriteFile(execCtx.config.CheckpointFile, []byte(fmt.Sprintf("%d", ckp)), 0o644)
 }
 
-func New(masterCtx context.Context, config Config) *ExecContext {
+func New(config Config) *ExecContext {
  execCtx := &ExecContext{
  config: config,
  }
- execCtx.wkrs, execCtx.eg = CreateWorkers(masterCtx, execCtx.config.NWorkers, WorkersOpt{
+ wCtx, cancel := context.WithCancel(context.Background())
+ execCtx.wkrs, execCtx.eg = CreateWorkers(wCtx, execCtx.config.NWorkers, WorkersOpt{
  ObserveTotalSize: &execCtx.total,
  RateLimitInMiB: execCtx.config.RateLimit,
+ OnStep: config.OnStep,
  })
  execCtx.start = time.Now()
  execCtx.chooser = RoundRobin(execCtx.wkrs)
  execCtx.cnt = uint64(0)
  execCtx.lastSent = execCtx.checkpoint()
+ execCtx.cancel = cancel
  return execCtx
 }
 
-func (execCtx *ExecContext) Run() {
- total := uint64(0)
+func (execCtx *ExecContext) RunAndClose(ctx context.Context) {
+ defer execCtx.cancel()
 
+ total := uint64(0)
  klog.InfoS("Using checkpoint.", "checkpoint", execCtx.lastSent, "time", time.UnixMilli(int64(execCtx.lastSent)).String())
 
  switch execCtx.config.Type {
  case "footer":
- WarmUpFooters(execCtx.config.Files, func(rf tasks.ReadFile) {
- execCtx.sendToWorker(rf)
+ WarmUpFooters(execCtx.config.Files, func(rf tasks.ReadFile) bool {
+ return execCtx.sendToWorker(ctx, rf)
  })
  case "whole":
- WarmUpWholeFile(execCtx.config.Files, func(rf tasks.ReadFile) {
- execCtx.sendToWorker(rf)
+ WarmUpWholeFile(execCtx.config.Files, func(rf tasks.ReadFile) bool {
+ return execCtx.sendToWorker(ctx, rf)
  })
  }
 
@@ -291,26 +310,41 @@ func (execCtx *ExecContext) Run() {
  klog.InfoS("Done.", "take", take, "total", total, "rate", fmt.Sprintf("%s/s", units.HumanSize(rate)))
 }
 
-func (execCtx *ExecContext) sendToWorker(rf tasks.ReadFile) {
- createTs := rf.File.ModTime().UnixMilli()
- if createTs > int64(execCtx.checkpoint()) {
- return
- }
+func (execCtx *ExecContext) checkpointTick(ctx context.Context) {
  execCtx.cnt += 1
  if execCtx.cnt%execCtx.config.CheckpointEvery == 0 {
  ckp := execCtx.lastSent
  now := time.Now()
- TrySync(execCtx.wkrs)
+ if err := TrySync(ctx, execCtx.wkrs); err != nil {
+ klog.ErrorS(err, "skip saving checkpoint: the worker is closing", "err", err)
+ return
+ }
  err := execCtx.saveCheckpoint(ckp)
  if err != nil {
  klog.ErrorS(err, "Failed to save checkpoint.", "checkpoint", ckp, "take", time.Since(now))
  }
  }
+}
+
+func (execCtx *ExecContext) sendToWorker(ctx context.Context, rf tasks.ReadFile) bool {
+ createTs := rf.File.ModTime().UnixMilli()
+ if execCtx.config.OnFireRequest != nil {
+ execCtx.config.OnFireRequest(&rf)
+ }
+ if ctx.Err() != nil {
+ klog.InfoS("early exit due to context canceled.", "err", ctx.Err())
+ return true
+ }
+ if createTs > int64(execCtx.lastSent) {
+ return false
+ }
+ execCtx.checkpointTick(ctx)
  rf.Direct = execCtx.config.Direct
  execCtx.chooser() <- rf
  if execCtx.lastSent < uint64(createTs) {
  klog.Warningln("unordered files: checkpoint is unavailable.", "checkpoint=", execCtx.checkpoint(),
  "createTs=", uint64(createTs), "lastSent=", execCtx.lastSent)
  }
  execCtx.lastSent = uint64(createTs)
+ return false
 }