integrate Select

Author: Wondertan

fibre/client_download.go

@@ -4,8 +4,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"math/rand/v2"
-	"slices"
 	"sync/atomic"
 
 	"github.com/celestiaorg/celestia-app/v6/fibre/validator"
@@ -26,14 +24,14 @@ var (
 
 // Download retrieves and reconstructs [Blob] by [Commitment] from the [Server]s.
 //
-// The algorithm is simple: take a randomized 2/3 subset of the latest validators (instead of the complete set)
-// and request blob shards from them. In the happy path, this is sufficient to reconstruct blobs.
-// For the unhappy path, we add recovery requests to additional validators if any of the initial requests fail,
-// so that in the worst case, we cover all validators.
+// The algorithm selects minimal required number of validators,
+// shuffled by stake weight for load balancing and requests them for shards.
+// If any of the requests fails, more validators are requested until enough shards are retrieved or
+// the maximum number of validators is reached. In the happy case, the operation succeeds in a single roundtrip.
 //
 // Errors:
 //   - [ErrNotFound]: no shard was retrieved for the blob
-//   - [ErrNotEnoughShards]: not enough rows were retrieved to reconstruct the original data
+//   - [ErrNotEnoughShards]: not enough shards were retrieved to reconstruct the original data
 //   - [ErrInvalidCommitment]: the commitment doesn't match the reconstructed blob
 func (c *Client) Download(ctx context.Context, commitment Commitment) (*Blob, error) {
 	if c.closed.Load() {
@@ -192,19 +190,9 @@ func (c *Client) downloadBlob(
 		downloadedCh = make(chan struct{}) // closes when downloadTarget amount of responses complete
 	)
 
-	var (
-		// limit to download minimum required amount of shards in the best case, instead of everything
-		downloadTarget = max(valSet.Size()*int(c.cfg.UploadTargetSignaturesCount.Numerator)/
-			int(c.cfg.UploadTargetSignaturesCount.Denominator), 1)
-		downloadLimitCh = make(chan struct{}, downloadTarget)
-	)
-
-	// shuffle validators for random prioritization
-	// TODO(@Wondertan): Order validators based on their performance over time using EWMA
-	validators := slices.Clone(valSet.Validators)
-	rand.Shuffle(len(validators), func(i, j int) {
-		validators[i], validators[j] = validators[j], validators[i]
-	})
+	// select validators shuffled by stake for load balancing
+	validators, downloadTarget := valSet.Select(blob.Config().OriginalRows, c.cfg.MinRowsPerValidator, c.cfg.LivenessThreshold)
+	downloadLimitCh := make(chan struct{}, downloadTarget)
 
 loop:
 	for _, val := range validators {

fibre/client_download_test.go

@@ -97,7 +97,7 @@ func testClientDownloadClosedClient(t *testing.T) {
 
 func testClientDownloadExactTargetCount(t *testing.T) {
 	// test that we download from exactly downloadTarget validators (no more)
-	// with 10 validators and 2/3 target, downloadTarget = 6
+	// with 10 equal-stake validators and livenessThreshold=1/3, downloadTarget = 4
 	const numValidators = 10
 
 	blob := makeTestBlobV0(t, 256*1024)
@@ -112,13 +112,14 @@ func testClientDownloadExactTargetCount(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, blob.Data(), downloaded.Data())
 
-	// downloadTarget = 10 * 2/3 = 6
-	// we should have exactly 6 successful downloads (no over-fetching in happy path)
-	require.Equal(t, int64(6), counter.Load(), "should download from exactly downloadTarget validators")
+	// Select returns minRequired = 4 for 10 equal-stake validators with livenessThreshold=1/3
+	// we should have exactly 4 successful downloads (no over-fetching in happy path)
+	require.Equal(t, int64(4), counter.Load(), "should download from exactly downloadTarget validators")
 }
 
 func testClientDownloadFaultTolerance(t *testing.T) {
-	// test failure tolerance boundaries with 10 validators and 2/3 target
+	// test failure tolerance boundaries with 10 validators
+	// Select uses livenessThreshold (1/3), so downloadTarget = 4 for 10 equal-stake validators
 	const numValidators = 10
 	blob := makeTestBlobV0(t, 256*1024)
 
@@ -127,10 +128,10 @@ func testClientDownloadFaultTolerance(t *testing.T) {
 		expectErr error
 	}{
 		{10, fibre.ErrNotFound},
-		{6, fibre.ErrNotEnoughShards}, // 4 successes, need 6
-		{5, fibre.ErrNotEnoughShards}, // 5 successes, need 6
-		{4, nil},                      // 6 successes, exactly enough
-		{3, nil},                      // 7 successes, more than enough
+		{7, fibre.ErrNotEnoughShards}, // 3 successes, need 4
+		{6, nil},                      // 4 successes, exactly enough
+		{5, nil},                      // 5 successes, more than enough
+		{4, nil},                      // 6 successes, more than enough
 	}
 
 	for _, tc := range tests {