fix: Set context timeout for queryPeer

libp2p · Oct 23, 2024 · 21b4c47 · 21b4c47
1 parent 8170ac4
commit 21b4c47
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/crawler/crawler.go b/crawler/crawler.go
@@ -34,6 +34,7 @@ type (
 	DefaultCrawler struct {
 		parallelism          int
 		connectTimeout       time.Duration
+		queryTimeout         time.Duration
 		host                 host.Host
 		dhtRPC               *pb.ProtocolMessenger
 		dialAddressExtendDur time.Duration
@@ -60,6 +61,7 @@ func NewDefaultCrawler(host host.Host, opts ...Option) (*DefaultCrawler, error)
 	return &DefaultCrawler{
 		parallelism:          o.parallelism,
 		connectTimeout:       o.connectTimeout,
+		queryTimeout:         3 * o.connectTimeout,
 		host:                 host,
 		dhtRPC:               pm,
 		dialAddressExtendDur: o.dialAddressExtendDur,
@@ -144,6 +146,8 @@ func (c *DefaultCrawler) Run(ctx context.Context, startingPeers []*peer.AddrInfo
 	for i := 0; i < c.parallelism; i++ {
 		go func() {
 			defer wg.Done()
+			ctx, cancel := context.WithTimeout(ctx, c.queryTimeout)
+			defer cancel()
 			for p := range jobs {
 				res := c.queryPeer(ctx, p)
 				results <- res