add type rate

Signed-off-by: yeya24 <[email protected]>

cmd/dfdaemon/app/init.go

@@ -65,7 +65,7 @@ func initDfdaemon(cfg config.Properties) error {
 	if err != nil {
 		return errors.Wrap(err, "get dfget version")
 	}
-	logrus.Infof("use dfget %s from %s", bytes.TrimSpace(dfgetVersion), cfg.DFPath)
+	logrus.Infof("use %s from %s", bytes.TrimSpace(dfgetVersion), cfg.DFPath)
 
 	return nil
 }

cmd/dfdaemon/app/root.go

@@ -22,12 +22,14 @@ import (
 	"os/exec"
 	"path/filepath"
 	"reflect"
+	"time"
 
 	"github.com/dragonflyoss/Dragonfly/dfdaemon"
 	"github.com/dragonflyoss/Dragonfly/dfdaemon/config"
 	"github.com/dragonflyoss/Dragonfly/dfdaemon/constant"
 	dferr "github.com/dragonflyoss/Dragonfly/pkg/errortypes"
 	"github.com/dragonflyoss/Dragonfly/pkg/netutils"
+	"github.com/dragonflyoss/Dragonfly/pkg/rate"
 
 	"github.com/mitchellh/mapstructure"
 	"github.com/pkg/errors"
@@ -92,7 +94,7 @@ func init() {
 	// dfget download config
 	rf.String("localrepo", filepath.Join(os.Getenv("HOME"), ".small-dragonfly/dfdaemon/data/"), "temp output dir of dfdaemon")
 	rf.String("dfpath", defaultDfgetPath, "dfget path")
-	rf.String("ratelimit", netutils.NetLimit(), "net speed limit,format:xxxM/K")
+	rf.Var(netutils.NetLimit(), "ratelimit", "net speed limit")
 	rf.StringSlice("node", nil, "specify the addresses(host:port) of supernodes that will be passed to dfget.")
 
 	exitOnError(bindRootFlags(viper.GetViper()), "bind root command flags")
@@ -152,6 +154,8 @@ func getConfigFromViper(v *viper.Viper) (*config.Properties, error) {
 			reflect.TypeOf(config.Regexp{}),
 			reflect.TypeOf(config.URL{}),
 			reflect.TypeOf(config.CertPool{}),
+			reflect.TypeOf(time.Second),
+			reflect.TypeOf(rate.B),
 		)
 	}); err != nil {
 		return nil, errors.Wrap(err, "unmarshal yaml")

cmd/dfget/app/root.go

@@ -20,7 +20,6 @@ import (
 	"fmt"
 	"os"
 	"path"
-	"strconv"
 	"strings"
 	"time"
 
@@ -36,12 +35,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
-var (
-	localLimit string
-	totalLimit string
-	minRate    string
-	filter     string
-)
+var filter string
 
 var cfg = config.NewConfig()
 
@@ -75,9 +69,7 @@ func runDfget() error {
 		return err
 	}
 
-	if err := transParams(); err != nil {
-		return err
-	}
+	cfg.Filter = transFilter(filter)
 
 	// get config from property files
 	initProperties()
@@ -145,27 +137,6 @@ func initProperties() {
 	}
 }
 
-// transParams trans the user-friendly parameter formats
-// to the format corresponding to the `Config` struct.
-func transParams() error {
-	cfg.Filter = transFilter(filter)
-
-	var err error
-	if cfg.LocalLimit, err = transLimit(localLimit); err != nil {
-		return errors.Wrapf(errortypes.ErrConvertFailed, "locallimit: %v", err)
-	}
-
-	if cfg.MinRate, err = transLimit(minRate); err != nil {
-		return errors.Wrapf(errortypes.ErrConvertFailed, "minrate: %v", err)
-	}
-
-	if cfg.TotalLimit, err = transLimit(totalLimit); err != nil {
-		return errors.Wrapf(errortypes.ErrConvertFailed, "totallimit: %v", err)
-	}
-
-	return nil
-}
-
 // initClientLog initializes dfget client's logger.
 // There are two kinds of logger dfget client uses: logfile and console.
 // logfile is used to stored generated log in local filesystem,
@@ -198,13 +169,13 @@ func initFlags() {
 		"Destination path which is used to store the requested downloading file. It must contain detailed directory and specific filename, for example, '/tmp/file.mp4'")
 
 	// localLimit & minRate & totalLimit & timeout
-	flagSet.StringVarP(&localLimit, "locallimit", "s", "",
-		"network bandwidth rate limit for single download task, in format of 20M/m/K/k")
-	flagSet.StringVar(&minRate, "minrate", "",
-		"minimal network bandwidth rate for downloading a file, in format of 20M/m/K/k")
-	flagSet.StringVar(&totalLimit, "totallimit", "",
-		"network bandwidth rate limit for the whole host, in format of 20M/m/K/k")
-	flagSet.IntVarP(&cfg.Timeout, "timeout", "e", 0,
+	flagSet.VarP(&cfg.LocalLimit, "locallimit", "s",
+		"network bandwidth rate limit for single download task, in format of G(B)/g/M(B)/m/K(B)/k/B, pure number will also be parsed as Byte")
+	flagSet.Var(&cfg.MinRate, "minrate",
+		"minimal network bandwidth rate for downloading a file, in format of G(B)/g/M(B)/m/K(B)/k/B, pure number will also be parsed as Byte")
+	flagSet.Var(&cfg.TotalLimit, "totallimit",
+		"network bandwidth rate limit for the whole host, in format of G(B)/g/M(B)/m/K(B)/k/B, pure number will also be parsed as Byte")
+	flagSet.DurationVarP(&cfg.Timeout, "timeout", "e", 0,
 		"Timeout set for file downloading task. If dfget has not finished downloading all pieces of file before --timeout, the dfget will throw an error and exit")
 
 	// md5 & identifier
@@ -257,29 +228,6 @@ func initFlags() {
 	flagSet.MarkDeprecated("exceed", "please use '--timeout' or '-e' instead")
 }
 
-// Helper functions.
-func transLimit(limit string) (int, error) {
-	if stringutils.IsEmptyStr(limit) {
-		return 0, nil
-	}
-	l := len(limit)
-	i, err := strconv.Atoi(limit[:l-1])
-
-	if err != nil {
-		return 0, err
-	}
-
-	unit := limit[l-1]
-	if unit == 'k' || unit == 'K' {
-		return i * 1024, nil
-	}
-	if unit == 'm' || unit == 'M' {
-		return i * 1024 * 1024, nil
-	}
-	return 0, fmt.Errorf("invalid unit '%c' of '%s', 'KkMm' are supported",
-		unit, limit)
-}
-
 func transFilter(filter string) []string {
 	if stringutils.IsEmptyStr(filter) {
 		return nil

cmd/dfget/app/root_test.go

@@ -18,18 +18,15 @@ package app
 
 import (
 	"bytes"
-	"fmt"
 	"io/ioutil"
 	"os"
 	"path/filepath"
-	"strconv"
-	"strings"
 	"testing"
 	"time"
 
 	"github.com/dragonflyoss/Dragonfly/dfget/config"
 	"github.com/dragonflyoss/Dragonfly/pkg/errortypes"
-	"github.com/dragonflyoss/Dragonfly/pkg/stringutils"
+	"github.com/dragonflyoss/Dragonfly/pkg/rate"
 
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/suite"
@@ -41,11 +38,10 @@ type dfgetSuit struct {
 
 func (suit *dfgetSuit) Test_initFlagsNoArguments() {
 	suit.Nil(cfg.Node)
-	suit.Equal(cfg.LocalLimit, 0)
-	suit.Equal(cfg.TotalLimit, 0)
+	suit.Equal(cfg.LocalLimit, 20*rate.MB)
+	suit.Equal(cfg.TotalLimit, rate.Rate(0))
 	suit.Equal(cfg.Notbs, false)
 	suit.Equal(cfg.DFDaemon, false)
-	suit.Equal(cfg.ShowBar, false)
 	suit.Equal(cfg.Console, false)
 	suit.Equal(cfg.Verbose, false)
 	suit.Equal(cfg.URL, "")
@@ -59,7 +55,7 @@ func (suit *dfgetSuit) Test_initProperties() {
 	iniFile := filepath.Join(dirName, "dragonfly.ini")
 	yamlFile := filepath.Join(dirName, "dragonfly.yaml")
 	iniContent := []byte("[node]\naddress=1.1.1.1")
-	yamlContent := []byte("nodes:\n  - 1.1.1.2\nlocalLimit: 1024000")
+	yamlContent := []byte("nodes:\n  - 1.1.1.2\nlocalLimit: 1000K")
 	ioutil.WriteFile(iniFile, iniContent, os.ModePerm)
 	ioutil.WriteFile(yamlFile, yamlContent, os.ModePerm)
 
@@ -75,20 +71,18 @@ func (suit *dfgetSuit) Test_initProperties() {
 		{configs: []string{iniFile, yamlFile},
 			expected: newProp(0, 0, 0, "1.1.1.1")},
 		{configs: []string{yamlFile, iniFile},
-			expected: newProp(1024000, 0, 0, "1.1.1.2")},
+			expected: newProp(int(rate.MB*20), 0, 0, "1.1.1.2")},
 		{configs: []string{filepath.Join(dirName, "x"), yamlFile},
-			expected: newProp(1024000, 0, 0, "1.1.1.2")},
+			expected: newProp(int(rate.MB*20), 0, 0, "1.1.1.2")},
 	}
 
 	for _, v := range cases {
 		cfg = config.NewConfig()
 		buf.Reset()
 		cfg.ConfigFiles = v.configs
-		localLimitStr := strconv.FormatInt(int64(v.expected.LocalLimit/1024), 10)
-		totalLimitStr := strconv.FormatInt(int64(v.expected.TotalLimit/1024), 10)
 		rootCmd.Flags().Parse([]string{
-			"--locallimit", fmt.Sprintf("%sk", localLimitStr),
-			"--totallimit", fmt.Sprintf("%sk", totalLimitStr)})
+			"--locallimit", v.expected.LocalLimit.String(),
+			"--totallimit", v.expected.TotalLimit.String()})
 		initProperties()
 		suit.EqualValues(cfg.Node, v.expected.Nodes)
 		suit.Equal(cfg.LocalLimit, v.expected.LocalLimit)
@@ -97,33 +91,6 @@ func (suit *dfgetSuit) Test_initProperties() {
 	}
 }
 
-func (suit *dfgetSuit) Test_transLimit() {
-	var cases = map[string]struct {
-		i   int
-		err string
-	}{
-		"20M":   {20971520, ""},
-		"20m":   {20971520, ""},
-		"10k":   {10240, ""},
-		"10K":   {10240, ""},
-		"10x":   {0, "invalid unit 'x' of '10x', 'KkMm' are supported"},
-		"10.0x": {0, "invalid syntax"},
-		"ab":    {0, "invalid syntax"},
-		"abM":   {0, "invalid syntax"},
-	}
-
-	for k, v := range cases {
-		i, e := transLimit(k)
-		suit.Equal(i, v.i)
-		if stringutils.IsEmptyStr(v.err) {
-			suit.Nil(e)
-		} else {
-			suit.NotNil(e)
-			suit.True(strings.Contains(e.Error(), v.err), true)
-		}
-	}
-}
-
 func (suit *dfgetSuit) Test_transFilter() {
 	var cases = []string{
 		"a&b&c",
@@ -170,10 +137,10 @@ func newProp(local int, total int, size int, nodes ...string) *config.Properties
 		p.Nodes = nodes
 	}
 	if local != 0 {
-		p.LocalLimit = local
+		p.LocalLimit = rate.Rate(local)
 	}
 	if total != 0 {
-		p.TotalLimit = total
+		p.TotalLimit = rate.Rate(total)
 	}
 	if size != 0 {
 		p.ClientQueueSize = size

cmd/supernode/app/root.go

@@ -80,11 +80,11 @@ func setupFlags(cmd *cobra.Command, opt *Options) {
 	flagSet.StringVar(&opt.DownloadPath, "download-path", opt.DownloadPath,
 		"specifies the path where to store downloaded filed from source address")
 
-	flagSet.IntVar(&opt.SystemReservedBandwidth, "system-bandwidth", opt.SystemReservedBandwidth,
-		"Network rate reserved for system (unit: MB/s)")
+	flagSet.Var(&opt.SystemReservedBandwidth, "system-bandwidth",
+		"Network rate reserved for system")
 
-	flagSet.IntVar(&opt.MaxBandwidth, "max-bandwidth", opt.MaxBandwidth,
-		"network rate that supernode can use (unit: MB/s)")
+	flagSet.Var(&opt.MaxBandwidth, "max-bandwidth",
+		"network rate that supernode can use")
 
 	flagSet.IntVar(&opt.SchedulerCorePoolSize, "pool-size", opt.SchedulerCorePoolSize,
 		"the core pool size of ScheduledExecutorService")

dfdaemon/config/config.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/dragonflyoss/Dragonfly/dfdaemon/constant"
 	dferr "github.com/dragonflyoss/Dragonfly/pkg/errortypes"
+	"github.com/dragonflyoss/Dragonfly/pkg/rate"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/afero"
@@ -97,11 +98,11 @@ type Properties struct {
 	MaxProcs int `yaml:"maxprocs" json:"maxprocs"`
 
 	// dfget config
-	DfgetFlags []string `yaml:"dfget_flags" json:"dfget_flags"`
-	SuperNodes []string `yaml:"supernodes" json:"supernodes"`
-	RateLimit  string   `yaml:"ratelimit" json:"ratelimit"`
-	DFRepo     string   `yaml:"localrepo" json:"localrepo"`
-	DFPath     string   `yaml:"dfpath" json:"dfpath"`
+	DfgetFlags []string  `yaml:"dfget_flags" json:"dfget_flags"`
+	SuperNodes []string  `yaml:"supernodes" json:"supernodes"`
+	RateLimit  rate.Rate `yaml:"ratelimit" json:"ratelimit"`
+	DFRepo     string    `yaml:"localrepo" json:"localrepo"`
+	DFPath     string    `yaml:"dfpath" json:"dfpath"`
 }
 
 // Validate validates the config
@@ -127,13 +128,6 @@ func (p *Properties) Validate() error {
 		)
 	}
 
-	if ok, _ := regexp.MatchString("^[[:digit:]]+[MK]$", p.RateLimit); !ok {
-		return dferr.Newf(
-			constant.CodeExitRateLimitInvalid,
-			"invalid rate limit %s", p.RateLimit,
-		)
-	}
-
 	return nil
 }
 
@@ -150,7 +144,7 @@ func (p *Properties) DFGetConfig() DFGetConfig {
 	dfgetConfig := DFGetConfig{
 		DfgetFlags: dfgetFlags,
 		SuperNodes: p.SuperNodes,
-		RateLimit:  p.RateLimit,
+		RateLimit:  p.RateLimit.String(),
 		DFRepo:     p.DFRepo,
 		DFPath:     p.DFPath,
 	}

dfdaemon/config/config_test.go

@@ -28,6 +28,7 @@ import (
 
 	"github.com/dragonflyoss/Dragonfly/dfdaemon/constant"
 	"github.com/dragonflyoss/Dragonfly/pkg/errortypes"
+	"github.com/dragonflyoss/Dragonfly/pkg/rate"
 
 	"github.com/pkg/errors"
 	"github.com/spf13/afero"
@@ -95,21 +96,6 @@ func (ts *configTestSuite) TestValidateDFPath() {
 	r.Equal(constant.CodeExitDfgetNotFound, getCode(c.Validate()))
 }
 
-func (ts *configTestSuite) TestValidateRateLimit() {
-	c := defaultConfig()
-	r := ts.Require()
-
-	for _, l := range []string{"M", "K", "1KB"} {
-		c.RateLimit = l
-		r.Equal(constant.CodeExitRateLimitInvalid, getCode(c.Validate()))
-	}
-
-	for _, l := range []string{"1M", "20K", "20M"} {
-		c.RateLimit = l
-		r.Nil(c.Validate())
-	}
-}
-
 func (ts *configTestSuite) TestURLNew() {
 	r := ts.Require()
 
@@ -366,7 +352,7 @@ func defaultConfig() *Properties {
 		HostIP:    "127.0.0.1",
 		DFRepo:    "/tmp",
 		DFPath:    "/tmp",
-		RateLimit: "20M",
+		RateLimit: 20 * rate.MB,
 	}
 }
 

dfdaemon/proxy/proxy.go

@@ -146,7 +146,7 @@ func NewFromConfig(c config.Properties) (*Proxy, error) {
 	if len(c.SuperNodes) > 0 {
 		logrus.Infof("use supernodes: %s", strings.Join(c.SuperNodes, ","))
 	}
-	logrus.Infof("rate limit set to %s", c.RateLimit)
+	logrus.Infof("rate limit set to %s", c.RateLimit.String())
 
 	if c.HijackHTTPS != nil {
 		opts = append(opts, WithHTTPSHosts(c.HijackHTTPS.Hosts...))