cmd/sync: add global traffic control for sync

Author: zhijian-pro

cmd/sync.go

@@ -265,6 +265,10 @@ func syncStorageFlags() []cli.Flag {
 			Name:  "bwlimit",
 			Usage: "limit bandwidth in Mbps (0 means unlimited)",
 		},
+		&cli.StringFlag{
+			Name:  "traffic-control-addr",
+			Usage: "the address of the traffic control",
+		},
 	})
 }
 

pkg/sync/config.go

@@ -30,45 +30,46 @@ import (
 )
 
 type Config struct {
-	StorageClass   string
-	Start          string
-	End            string
-	Threads        int
-	Update         bool
-	ForceUpdate    bool
-	Perms          bool
-	MaxFailure     int64
-	Dry            bool
-	DeleteSrc      bool
-	DeleteDst      bool
-	MatchFullPath  bool
-	Dirs           bool
-	Exclude        []string
-	Include        []string
-	Existing       bool
-	IgnoreExisting bool
-	Links          bool
-	Inplace        bool
-	Limit          int64
-	Manager        string
-	Workers        []string
-	ManagerAddr    string
-	ListThreads    int
-	ListDepth      int
-	BWLimit        int64
-	NoHTTPS        bool
-	Verbose        bool
-	Quiet          bool
-	CheckAll       bool
-	CheckNew       bool
-	CheckChange    bool
-	MaxSize        int64
-	MinSize        int64
-	MaxAge         time.Duration
-	MinAge         time.Duration
-	StartTime      time.Time
-	EndTime        time.Time
-	Env            map[string]string
+	StorageClass       string
+	Start              string
+	End                string
+	Threads            int
+	Update             bool
+	ForceUpdate        bool
+	Perms              bool
+	MaxFailure         int64
+	Dry                bool
+	DeleteSrc          bool
+	DeleteDst          bool
+	MatchFullPath      bool
+	Dirs               bool
+	Exclude            []string
+	Include            []string
+	Existing           bool
+	IgnoreExisting     bool
+	Links              bool
+	Inplace            bool
+	Limit              int64
+	Manager            string
+	Workers            []string
+	ManagerAddr        string
+	ListThreads        int
+	ListDepth          int
+	BWLimit            int64
+	TrafficControlAddr string
+	NoHTTPS            bool
+	Verbose            bool
+	Quiet              bool
+	CheckAll           bool
+	CheckNew           bool
+	CheckChange        bool
+	MaxSize            int64
+	MinSize            int64
+	MaxAge             time.Duration
+	MinAge             time.Duration
+	StartTime          time.Time
+	EndTime            time.Time
+	Env                map[string]string
 
 	FilesFrom string
 
@@ -171,46 +172,47 @@ func NewConfigFromCli(c *cli.Context) *Config {
 		}
 	}
 	cfg := &Config{
-		StorageClass:   c.String("storage-class"),
-		Start:          c.String("start"),
-		End:            c.String("end"),
-		Threads:        c.Int("threads"),
-		ListThreads:    c.Int("list-threads"),
-		ListDepth:      c.Int("list-depth"),
-		Update:         c.Bool("update"),
-		ForceUpdate:    c.Bool("force-update"),
-		Perms:          c.Bool("perms"),
-		Dirs:           c.Bool("dirs"),
-		Dry:            c.Bool("dry"),
-		MaxFailure:     c.Int64("max-failure"),
-		DeleteSrc:      c.Bool("delete-src"),
-		DeleteDst:      c.Bool("delete-dst"),
-		Exclude:        c.StringSlice("exclude"),
-		Include:        c.StringSlice("include"),
-		MatchFullPath:  c.Bool("match-full-path"),
-		Existing:       c.Bool("existing"),
-		IgnoreExisting: c.Bool("ignore-existing"),
-		Links:          c.Bool("links"),
-		Inplace:        c.Bool("inplace"),
-		Limit:          c.Int64("limit"),
-		Workers:        c.StringSlice("worker"),
-		ManagerAddr:    c.String("manager-addr"),
-		Manager:        c.String("manager"),
-		BWLimit:        utils.ParseMbps(c, "bwlimit"),
-		NoHTTPS:        c.Bool("no-https"),
-		Verbose:        c.Bool("verbose"),
-		Quiet:          c.Bool("quiet"),
-		CheckAll:       c.Bool("check-all"),
-		CheckNew:       c.Bool("check-new"),
-		CheckChange:    c.Bool("check-change"),
-		MaxSize:        int64(utils.ParseBytes(c, "max-size", 'B')),
-		MinSize:        int64(utils.ParseBytes(c, "min-size", 'B')),
-		MaxAge:         utils.Duration(c.String("max-age")),
-		MinAge:         utils.Duration(c.String("min-age")),
-		StartTime:      startTime,
-		EndTime:        endTime,
-		FilesFrom:      c.String("files-from"),
-		Env:            make(map[string]string),
+		StorageClass:       c.String("storage-class"),
+		Start:              c.String("start"),
+		End:                c.String("end"),
+		Threads:            c.Int("threads"),
+		ListThreads:        c.Int("list-threads"),
+		ListDepth:          c.Int("list-depth"),
+		Update:             c.Bool("update"),
+		ForceUpdate:        c.Bool("force-update"),
+		Perms:              c.Bool("perms"),
+		Dirs:               c.Bool("dirs"),
+		Dry:                c.Bool("dry"),
+		MaxFailure:         c.Int64("max-failure"),
+		DeleteSrc:          c.Bool("delete-src"),
+		DeleteDst:          c.Bool("delete-dst"),
+		Exclude:            c.StringSlice("exclude"),
+		Include:            c.StringSlice("include"),
+		MatchFullPath:      c.Bool("match-full-path"),
+		Existing:           c.Bool("existing"),
+		IgnoreExisting:     c.Bool("ignore-existing"),
+		Links:              c.Bool("links"),
+		Inplace:            c.Bool("inplace"),
+		Limit:              c.Int64("limit"),
+		Workers:            c.StringSlice("worker"),
+		ManagerAddr:        c.String("manager-addr"),
+		Manager:            c.String("manager"),
+		BWLimit:            utils.ParseMbps(c, "bwlimit"),
+		TrafficControlAddr: c.String("traffic-control-addr"),
+		NoHTTPS:            c.Bool("no-https"),
+		Verbose:            c.Bool("verbose"),
+		Quiet:              c.Bool("quiet"),
+		CheckAll:           c.Bool("check-all"),
+		CheckNew:           c.Bool("check-new"),
+		CheckChange:        c.Bool("check-change"),
+		MaxSize:            int64(utils.ParseBytes(c, "max-size", 'B')),
+		MinSize:            int64(utils.ParseBytes(c, "min-size", 'B')),
+		MaxAge:             utils.Duration(c.String("max-age")),
+		MinAge:             utils.Duration(c.String("min-age")),
+		StartTime:          startTime,
+		EndTime:            endTime,
+		FilesFrom:          c.String("files-from"),
+		Env:                make(map[string]string),
 	}
 	if !c.IsSet("max-size") {
 		cfg.MaxSize = math.MaxInt64

pkg/sync/sync.go

@@ -20,11 +20,13 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"hash/crc32"
 	"io"
 	"math"
+	"net/http"
 	"os"
 	"path"
 	"runtime"
@@ -66,9 +68,125 @@ var (
 	deleted, failed         *utils.Bar
 	listedPrefix            *utils.Bar
 	concurrent              chan int
-	limiter                 *ratelimit.Bucket
+	limiter                 *mixedLimiter
 	totalHandled            atomic.Int64
 )
+
+type mixedLimiter struct {
+	global *globalLimit
+	local  *ratelimit.Bucket
+}
+
+func (l *mixedLimiter) Wait(count int64) {
+	if l.local != nil {
+		l.local.Wait(count)
+	}
+	if l.global != nil {
+		l.global.wait(count)
+	}
+}
+
+type globalLimit struct {
+	sync.Mutex
+	balance int64
+	due     time.Time
+	need    int64
+	waiters []*sync.Cond
+
+	address string
+}
+type req struct {
+	// Positive numbers indicate a request, negative numbers indicate a payback.
+	Bytes int64 `json:"bytes"`
+}
+
+type resp struct {
+	Granted int64 `json:"granted"` // bytes
+	Expired int64 `json:"expired"` // Millisecond
+}
+
+func (l *globalLimit) request(ask int64) (int64, int64, error) {
+	r := req{Bytes: ask}
+	data, err := json.Marshal(r)
+	if err != nil {
+		return 0, 0, err
+	}
+	result, err := http.Post(l.address, "application/json", bytes.NewReader(data))
+	if err != nil || result.StatusCode != http.StatusOK {
+		var status string
+		if result != nil {
+			status = http.StatusText(result.StatusCode)
+		}
+		logger.Errorf("request traffic control %s failed: %s, http status: %s", l.address, err, status)
+		return 0, 0, err
+	}
+	content, err := io.ReadAll(result.Body)
+	if err != nil {
+		return 0, 0, err
+	}
+	res := resp{}
+	if err := json.Unmarshal(content, &res); err != nil {
+		return 0, 0, err
+	}
+	return res.Granted, res.Expired, nil
+}
+
+func (l *globalLimit) wait(bytes int64) {
+	l.Lock()
+	defer l.Unlock()
+	if bytes <= 0 || l.balance >= bytes && len(l.waiters) == 0 {
+		l.balance -= bytes
+		return
+	}
+	l.need += bytes
+
+	var me = sync.NewCond(l)
+	l.waiters = append(l.waiters, me)
+	for l.waiters[0] != me {
+		me.Wait()
+	}
+
+	if l.balance < bytes {
+		// request credit for other waiters together
+		ask := l.need - l.balance
+		if ask >= bytes*10 {
+			// don't wait for too long
+			ask = bytes * 10
+		}
+		l.Unlock()
+		granted, expire, err := l.request(ask)
+		l.Lock()
+		if err == nil {
+			l.balance += granted
+			l.due = time.Now().Add(time.Millisecond * time.Duration(expire))
+			logger.Debugf("grant %d from %s until %s", granted, l.address, l.due)
+		}
+	}
+
+	l.balance -= bytes
+	l.need -= bytes
+	l.waiters = l.waiters[1:]
+	if len(l.waiters) > 0 {
+		l.waiters[0].Signal()
+	}
+}
+
+func (l *globalLimit) checkBalance() {
+	now := time.Now()
+	l.Lock()
+	if l.balance > 0 && l.need == 0 && l.due.Before(now) {
+		payback := l.balance
+		if payback > 1<<30 {
+			payback = 1 << 30
+		}
+		l.balance -= payback
+		l.Unlock()
+		_, _, _ = l.request(-payback)
+	} else {
+		l.Unlock()
+	}
+}
+
 var crcTable = crc32.MakeTable(crc32.Castagnoli)
 var logger = utils.GetLogger("juicefs")
 var ctx = context.Background()
@@ -1711,9 +1829,26 @@ func Sync(src, dst object.ObjectStorage, config *Config) error {
 	tasks := make(chan object.Object, bufferSize)
 	wg := sync.WaitGroup{}
 	concurrent = make(chan int, config.Threads)
+	var localLimit *ratelimit.Bucket
 	if config.BWLimit > 0 {
 		bps := float64(config.BWLimit*1e6/8) * 0.85 // 15% overhead
-		limiter = ratelimit.NewBucketWithRate(bps, int64(bps)/10)
+		localLimit = ratelimit.NewBucketWithRate(bps, int64(bps)/10)
+	}
+	var gLimit *globalLimit
+	if config.TrafficControlAddr != "" {
+		gLimit = &globalLimit{address: config.TrafficControlAddr}
+		go func() {
+			for {
+				time.Sleep(time.Millisecond * 10)
+				gLimit.checkBalance()
+			}
+		}()
+	}
+	if localLimit != nil || gLimit != nil {
+		limiter = &mixedLimiter{
+			global: gLimit,
+			local:  localLimit,
+		}
 	}
 
 	progress := utils.NewProgress(config.Verbose || config.Quiet || config.Manager != "")