feat(persistentcache): add concurrent piece count option for peers

Signed-off-by: Gaius <gaius.qi@gmail.com>

Author: gaius-qi

go.mod

@@ -4,7 +4,7 @@ go 1.23.8
 
 require (
 	cloud.google.com/go/storage v1.50.0
-	d7y.io/api/v2 v2.1.60
+	d7y.io/api/v2 v2.1.65
 	github.com/MysteriousPotato/go-lockable v1.0.0
 	github.com/Showmax/go-fqdn v1.0.0
 	github.com/VividCortex/mysqlerr v1.0.0
@@ -96,7 +96,7 @@ require (
 	golang.org/x/time v0.12.0
 	google.golang.org/api v0.248.0
 	google.golang.org/grpc v1.75.0
-	google.golang.org/protobuf v1.36.7
+	google.golang.org/protobuf v1.36.8
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0
 	gopkg.in/yaml.v3 v3.0.1
 	gorm.io/driver/mysql v1.4.7

go.sum

@@ -63,8 +63,8 @@ cloud.google.com/go/storage v1.50.0 h1:3TbVkzTooBvnZsk7WaAQfOsNrdoM8QHusXA1cpk6Q
 cloud.google.com/go/storage v1.50.0/go.mod h1:l7XeiD//vx5lfqE3RavfmU9yvk5Pp0Zhcv482poyafY=
 cloud.google.com/go/trace v1.11.6 h1:2O2zjPzqPYAHrn3OKl029qlqG6W8ZdYaOWRyr8NgMT4=
 cloud.google.com/go/trace v1.11.6/go.mod h1:GA855OeDEBiBMzcckLPE2kDunIpC72N+Pq8WFieFjnI=
-d7y.io/api/v2 v2.1.60 h1:Imk1mw30wmTJNdqKFF5VO/eqI/t1H5PfIPPK4ryKX2U=
-d7y.io/api/v2 v2.1.60/go.mod h1:WzakPywEgs27gr/TwrlRarPRRK2o5nN0YrTzHDdFii8=
+d7y.io/api/v2 v2.1.65 h1:1o13i+aF3E566tu6zrDTl3kV3/l09VkhSgythVFc2wM=
+d7y.io/api/v2 v2.1.65/go.mod h1:SQgYPAV1yW4Jo0cXFE7lJyGFy4CBn5cLpQTBLed6+6o=
 dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7gZCb22OMCxBHrMx7a5I7Hp++hsVxbQ4BYO7hU=
 filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
 filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
@@ -2186,8 +2186,8 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
-google.golang.org/protobuf v1.36.7 h1:IgrO7UwFQGJdRNXH/sQux4R1Dj1WAKcLElzeeRaXV2A=
-google.golang.org/protobuf v1.36.7/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY=
+google.golang.org/protobuf v1.36.8 h1:xHScyCOEuuwZEc6UtSOvPbAT4zRh0xcNRYekJwfqyMc=
+google.golang.org/protobuf v1.36.8/go.mod h1:fuxRtAxBytpl4zzqUh6/eyUujkJdNiuEkXntxiD/uRU=
 gopkg.in/airbrake/gobrake.v2 v2.0.9/go.mod h1:/h5ZAUhDkGaJfjzjKLSjv6zCL6O0LLBxU4K+aSYdM/U=
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

scheduler/resource/persistentcache/peer.go

@@ -26,6 +26,12 @@ import (
 	logger "d7y.io/dragonfly/v2/internal/dflog"
 )
 
+const (
+	// defaultConcurrentPieceCount is the fallback value for concurrent pieces per peer
+	// when ConcurrentPieceCount is not reported by the client.
+	defaultConcurrentPieceCount uint32 = 8
+)
+
 const (
 	// Peer has been created but did not start running.
 	PeerStatePending = "Pending"
@@ -69,6 +75,16 @@ const (
 	PeerEventFailed = "Failed"
 )
 
+// PeerOption is a functional option for persistent cache peer.
+type PeerOption func(peer *Peer)
+
+// WithConcurrentPieceCount set ConcurrentPieceCount for peer.
+func WithConcurrentPieceCount(count uint32) PeerOption {
+	return func(p *Peer) {
+		p.ConcurrentPieceCount = count
+	}
+}
+
 // Peer contains content for persistent cache peer.
 type Peer struct {
 	// ID is persistent cache peer id.
@@ -77,6 +93,9 @@ type Peer struct {
 	// Persistent is whether the peer is persistent.
 	Persistent bool
 
+	// ConcurrentPieceCount is the count of pieces that can be downloaded concurrently.
+	ConcurrentPieceCount uint32
+
 	// Pieces is finished pieces bitset.
 	FinishedPieces *bitset.BitSet
 
@@ -107,18 +126,19 @@ type Peer struct {
 
 // New persistent cache peer instance.
 func NewPeer(id, state string, persistent bool, finishedPieces *bitset.BitSet, blockParents []string, task *Task, host *Host,
-	cost time.Duration, createdAt, updatedAt time.Time, log *logger.SugaredLoggerOnWith) *Peer {
+	cost time.Duration, createdAt, updatedAt time.Time, log *logger.SugaredLoggerOnWith, options ...PeerOption) *Peer {
 	p := &Peer{
-		ID:             id,
-		Persistent:     persistent,
-		FinishedPieces: finishedPieces,
-		Task:           task,
-		Host:           host,
-		BlockParents:   blockParents,
-		Cost:           cost,
-		CreatedAt:      createdAt,
-		UpdatedAt:      updatedAt,
-		Log:            logger.WithPeer(host.ID, task.ID, id),
+		ID:                   id,
+		Persistent:           persistent,
+		ConcurrentPieceCount: defaultConcurrentPieceCount,
+		FinishedPieces:       finishedPieces,
+		Task:                 task,
+		Host:                 host,
+		BlockParents:         blockParents,
+		Cost:                 cost,
+		CreatedAt:            createdAt,
+		UpdatedAt:            updatedAt,
+		Log:                  logger.WithPeer(host.ID, task.ID, id),
 	}
 
 	// Initialize state machine.
@@ -155,5 +175,9 @@ func NewPeer(id, state string, persistent bool, finishedPieces *bitset.BitSet, b
 	)
 	p.FSM.SetState(state)
 
+	for _, opt := range options {
+		opt(p)
+	}
+
 	return p
 }

scheduler/resource/persistentcache/peer_manager.go

@@ -157,6 +157,12 @@ func (p *peerManager) Load(ctx context.Context, peerID string) (*Peer, bool) {
 		return nil, false
 	}
 
+	concurrentPieceCount, err := strconv.ParseUint(rawPeer["concurrent_piece_count"], 10, 32)
+	if err != nil {
+		log.Errorf("parsing concurrent piece count failed: %v", err)
+		return nil, false
+	}
+
 	return NewPeer(
 		rawPeer["id"],
 		rawPeer["state"],
@@ -169,6 +175,7 @@ func (p *peerManager) Load(ctx context.Context, peerID string) (*Peer, bool) {
 		createdAt,
 		updatedAt,
 		logger.WithPeer(host.ID, task.ID, rawPeer["id"]),
+		WithConcurrentPieceCount(uint32(concurrentPieceCount)),
 	), true
 }
 
@@ -211,6 +218,7 @@ local cost = ARGV[8]
 local created_at = ARGV[9]
 local updated_at = ARGV[10]
 local ttl_seconds = tonumber(ARGV[11])
+local concurrent_piece_count = ARGV[12]
 
 -- Store peer information
 redis.call("HSET", peer_key,
@@ -223,7 +231,8 @@ redis.call("HSET", peer_key,
     "host_id", host_id,
     "cost", cost,
     "created_at", created_at,
-    "updated_at", updated_at)
+    "updated_at", updated_at,
+		"concurrent_piece_count", concurrent_piece_count)
 
 -- Set expiration for the peer key
 redis.call("EXPIRE", peer_key, ttl_seconds)
@@ -268,6 +277,7 @@ return true
 		peer.Cost.Nanoseconds(),
 		peer.CreatedAt.Format(time.RFC3339),
 		peer.UpdatedAt.Format(time.RFC3339),
+		peer.ConcurrentPieceCount,
 		remainingTTLSeconds,
 	}