From e95eeb2ae5f16a5c3262b751e2abfdde5791a5f9 Mon Sep 17 00:00:00 2001
From: Adin Schmahmann <adin.schmahmann@gmail.com>
Date: Fri, 26 Jul 2024 06:47:59 -0400
Subject: [PATCH] feat(routing/providerquerymanager): allow for configurable
 WithMaxProviders option

---
 bitswap/client/client.go                           |  2 +-
 .../providerquerymanager/providerquerymanager.go   | 14 ++++++++++++--
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/bitswap/client/client.go b/bitswap/client/client.go
index 35aff260c..ecf80c765 100644
--- a/bitswap/client/client.go
+++ b/bitswap/client/client.go
@@ -170,7 +170,7 @@ func New(parent context.Context, network bsnet.BitSwapNetwork, bstore blockstore
 	var defaultQueryManager *rpqm.ProviderQueryManager
 	if bs.useDefaultLookupManagement {
 		var err error
-		defaultQueryManager, err = rpqm.New(ctx, network)
+		defaultQueryManager, err = rpqm.New(ctx, network, rpqm.WithMaxProviders(10))
 		if err != nil {
 			// Should not be possible to hit this
 			panic(err)
diff --git a/routing/providerquerymanager/providerquerymanager.go b/routing/providerquerymanager/providerquerymanager.go
index af3ae7813..38c4009cd 100644
--- a/routing/providerquerymanager/providerquerymanager.go
+++ b/routing/providerquerymanager/providerquerymanager.go
@@ -16,8 +16,8 @@ import (
 var log = logging.Logger("routing/provqrymgr")
 
 const (
-	maxProviders                = 10
 	defaultMaxInProcessRequests = 6
+	defaultMaxProviders         = 0
 	defaultTimeout              = 10 * time.Second
 )
 
@@ -85,6 +85,7 @@ type ProviderQueryManager struct {
 	findProviderTimeout time.Duration
 	timeoutMutex        sync.RWMutex
 
+	maxProviders         int
 	maxInProcessRequests int
 
 	// do not touch outside the run loop
@@ -108,6 +109,14 @@ func WithMaxInProcessRequests(count int) Option {
 	}
 }
 
+// WithMaxProviders is the maximum number of providers that will be looked up per query
+func WithMaxProviders(count int) Option {
+	return func(mgr *ProviderQueryManager) error {
+		mgr.maxProviders = count
+		return nil
+	}
+}
+
 // New initializes a new ProviderQueryManager for a given context and a given
 // network provider.
 func New(ctx context.Context, network ProviderQueryNetwork, opts ...Option) (*ProviderQueryManager, error) {
@@ -120,6 +129,7 @@ func New(ctx context.Context, network ProviderQueryNetwork, opts ...Option) (*Pr
 		inProgressRequestStatuses:    make(map[cid.Cid]*inProgressRequestStatus),
 		findProviderTimeout:          defaultTimeout,
 		maxInProcessRequests:         defaultMaxInProcessRequests,
+		maxProviders:                 defaultMaxProviders,
 	}
 
 	for _, o := range opts {
@@ -275,7 +285,7 @@ func (pqm *ProviderQueryManager) findProviderWorker() {
 			pqm.timeoutMutex.RUnlock()
 			span := trace.SpanFromContext(findProviderCtx)
 			span.AddEvent("StartFindProvidersAsync")
-			providers := pqm.network.FindProvidersAsync(findProviderCtx, k, maxProviders)
+			providers := pqm.network.FindProvidersAsync(findProviderCtx, k, pqm.maxProviders)
 			wg := &sync.WaitGroup{}
 			for p := range providers {
 				wg.Add(1)