feat: add ratelimit implementation in download process.

cmd/sanqiu.go

@@ -66,7 +66,7 @@ func init() {
 	flags.BoolVarP(&argument.Extract, "extract", "e", argument.Extract, "Extract the archive file for filtering.")
 	flags.StringVarP(&argument.DownloadPath, "download", "d", argument.DownloadPath,
 		"The book directory you want to use, default would be current working directory.")
-	flags.IntVarP(&argument.InitialBookID, "initial", "i", argument.InitialBookID,
+	flags.Int64VarP(&argument.InitialBookID, "initial", "i", argument.InitialBookID,
 		"The book id you want to start download. It should exceed 0.")
 	flags.BoolVarP(&argument.Rename, "rename", "r", argument.Rename, "Rename the book file by book ID.")
 	flags.IntVarP(&argument.Thread, "thread", "t", argument.Thread, "The number of concurrent download thead.")

cmd/talebook/download.go

@@ -60,7 +60,7 @@ func init() {
 	flags.StringSliceVarP(&argument.Formats, "format", "f", argument.Formats, "The file formats you want to download.")
 	flags.StringVarP(&argument.DownloadPath, "download", "d", argument.DownloadPath,
 		"The book directory you want to use, default would be current working directory.")
-	flags.IntVarP(&argument.InitialBookID, "initial", "i", argument.InitialBookID,
+	flags.Int64VarP(&argument.InitialBookID, "initial", "i", argument.InitialBookID,
 		"The book id you want to start download. It should exceed 0.")
 	flags.BoolVarP(&argument.Rename, "rename", "r", argument.Rename, "Rename the book file by book ID.")
 	flags.IntVarP(&argument.Thread, "thread", "t", argument.Thread, "The number of concurrent download thead.")

cmd/telegram.go

@@ -91,7 +91,7 @@ func init() {
 	flags.BoolVarP(&argument.Extract, "extract", "e", argument.Extract, "Extract the archive file for filtering.")
 	flags.StringVarP(&argument.DownloadPath, "download", "d", argument.DownloadPath,
 		"The book directory you want to use, default would be current working directory.")
-	flags.IntVarP(&argument.InitialBookID, "initial", "i", argument.InitialBookID,
+	flags.Int64VarP(&argument.InitialBookID, "initial", "i", argument.InitialBookID,
 		"The book id you want to start download. It should exceed 0.")
 	flags.BoolVarP(&argument.Rename, "rename", "r", argument.Rename, "Rename the book file by book ID.")
 	flags.IntVarP(&argument.Thread, "thread", "t", argument.Thread, "The number of concurrent download thead.")

go.mod

@@ -11,9 +11,11 @@ require (
 	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db
 	github.com/schollz/progressbar/v3 v3.12.1
 	github.com/spf13/cobra v1.6.1
+	go.uber.org/ratelimit v0.2.0
 )
 
 require (
+	github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 // indirect
 	github.com/inconshreveable/mousetrap v1.0.1 // indirect
 	github.com/mattn/go-isatty v0.0.16 // indirect
 	github.com/mattn/go-runewidth v0.0.14 // indirect

go.sum

@@ -1,3 +1,5 @@
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129 h1:MzBOUgng9orim59UnfUTLRjMpd09C5uEVQ6RPGeCaVI=
+github.com/andres-erbsen/clock v0.0.0-20160526145045-9e14626cd129/go.mod h1:rFgpPQZYZ8vdbc+48xibu8ALc3yeyd64IhHS+PU6Yyg=
 github.com/bits-and-blooms/bitset v1.3.3 h1:R1XWiopGiXf66xygsiLpzLo67xEYvMkHw3w+rCOSAwg=
 github.com/bits-and-blooms/bitset v1.3.3/go.mod h1:gIdJ4wp64HaoK2YrL1Q5/N7Y16edYb8uY+O0FJTyyDA=
 github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
@@ -44,11 +46,16 @@ github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+
 github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw=
 github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.4/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
+go.uber.org/atomic v1.7.0 h1:ADUqmZGgLDDfbSL9ZmPxKTybcoEYHgpYfELNoN+7hsw=
+go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
+go.uber.org/ratelimit v0.2.0 h1:UQE2Bgi7p2B85uP5dC2bbRtig0C+OeNRnNEafLjsLPA=
+go.uber.org/ratelimit v0.2.0/go.mod h1:YYBV4e4naJvhpitQrWJu1vCpgB7CboMe0qhltKt6mUg=
 golang.org/x/net v0.0.0-20211029224645-99673261e6eb/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.2.0 h1:sZfSu1wtKLGlWI4ZZayP0ck9Y73K1ynO6gqzTdBVdPU=
 golang.org/x/net v0.2.0/go.mod h1:KqCZLdyyvdV855qA2rE3GC2aiw5xGR5TEjj8smXukLY=

internal/argument/arguments.go

@@ -30,7 +30,7 @@ var (
 	}
 	Extract       = false
 	DownloadPath  = ""
-	InitialBookID = 1
+	InitialBookID = int64(1)
 	Rename        = false
 	Thread        = 1
 	RateLimit     = 30

internal/fetcher/common.go

@@ -39,7 +39,7 @@ type Config struct {
 	Formats       []Format // The formats that the user wants.
 	Extract       bool     // Extract the archives after download.
 	DownloadPath  string   // The path for storing the file.
-	InitialBookID int      // The book id start to download.
+	InitialBookID int64    // The book id start to download.
 	Rename        bool     // Rename the file by using book ID.
 	Thread        int      // The number of download threads.
 	RateLimit     int      // Request per minute.

internal/progress/progress.go

@@ -5,8 +5,10 @@ import (
 	"fmt"
 	"os"
 	"sync"
+	"time"
 
 	"github.com/bits-and-blooms/bitset"
+	"go.uber.org/ratelimit"
 )
 
 const NoBookToDownload = -1
@@ -19,14 +21,15 @@ var (
 
 // Progress is a bit based
 type Progress struct {
-	progress *bitset.BitSet // progress is used for file Progress.
-	assigned *bitset.BitSet // the assign status, memory based.
-	lock     *sync.Mutex    // lock is used for concurrent request.
-	file     string         // The Progress file path for download progress.
+	limit    ratelimit.Limiter // The ratelimit for acquiring a book ID.
+	progress *bitset.BitSet    // progress is used for file Progress.
+	assigned *bitset.BitSet    // the assign status, memory based.
+	lock     *sync.Mutex       // lock is used for concurrent request.
+	file     string            // The Progress file path for download progress.
 }
 
 // NewProgress Create a storge for save the download progress.
-func NewProgress(start, size int64, path string) (*Progress, error) {
+func NewProgress(start, size int64, rate int, path string) (*Progress, error) {
 	if start < 1 {
 		return nil, ErrStartBookID
 	}
@@ -84,7 +87,11 @@ func NewProgress(start, size int64, path string) (*Progress, error) {
 	assigned := bitset.New(progress.Len())
 	progress.Copy(assigned)
 
+	// Create ratelimit
+	limit := ratelimit.New(rate, ratelimit.Per(time.Minute))
+
 	return &Progress{
+		limit:    limit,
 		progress: progress,
 		assigned: assigned,
 		lock:     new(sync.Mutex),
@@ -111,6 +118,8 @@ func (storage *Progress) AcquireBookID() int64 {
 	storage.lock.Lock()
 	defer storage.lock.Unlock()
 
+	// Block until the rate meets the given config.
+	storage.limit.Take()
 	for i := uint(0); i < storage.assigned.Len(); i++ {
 		if !storage.assigned.Test(i) {
 			storage.assigned.Set(i)

internal/progress/progress_test.go

@@ -17,7 +17,7 @@ func TestProgress_AcquireBookID(t *testing.T) {
 	file := tempFile()
 	defer func() { _ = os.Remove(file) }()
 
-	s, err := NewProgress(1, 10000, file)
+	s, err := NewProgress(1, 10000, 10000, file)
 	if err != nil {
 		t.Errorf("Error in creating Progress: %v", err)
 	}
@@ -34,7 +34,7 @@ func TestProgress_SaveBookID(t *testing.T) {
 	file := tempFile()
 	defer func() { _ = os.Remove(file) }()
 
-	s, err := NewProgress(1, 1000, file)
+	s, err := NewProgress(1, 1000, 10000, file)
 	if err != nil {
 		t.Errorf("Error in creating Progress: %v", err)
 	}
@@ -51,7 +51,7 @@ func TestProgress_SaveBookID(t *testing.T) {
 		}
 	}
 
-	s2, err := NewProgress(1, 1000, file)
+	s2, err := NewProgress(1, 1000, 10000, file)
 	if err != nil {
 		t.Errorf("Error in creating Progress: %v", err)
 	}