refactor(cmd): move pool flag to global set

iyear · Sep 9, 2023 · 37dc867 · 37dc867
1 parent 4190b6f
commit 37dc867
Show file tree

Hide file tree

Showing 4 changed files with 3 additions and 3 deletions.
diff --git a/app/dl/dl.go b/app/dl/dl.go
@@ -34,7 +34,6 @@ type Options struct {
 	Exclude    []string
 	Desc       bool
 	Takeout    bool
-	PoolSize   int64
 
 	// resume opts
 	Continue, Restart bool
@@ -52,7 +51,7 @@ func Run(ctx context.Context, opts *Options) error {
 	}
 
 	return tgc.RunWithAuth(ctx, c, func(ctx context.Context) (rerr error) {
-		pool := dcpool.NewPool(c, opts.PoolSize, floodwait.NewSimpleWaiter())
+		pool := dcpool.NewPool(c, int64(viper.GetInt(consts.FlagPoolSize)), floodwait.NewSimpleWaiter())
 		defer multierr.AppendInvoke(&rerr, multierr.Close(pool))
 
 		parsers := []parser{

diff --git a/cmd/dl.go b/cmd/dl.go
@@ -58,7 +58,6 @@ func NewDownload() *cobra.Command {
 	// do not match extension, because some files' extension is corrected by --rewrite-ext flag
 	cmd.Flags().BoolVar(&opts.SkipSame, "skip-same", false, "skip files with the same name(without extension) and size")
 
-	cmd.Flags().Int64Var(&opts.PoolSize, "pool", 3, "specify the size of the DC pool")
 	cmd.Flags().BoolVar(&opts.Desc, "desc", false, "download files from the newest to the oldest ones (may affect resume download)")
 	cmd.Flags().BoolVar(&opts.Takeout, "takeout", false, "takeout sessions let you export data from your account with lower flood wait limits.")
 

diff --git a/cmd/root.go b/cmd/root.go
@@ -54,6 +54,7 @@ func New() *cobra.Command {
 	cmd.PersistentFlags().IntP(consts.FlagPartSize, "s", 512*1024, "part size for transfer, max is 512*1024")
 	cmd.PersistentFlags().IntP(consts.FlagThreads, "t", 4, "max threads for transfer one item")
 	cmd.PersistentFlags().IntP(consts.FlagLimit, "l", 2, "max number of concurrent tasks")
+	cmd.PersistentFlags().Int(consts.FlagPoolSize, 3, "specify the size of the DC pool")
 
 	cmd.PersistentFlags().String(consts.FlagNTP, "", "ntp server host, if not set, use system time")
 	cmd.PersistentFlags().Duration(consts.FlagReconnectTimeout, 2*time.Minute, "Telegram client reconnection backoff timeout, infinite if set to 0") // #158

diff --git a/pkg/consts/flag.go b/pkg/consts/flag.go
@@ -7,6 +7,7 @@ const (
 	FlagPartSize         = "size"
 	FlagThreads          = "threads"
 	FlagLimit            = "limit"
+	FlagPoolSize         = "pool"
 	FlagNTP              = "ntp"
 	FlagReconnectTimeout = "reconnect-timeout"
 	FlagDlTemplate       = "template"