From 9870a5682e75782b4ab3bbd39c9028e67d34d9d7 Mon Sep 17 00:00:00 2001
From: ljx373327 <ljx373327@alibaba-inc.com>
Date: Thu, 16 Jun 2022 14:19:43 +0800
Subject: [PATCH] fix ratelimit

Signed-off-by: ljx373327 <ljx373327@alibaba-inc.com>
---
 client/config/dfget.go        |  4 ++--
 client/config/dfget_darwin.go | 18 ++++++++++++------
 client/config/dfget_linux.go  | 11 ++++++++---
 client/dfget/dfget.go         |  2 +-
 cmd/dfget/cmd/root.go         |  2 +-
 5 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/client/config/dfget.go b/client/config/dfget.go
index 88d5b88e85c..bb2d819ba79 100644
--- a/client/config/dfget.go
+++ b/client/config/dfget.go
@@ -29,8 +29,8 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
-	"golang.org/x/time/rate"
 
+	"d7y.io/dragonfly/v2/client/clientutil"
 	"d7y.io/dragonfly/v2/cmd/dependency/base"
 	"d7y.io/dragonfly/v2/internal/dferrors"
 	logger "d7y.io/dragonfly/v2/internal/dflog"
@@ -103,7 +103,7 @@ type ClientOption struct {
 	// WorkHome is working directory of dfget.
 	WorkHome string `yaml:"workHome,omitempty" mapstructure:"workHome,omitempty"`
 
-	RateLimit rate.Limit `yaml:"rateLimit,omitempty" mapstructure:"rateLimit,omitempty"`
+	RateLimit clientutil.RateLimit `yaml:"rateLimit,omitempty" mapstructure:"rateLimit,omitempty"`
 
 	// Config file paths,
 	// default:["/etc/dragonfly/dfget.yaml","/etc/dragonfly.conf"].
diff --git a/client/config/dfget_darwin.go b/client/config/dfget_darwin.go
index c6805169da0..176055e513a 100644
--- a/client/config/dfget_darwin.go
+++ b/client/config/dfget_darwin.go
@@ -19,14 +19,20 @@
 
 package config
 
-import "d7y.io/dragonfly/v2/pkg/unit"
+import (
+	"d7y.io/dragonfly/v2/client/clientutil"
+	"d7y.io/dragonfly/v2/pkg/unit"
+	"golang.org/x/time/rate"
+)
 
 var dfgetConfig = ClientOption{
-	URL:               "",
-	Output:            "",
-	Timeout:           0,
-	BenchmarkRate:     128 * unit.KB,
-	RateLimit:         0,
+	URL:           "",
+	Output:        "",
+	Timeout:       0,
+	BenchmarkRate: 128 * unit.KB,
+	RateLimit: clientutil.RateLimit{
+		Limit: rate.Limit(DefaultTotalDownloadLimit),
+	},
 	Md5:               "",
 	DigestMethod:      "",
 	DigestValue:       "",
diff --git a/client/config/dfget_linux.go b/client/config/dfget_linux.go
index 437aeb5591c..61ece8f5f24 100644
--- a/client/config/dfget_linux.go
+++ b/client/config/dfget_linux.go
@@ -19,10 +19,15 @@
 
 package config
 
+import "d7y.io/dragonfly/v2/client/clientutil"
+
 var dfgetConfig = ClientOption{
-	URL:               "",
-	Output:            "",
-	Timeout:           0,
+	URL:     "",
+	Output:  "",
+	Timeout: 0,
+	RateLimit: clientutil.RateLimit{
+		Limit: rate.Limit(DefaultTotalDownloadLimit),
+	},
 	Md5:               "",
 	DigestMethod:      "",
 	DigestValue:       "",
diff --git a/client/dfget/dfget.go b/client/dfget/dfget.go
index 3755f574703..fa4f00a54a2 100644
--- a/client/dfget/dfget.go
+++ b/client/dfget/dfget.go
@@ -220,7 +220,7 @@ func newDownRequest(cfg *config.DfgetConfig, hdr map[string]string) *dfdaemon.Do
 		Url:               cfg.URL,
 		Output:            cfg.Output,
 		Timeout:           uint64(cfg.Timeout),
-		Limit:             float64(cfg.RateLimit),
+		Limit:             float64(cfg.RateLimit.Limit),
 		DisableBackSource: cfg.DisableBackSource,
 		UrlMeta: &base.UrlMeta{
 			Digest: cfg.Digest,
diff --git a/cmd/dfget/cmd/root.go b/cmd/dfget/cmd/root.go
index 16c9ff850d6..1e016bb8061 100644
--- a/cmd/dfget/cmd/root.go
+++ b/cmd/dfget/cmd/root.go
@@ -139,7 +139,7 @@ func init() {
 
 	flagSet.Duration("timeout", dfgetConfig.Timeout, "Timeout for the downloading task, 0 is infinite")
 
-	flagSet.String("limit", unit.Bytes(dfgetConfig.RateLimit).String(),
+	flagSet.String("ratelimit", unit.Bytes(dfgetConfig.RateLimit.Limit).String(),
 		"The downloading network bandwidth limit per second in format of G(B)/g/M(B)/m/K(B)/k/B, pure number will be parsed as Byte, 0 is infinite")
 
 	flagSet.String("digest", dfgetConfig.Digest,