feat: ed2k protocol support (#1316)

pkg/download/downloader.go

@@ -31,6 +31,8 @@ const (
 	bucketTask = "task"
 	// task download data bucket
 	bucketSave = "save"
+	// protocol-level shared client state bucket
+	bucketProtocolState = "protocol_state"
 	// downloader config bucket
 	bucketConfig = "config"
 	// downloader extension bucket
@@ -141,7 +143,7 @@ func NewDownloader(cfg *DownloaderConfig) *Downloader {
 
 func (d *Downloader) Setup() error {
 	// setup storage
-	if err := d.storage.Setup([]string{bucketTask, bucketSave, bucketConfig, bucketExtension, bucketExtensionStorage}); err != nil {
+	if err := d.storage.Setup([]string{bucketTask, bucketSave, bucketProtocolState, bucketConfig, bucketExtension, bucketExtensionStorage}); err != nil {
 		return err
 	}
 	// load config from storage
@@ -165,6 +167,12 @@ func (d *Downloader) Setup() error {
 		if _, ok := d.cfg.DownloaderStoreConfig.ProtocolConfig[protocol]; !ok {
 			d.cfg.DownloaderStoreConfig.ProtocolConfig[protocol] = fm.DefaultConfig()
 		}
+		if sfm, ok := fm.(fetcher.StatefulFetcherManager); ok {
+			sfm.SetStateStore(&protocolStateStore{
+				storage:  d.storage,
+				protocol: protocol,
+			})
+		}
 	}
 
 	// load tasks from storage
@@ -358,9 +366,16 @@ func (d *Downloader) saveTask(task *Task) error {
 		d.Logger.Error().Stack().Err(err).Msgf("serialize fetcher failed: %s", task.ID)
 		return err
 	}
-	if err := d.storage.Put(bucketSave, task.ID, data); err != nil {
-		d.Logger.Error().Stack().Err(err).Msgf("persist fetcher failed: %s", task.ID)
-		return err
+	if data != nil {
+		if err := d.storage.Put(bucketSave, task.ID, data); err != nil {
+			d.Logger.Error().Stack().Err(err).Msgf("persist fetcher failed: %s", task.ID)
+			return err
+		}
+	} else {
+		if err := d.storage.Delete(bucketSave, task.ID); err != nil {
+			d.Logger.Error().Stack().Err(err).Msgf("clear fetcher state failed: %s", task.ID)
+			return err
+		}
 	}
 	if err := d.storage.Put(bucketTask, task.ID, task); err != nil {
 		d.Logger.Error().Stack().Err(err).Msgf("persist task failed: %s", task.ID)
@@ -820,6 +835,26 @@ func (d *Downloader) Clear() error {
 	return nil
 }
 
+type protocolStateStore struct {
+	storage  Storage
+	protocol string
+}
+
+func (s *protocolStateStore) Load(v any) (bool, error) {
+	return s.storage.Get(bucketProtocolState, s.protocol, v)
+}
+
+func (s *protocolStateStore) Save(v any) error {
+	if v == nil {
+		return s.Delete()
+	}
+	return s.storage.Put(bucketProtocolState, s.protocol, v)
+}
+
+func (s *protocolStateStore) Delete() error {
+	return s.storage.Delete(bucketProtocolState, s.protocol)
+}
+
 func (d *Downloader) Listener(fn Listener) {
 	d.listener = fn
 }
@@ -1300,8 +1335,14 @@ func (d *Downloader) doPause(task *Task) (err error) {
 				return err
 			}
 		}
-		if err := d.storage.Put(bucketTask, task.ID, task.clone()); err != nil {
-			return err
+		if task.fetcherManager != nil && task.fetcher != nil {
+			if err := d.saveTask(task); err != nil {
+				return err
+			}
+		} else {
+			if err := d.storage.Put(bucketTask, task.ID, task.clone()); err != nil {
+				return err
+			}
 		}
 		d.emit(EventKeyPause, task)
 		return nil

pkg/download/model.go

@@ -8,6 +8,7 @@ import (
 	"github.com/GopeedLab/gopeed/internal/controller"
 	"github.com/GopeedLab/gopeed/internal/fetcher"
 	"github.com/GopeedLab/gopeed/internal/protocol/bt"
+	"github.com/GopeedLab/gopeed/internal/protocol/ed2k"
 	"github.com/GopeedLab/gopeed/internal/protocol/http"
 	"github.com/GopeedLab/gopeed/pkg/base"
 	"github.com/GopeedLab/gopeed/pkg/util"
@@ -110,6 +111,14 @@ func (t *Task) updateUploadSpeed(downloaded int64, usedTime float64) int64 {
 }
 
 func calcSpeed(speedArr *[]int64, downloaded int64, usedTime float64) int64 {
+	if usedTime <= 0 {
+		return 0
+	}
+	if downloaded < 0 {
+		*speedArr = (*speedArr)[:0]
+		return 0
+	}
+
 	*speedArr = append(*speedArr, downloaded)
 	// Record last 5 seconds of download speed to calculate the average speed
 	if len(*speedArr) > int(5.0/usedTime) {
@@ -156,6 +165,7 @@ func (cfg *DownloaderConfig) Init() *DownloaderConfig {
 		cfg.FetchManagers = []fetcher.FetcherManager{
 			new(http.FetcherManager),
 			new(bt.FetcherManager),
+			new(ed2k.FetcherManager),
 		}
 	}
 	if cfg.RefreshInterval == 0 {

pkg/download/model_test.go

@@ -0,0 +1,18 @@
+package download
+
+import "testing"
+
+func TestCalcSpeedResetOnRollback(t *testing.T) {
+	speedArr := []int64{1024, 2048, 4096}
+
+	if got := calcSpeed(&speedArr, -512, 1); got != 0 {
+		t.Fatalf("calcSpeed() = %d, want 0 after rollback", got)
+	}
+	if len(speedArr) != 0 {
+		t.Fatalf("speed window len = %d, want 0 after rollback", len(speedArr))
+	}
+
+	if got := calcSpeed(&speedArr, 1024, 1); got != 1024 {
+		t.Fatalf("calcSpeed() = %d, want 1024 after reset", got)
+	}
+}

pkg/protocol/ed2k/model.go

@@ -0,0 +1,14 @@
+package ed2k
+
+type Stats struct {
+	State         string `json:"state"`
+	Paused        bool   `json:"paused"`
+	ActivePeers   int    `json:"activePeers"`
+	TotalPeers    int    `json:"totalPeers"`
+	DownloadRate  int    `json:"downloadRate"`
+	Upload        int64  `json:"upload"`
+	UploadRate    int    `json:"uploadRate"`
+	TotalDone     int64  `json:"totalDone"`
+	TotalReceived int64  `json:"totalReceived"`
+	TotalWanted   int64  `json:"totalWanted"`
+}