Repository: libp2p/go-libp2p-kad-dht
Pull request #1081: fix/updates to use context passed in New function for context cancellation
Author: login: blockchainluffy
Created: 2025-05-01T10:20:29Z | Merged: 2025-05-06T11:03:02Z



Diff:
diff --git a/dht.go b/dht.go
index e82fa3f6f..1cf10f88d 100644
--- a/dht.go
+++ b/dht.go
@@ -195,7 +195,7 @@ func New(ctx context.Context, h host.Host, options ...Option) (*IpfsDHT, error)
 		return nil, err
 	}
 
-	dht, err := makeDHT(h, cfg)
+	dht, err := makeDHT(ctx, h, cfg)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create DHT, err=%s", err)
 	}
@@ -280,7 +280,7 @@ func NewDHTClient(ctx context.Context, h host.Host, dstore ds.Batching) *IpfsDHT
 	return dht
 }
 
-func makeDHT(h host.Host, cfg dhtcfg.Config) (*IpfsDHT, error) {
+func makeDHT(ctx context.Context, h host.Host, cfg dhtcfg.Config) (*IpfsDHT, error) {
 	var protocols, serverProtocols []protocol.ID
 
 	v1proto := cfg.ProtocolPrefix + kad1
@@ -352,20 +352,20 @@ func makeDHT(h host.Host, cfg dhtcfg.Config) (*IpfsDHT, error) {
 		dht.optProvJobsPool = make(chan struct{}, cfg.OptimisticProvideJobsPoolSize)
 	}
 
+	// create a tagged context derived from the original context
+	// the DHT context should be done when the process is closed
+	dht.ctx, dht.cancel = context.WithCancel(dht.newContextWithLocalTags(ctx))
+
 	// rt refresh manager
 	dht.rtRefreshManager, err = makeRtRefreshManager(dht, cfg, maxLastSuccessfulOutboundThreshold)
 	if err != nil {
 		return nil, fmt.Errorf("failed to construct RT Refresh Manager,err=%s", err)
 	}
 
-	// create a tagged context derived from the original context
-	// the DHT context should be done when the process is closed
-	dht.ctx, dht.cancel = context.WithCancel(dht.newContextWithLocalTags(context.Background()))
-
 	if cfg.ProviderStore != nil {
 		dht.providerStore = cfg.ProviderStore
 	} else {
-		dht.providerStore, err = providers.NewProviderManager(h.ID(), dht.peerstore, cfg.Datastore)
+		dht.providerStore, err = providers.NewProviderManager(dht.ctx, h.ID(), dht.peerstore, cfg.Datastore)
 		if err != nil {
 			return nil, fmt.Errorf("initializing default provider manager (%v)", err)
 		}
@@ -402,6 +402,7 @@ func makeRtRefreshManager(dht *IpfsDHT, cfg dhtcfg.Config, maxLastSuccessfulOutb
 	}
 
 	r, err := rtrefresh.NewRtRefreshManager(
+		dht.ctx,
 		dht.host, dht.routingTable, cfg.RoutingTable.AutoRefresh,
 		keyGenFnc,
 		queryFnc,
diff --git a/fullrt/dht.go b/fullrt/dht.go
index 3ac3a65db..9f716d0cb 100644
--- a/fullrt/dht.go
+++ b/fullrt/dht.go
@@ -184,7 +184,7 @@ func NewFullRT(h host.Host, protocolPrefix protocol.ID, options ...Option) (*Ful
 	ctx, cancel := context.WithCancel(context.Background())
 
 	self := h.ID()
-	pm, err := providers.NewProviderManager(self, h.Peerstore(), dhtcfg.Datastore, fullrtcfg.pmOpts...)
+	pm, err := providers.NewProviderManager(ctx, self, h.Peerstore(), dhtcfg.Datastore, fullrtcfg.pmOpts...)
 	if err != nil {
 		cancel()
 		return nil, err
diff --git a/internal/metrics/metrics.go b/internal/metrics/metrics.go
index 30da4ae5f..3cba483d4 100644
--- a/internal/metrics/metrics.go
+++ b/internal/metrics/metrics.go
@@ -171,6 +171,6 @@ func RecordMessageSendErr(ctx context.Context) {
 	sentMessageErrors.Add(ctx, 1, attrSetOpt)
 }
 
-func RecordNetworkSize(ns int64) {
-	networkSize.Record(context.Background(), int64(ns))
+func RecordNetworkSize(ctx context.Context, ns int64) {
+	networkSize.Record(ctx, int64(ns))
 }
diff --git a/lookup.go b/lookup.go
index 91d0bd4ce..90bb92c97 100644
--- a/lookup.go
+++ b/lookup.go
@@ -48,7 +48,7 @@ func (dht *IpfsDHT) GetClosestPeers(ctx context.Context, key string) ([]peer.ID,
 	}
 
 	if ns, err := dht.nsEstimator.NetworkSize(); err == nil {
-		metrics.RecordNetworkSize(int64(ns))
+		metrics.RecordNetworkSize(dht.ctx, int64(ns))
 	}
 
 	// Reset the refresh timer for this key's bucket since we've just
diff --git a/lookup_optim.go b/lookup_optim.go
index 015a48242..0973c56b3 100644
--- a/lookup_optim.go
+++ b/lookup_optim.go
@@ -114,12 +114,12 @@ func (dht *IpfsDHT) optimisticProvide(outerCtx context.Context, keyMH multihash.
 		return errors.New("can't lookup empty key")
 	}
 
-	// initialize new context for all putProvider operations.
+	// use dht.ctx for all putProvider operations.
 	// We don't want to give the outer context to the put operations as we return early before all
 	// put operations have finished to avoid the long tail of the latency distribution. If we
 	// provided the outer context the put operations may be cancelled depending on what happens
 	// with the context on the user side.
-	putCtx, putCtxCancel := context.WithTimeout(context.Background(), time.Minute)
+	putCtx, putCtxCancel := context.WithTimeout(dht.ctx, time.Minute)
 
 	es, err := dht.newOptimisticState(putCtx, key)
 	if err != nil {
@@ -177,7 +177,7 @@ func (dht *IpfsDHT) optimisticProvide(outerCtx context.Context, keyMH multihash.
 	}
 
 	if ns, err := dht.nsEstimator.NetworkSize(); err == nil {
-		metrics.RecordNetworkSize(int64(ns))
+		metrics.RecordNetworkSize(dht.ctx, int64(ns))
 	}
 
 	// refresh the cpl for this key as the query was successful
diff --git a/providers/providers_manager.go b/providers/providers_manager.go
index 67da926f8..ebe569230 100644
--- a/providers/providers_manager.go
+++ b/providers/providers_manager.go
@@ -117,7 +117,7 @@ type getProv struct {
 }
 
 // NewProviderManager constructor
-func NewProviderManager(local peer.ID, ps peerstore.Peerstore, dstore ds.Batching, opts ...Option) (*ProviderManager, error) {
+func NewProviderManager(ctx context.Context, local peer.ID, ps peerstore.Peerstore, dstore ds.Batching, opts ...Option) (*ProviderManager, error) {
 	pm := new(ProviderManager)
 	pm.self = local
 	pm.getprovs = make(chan *getProv)
@@ -133,7 +133,7 @@ func NewProviderManager(local peer.ID, ps peerstore.Peerstore, dstore ds.Batchin
 	if err := pm.applyOptions(opts...); err != nil {
 		return nil, err
 	}
-	pm.ctx, pm.cancel = context.WithCancel(context.Background())
+	pm.ctx, pm.cancel = context.WithCancel(ctx)
 	pm.run()
 	return pm, nil
 }
diff --git a/providers/providers_manager_test.go b/providers/providers_manager_test.go
index a8f755cba..3d6b2ac11 100644
--- a/providers/providers_manager_test.go
+++ b/providers/providers_manager_test.go
@@ -31,7 +31,7 @@ func TestProviderManager(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	p, err := NewProviderManager(mid, ps, dssync.MutexWrap(ds.NewMapDatastore()))
+	p, err := NewProviderManager(ctx, mid, ps, dssync.MutexWrap(ds.NewMapDatastore()))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -77,7 +77,7 @@ func TestProvidersDatastore(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	p, err := NewProviderManager(mid, ps, dssync.MutexWrap(ds.NewMapDatastore()))
+	p, err := NewProviderManager(ctx, mid, ps, dssync.MutexWrap(ds.NewMapDatastore()))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -164,7 +164,7 @@ func TestProvidesExpire(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	p, err := NewProviderManager(mid, ps, ds)
+	p, err := NewProviderManager(ctx, mid, ps, ds)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -278,7 +278,7 @@ func TestLargeProvidersSet(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	p, err := NewProviderManager(mid, ps, dstore)
+	p, err := NewProviderManager(ctx, mid, ps, dstore)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -318,7 +318,7 @@ func TestUponCacheMissProvidersAreReadFromDatastore(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	pm, err := NewProviderManager(p1, ps, dssync.MutexWrap(ds.NewMapDatastore()))
+	pm, err := NewProviderManager(ctx, p1, ps, dssync.MutexWrap(ds.NewMapDatastore()))
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -347,7 +347,7 @@ func TestWriteUpdatesCache(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	pm, err := NewProviderManager(p1, ps, dssync.MutexWrap(ds.NewMapDatastore()))
+	pm, err := NewProviderManager(ctx, p1, ps, dssync.MutexWrap(ds.NewMapDatastore()))
 	if err != nil {
 		t.Fatal(err)
 	}
diff --git a/rtrefresh/rt_refresh_manager.go b/rtrefresh/rt_refresh_manager.go
index a3280581a..e3a620e24 100644
--- a/rtrefresh/rt_refresh_manager.go
+++ b/rtrefresh/rt_refresh_manager.go
@@ -57,16 +57,16 @@ type RtRefreshManager struct {
 	refreshDoneCh chan struct{} // write to this channel after every refresh
 }
 
-func NewRtRefreshManager(h host.Host, rt *kbucket.RoutingTable, autoRefresh bool,
+func NewRtRefreshManager(ctx context.Context, h host.Host, rt *kbucket.RoutingTable, autoRefresh bool,
 	refreshKeyGenFnc func(cpl uint) (string, error),
 	refreshQueryFnc func(ctx context.Context, key string) error,
 	refreshPingFnc func(ctx context.Context, p peer.ID) error,
 	refreshQueryTimeout time.Duration,
 	refreshInterval time.Duration,
 	successfulOutboundQueryGracePeriod time.Duration,
-	refreshDoneCh chan struct{}) (*RtRefreshManager, error) {
-
-	ctx, cancel := context.WithCancel(context.Background())
+	refreshDoneCh chan struct{},
+) (*RtRefreshManager, error) {
+	ctx, cancel := context.WithCancel(ctx)
 	return &RtRefreshManager{
 		ctx:       ctx,
 		cancel:    cancel,


Reviews:
user: login: guillaumemichel
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-04T07:37:06Z
user: login: blockchainluffy
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T07:15:07Z
user: login: guillaumemichel
body: Reasoning for the changes I just pushed
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-06T09:25:10Z
user: login: guillaumemichel
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-05-06T11:02:50Z

Review comments:
diff_hunk: @@ -170,6 +171,8 @@ func (r *RtRefreshManager) pingAndEvictPeers(ctx context.Context) {
 			}
 
 			atomic.AddInt64(&alive, 1)
+
+			<-r.ctx.Done()
path: rtrefresh/rt_refresh_manager.go
user: login: guillaumemichel
body: Why is this necessary?
created_at: 2025-05-04T07:37:06Z
author_association: COLLABORATOR
diff_hunk: @@ -170,6 +171,8 @@ func (r *RtRefreshManager) pingAndEvictPeers(ctx context.Context) {
 			}
 
 			atomic.AddInt64(&alive, 1)
+
+			<-r.ctx.Done()
path: rtrefresh/rt_refresh_manager.go
user: login: blockchainluffy
body: This was just to make sure `pingAndEvictPeers` go routine also ends when context is cancelled.
LMK, if I am missing something!
created_at: 2025-05-05T07:15:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -151,7 +151,7 @@ func (pm *ProviderManager) run() {
 			if gcQuery != nil {
 				gcQuery.Close()
 			}
-			if err := pm.dstore.Flush(context.Background()); err != nil {
+			if err := pm.dstore.Flush(pm.ctx); err != nil {
path: providers/providers_manager.go
user: login: guillaumemichel
body: This needs to remain `context.Background` as it is executed after `pm.ctx` is cancelled.
created_at: 2025-05-06T09:16:31Z
author_association: COLLABORATOR
diff_hunk: @@ -119,7 +119,7 @@ func (dht *IpfsDHT) optimisticProvide(outerCtx context.Context, keyMH multihash.
 	// put operations have finished to avoid the long tail of the latency distribution. If we
 	// provided the outer context the put operations may be cancelled depending on what happens
 	// with the context on the user side.
path: lookup_optim.go
user: login: guillaumemichel
body: Updating comment to match code
created_at: 2025-05-06T09:17:00Z
author_association: COLLABORATOR
diff_hunk: @@ -2504,3 +2504,53 @@ func TestAddrFilter(t *testing.T) {
 	}
 	require.Equal(t, len(publicAddrs)+len(privAddrs), len(d3.host.Peerstore().Addrs(peerid)))
 }
+
+func TestContextCancellationStopsGoroutines(t *testing.T) {
path: dht_test.go
user: login: guillaumemichel
body: Removing test since we cannot afford to sleep for 10 seconds in the CI. Also there is a delay before `RtRefreshManager` and `ProviderManager` close, so we might just skip this test.
created_at: 2025-05-06T09:18:47Z
author_association: COLLABORATOR
diff_hunk: @@ -352,20 +352,20 @@ func makeDHT(h host.Host, cfg dhtcfg.Config) (*IpfsDHT, error) {
 		dht.optProvJobsPool = make(chan struct{}, cfg.OptimisticProvideJobsPoolSize)
 	}
 
+	// create a tagged context derived from the original context
+	// the DHT context should be done when the process is closed
+	dht.ctx, dht.cancel = context.WithCancel(dht.newContextWithLocalTags(ctx))
+
 	// rt refresh manager
 	dht.rtRefreshManager, err = makeRtRefreshManager(dht, cfg, maxLastSuccessfulOutboundThreshold)
 	if err != nil {
 		return nil, fmt.Errorf("failed to construct RT Refresh Manager,err=%s", err)
 	}
 
-	// create a tagged context derived from the original context
-	// the DHT context should be done when the process is closed
-	dht.ctx, dht.cancel = context.WithCancel(dht.newContextWithLocalTags(context.Background()))
-
 	if cfg.ProviderStore != nil {
 		dht.providerStore = cfg.ProviderStore
 	} else {
-		dht.providerStore, err = providers.NewProviderManager(h.ID(), dht.peerstore, cfg.Datastore)
+		dht.providerStore, err = providers.NewProviderManager(ctx, h.ID(), dht.peerstore, cfg.Datastore)
path: dht.go
user: login: guillaumemichel
body: ```suggestion
		dht.providerStore, err = providers.NewProviderManager(dht.ctx, h.ID(), dht.peerstore, cfg.Datastore)
```
created_at: 2025-05-06T09:19:38Z
author_association: COLLABORATOR

PR comments:
user: login: guillaumemichel
created_at: 2025-05-01T11:27:18Z
body: Thanks @blockchainluffy for your contribution!

Would you mind adding some tests to verify that all components stop after context is cancelled?
author_association: COLLABORATOR
user: login: blockchainluffy
created_at: 2025-05-02T12:04:49Z
body: > Thanks @blockchainluffy for your contribution!
> 
> Would you mind adding some tests to verify that all components stop after context is cancelled?

added test for context cancellation. LMK of any further changes required!
author_association: CONTRIBUTOR
user: login: guillaumemichel
created_at: 2025-05-06T09:25:59Z
body: @blockchainluffy I made some light changes to the PR. WDYT?
author_association: COLLABORATOR
user: login: blockchainluffy
created_at: 2025-05-06T10:59:12Z
body: > @blockchainluffy I made some light changes to the PR. WDYT?

Thank you for the changes, they look good to me!
author_association: CONTRIBUTOR