fix/updates to use context passed in New function for context cancellation (#1081)

* fix/updates to use context passed in New function for context cancellation

* updated test to fix CI

* fix context

---------

Co-authored-by: guillaumemichel <guillaume@michel.id>

Author: blockchainluffy

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,

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

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))
 }

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

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

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
 }

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)
 	}

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,