Repository: ipfs/boxo

Issue:
bitswap/session: context-aware lifecycle
I'm trying to trace a content retrieval request via `ipfs get ...` with the actual bitswap interactions.

This is tricky though because a bitswap client session creates a [new disparate context](https://github.com/ipfs/boxo/blob/0f1ebacc7662d946946e3315c49241f6fb188b7c/bitswap/client/internal/session/session.go#L164). Then I also found [this comment](https://github.com/ipfs/boxo/blob/0f1ebacc7662d946946e3315c49241f6fb188b7c/bitswap/client/internal/session/session.go#L209) for the `Close()` method which only cancels the session context:

```go
// Session lifecycle is independent of the context used to create requests -
// canceling a request context does not close the session. When context-based
// automation is desired, Client.NewSession(ctx) uses context.AfterFunc to
// automatically close sessions when the context is canceled.
func (s *Session) Close() {
	s.cancel()
}
```

@gammazero I saw that this comment is from you. Is the meaning descriptive or normative? Is it just describing how it's working or stating how it's supposed to work?

Looking at the session lifecycle (checking where new sessions are created and closed) I can see that they are always tied to the context that instantiated the session. This to me indicates that a separate context for the session isn't needed which would enable proper context propagation and improve the tracing. As a nice byproduct, the `retrievalState` wouldn't need to be passed manually through the stack.

New sessions are created here:

https://github.com/ipfs/boxo/blob/0f1ebacc7662d946946e3315c49241f6fb188b7c/bitswap/client/client.go#L442-L483

and here:

https://github.com/ipfs/boxo/blob/0f1ebacc7662d946946e3315c49241f6fb188b7c/bitswap/client/client.go#L702-L725

Both calls tie the new session to the incoming context.

The same argument then actually trickles down to the `sessionWantSender` as its lifecycle is tied to the sessions':

https://github.com/ipfs/boxo/blob/0f1ebacc7662d946946e3315c49241f6fb188b7c/bitswap/client/internal/session/sessionwantsender.go#L117-L142

Discussion:
@lidel: Triage notes:
- we will look into this a part of wider investigation into traces regressions / need of improvement
- revisit once boxo has mentioned fixes@gammazero: > Is the meaning descriptive or normative

Is is supposed to be descriptive, but we may have an issue here causing this to not work as expected (in which case the comment becomes normative). We are currently working on the issue which is affecting our ability to collect tracing results.@gammazero: Here is the necessary fix: https://github.com/ipfs/boxo/pull/1059

Solution patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 951dbd7e..a6991ba7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -29,6 +29,7 @@ The following emojis are used to highlight certain changes:
 ### Fixed
 
 - `gateway`: Fixed duplicate peer IDs appearing in retrieval timeout error messages
+- `bitswap/client`: fix tracing by using context to pass trace and retrieval state to session [#1059](https://github.com/ipfs/boxo/pull/1059)
 
 ### Security
 
diff --git a/bitswap/client/client.go b/bitswap/client/client.go
index 01b836e7..53f2300b 100644
--- a/bitswap/client/client.go
+++ b/bitswap/client/client.go
@@ -299,6 +299,7 @@ func New(parent context.Context, network bsnet.BitSwapNetwork, providerFinder ro
 	}
 
 	sessionFactory := func(
+		sessctx context.Context,
 		sessmgr bssession.SessionManager,
 		id uint64,
 		spm bssession.SessionPeerManager,
@@ -309,7 +310,6 @@ func New(parent context.Context, network bsnet.BitSwapNetwork, providerFinder ro
 		provSearchDelay time.Duration,
 		rebroadcastDelay time.Duration,
 		self peer.ID,
-		retrievalState *retrieval.State,
 	) bssm.Session {
 		// careful when bs.pqm is nil. Since we are type-casting it
 		// into routing.ContentDiscovery when passing it, it will become
@@ -321,7 +321,7 @@ func New(parent context.Context, network bsnet.BitSwapNetwork, providerFinder ro
 		} else if providerFinder != nil {
 			sessionProvFinder = providerFinder
 		}
-		return bssession.New(sessmgr, id, spm, sessionProvFinder, sim, pm, bpm, notif, provSearchDelay, rebroadcastDelay, self, bs.havesReceivedGauge, retrievalState)
+		return bssession.New(sessctx, sessmgr, id, spm, sessionProvFinder, sim, pm, bpm, notif, provSearchDelay, rebroadcastDelay, self, bs.havesReceivedGauge)
 	}
 	sessionPeerManagerFactory := func(id uint64) bssession.SessionPeerManager {
 		return bsspm.New(id, network)
@@ -446,11 +446,18 @@ func (bs *Client) GetBlocks(ctx context.Context, keys []cid.Cid) (<-chan blocks.
 	ctx, span := internal.StartSpan(ctx, "GetBlocks", trace.WithAttributes(attribute.Int("NumKeys", len(keys))))
 	defer span.End()
 
-	session := bs.sm.NewSession(bs.provSearchDelay, bs.rebroadcastDelay, retrieval.StateFromContext(ctx))
+	// Temporary session closed indepentendly of cancellation ctx.
+	sessCtx, cancelSession := context.WithCancel(context.Background())
+
+	// Preserve retrieval.State from the original request context if present
+	if retrievalState := retrieval.StateFromContext(ctx); retrievalState != nil {
+		sessCtx = context.WithValue(sessCtx, retrieval.ContextKey, retrievalState)
+	}
+	session := bs.sm.NewSession(sessCtx, bs.provSearchDelay, bs.rebroadcastDelay)
 
 	blocksChan, err := session.GetBlocks(ctx, keys)
 	if err != nil {
-		session.Close()
+		cancelSession()
 		return nil, err
 	}
 
@@ -458,7 +465,7 @@ func (bs *Client) GetBlocks(ctx context.Context, keys []cid.Cid) (<-chan blocks.
 	go func() {
 		defer func() {
 			close(out)
-			session.Close()
+			cancelSession()
 		}()
 
 		ctxDone := ctx.Done()
@@ -717,9 +724,5 @@ func (bs *Client) NewSession(ctx context.Context) exchange.Fetcher {
 	ctx, span := internal.StartSpan(ctx, "NewSession")
 	defer span.End()
 
-	session := bs.sm.NewSession(bs.provSearchDelay, bs.rebroadcastDelay, retrieval.StateFromContext(ctx))
-	context.AfterFunc(ctx, func() {
-		session.Close()
-	})
-	return session
+	return bs.sm.NewSession(ctx, bs.provSearchDelay, bs.rebroadcastDelay)
 }
diff --git a/bitswap/client/internal/session/session.go b/bitswap/client/internal/session/session.go
index 1c61126b..9c8c1936 100644
--- a/bitswap/client/internal/session/session.go
+++ b/bitswap/client/internal/session/session.go
@@ -10,7 +10,6 @@ import (
 	notifications "github.com/ipfs/boxo/bitswap/client/internal/notifications"
 	bspm "github.com/ipfs/boxo/bitswap/client/internal/peermanager"
 	bssim "github.com/ipfs/boxo/bitswap/client/internal/sessioninterestmanager"
-	"github.com/ipfs/boxo/retrieval"
 	blocks "github.com/ipfs/go-block-format"
 	cid "github.com/ipfs/go-cid"
 	logging "github.com/ipfs/go-log/v2"
@@ -133,13 +132,11 @@ type Session struct {
 	self peer.ID
 }
 
-// New creates a new bitswap session.
+// New creates a new bitswap session whose lifetime is bounded by the
+// given context.
 //
-// The session maintains its own internal context for operations and is not
-// tied to any external context lifecycle. The caller MUST call Close() when
-// the session is no longer needed to ensure proper cleanup. When sessions are
-// created via Client.NewSession(ctx), automatic cleanup via context.AfterFunc
-// is provided.
+// The caller MUST call Close() or cancel the context when the session is no
+// longer needed to ensure proper cleanup.
 //
 // The retrievalState parameter, if provided, enables diagnostic tracking of
 // the retrieval process. It is attached to the session's internal context and
@@ -147,6 +144,7 @@ type Session struct {
 // phases. This is particularly useful for debugging timeout errors and
 // understanding retrieval performance.
 func New(
+	ctx context.Context,
 	sm SessionManager,
 	id uint64,
 	sprm SessionPeerManager,
@@ -159,14 +157,8 @@ func New(
 	periodicSearchDelay time.Duration,
 	self peer.ID,
 	havesReceivedGauge bspm.Gauge,
-	retrievalState *retrieval.State,
 ) *Session {
-	ctx, cancel := context.WithCancel(context.Background())
-
-	// If a retrieval state is passed in, then keep it in the context.
-	if retrievalState != nil {
-		ctx = context.WithValue(ctx, retrieval.ContextKey, retrievalState)
-	}
+	ctx, cancel := context.WithCancel(ctx)
 
 	s := &Session{
 		sw:                  newSessionWants(broadcastLiveWantsLimit),
@@ -198,14 +190,10 @@ func (s *Session) ID() uint64 {
 	return s.id
 }
 
-// Close terminates the session and cleans up its resources. This method MUST
-// be called when the session is no longer needed to avoid resource leaks.
-// After calling Close, the session should not be used anymore.
-//
-// Session lifecycle is independent of the context used to create requests -
-// canceling a request context does not close the session. When context-based
-// automation is desired, Client.NewSession(ctx) uses context.AfterFunc to
-// automatically close sessions when the context is canceled.
+// Close terminates the session and cleans up its resources. This method must
+// be called, or the context used to create the session must be canceled, when
+// the session is no longer needed to avoid resource leaks. After calling
+// Close, the session should not be used anymore.
 func (s *Session) Close() {
 	s.cancel()
 }
diff --git a/bitswap/client/internal/sessionmanager/sessionmanager.go b/bitswap/client/internal/sessionmanager/sessionmanager.go
index 983b192b..33c3d438 100644
--- a/bitswap/client/internal/sessionmanager/sessionmanager.go
+++ b/bitswap/client/internal/sessionmanager/sessionmanager.go
@@ -12,7 +12,6 @@ import (
 	bssession "github.com/ipfs/boxo/bitswap/client/internal/session"
 	bssim "github.com/ipfs/boxo/bitswap/client/internal/sessioninterestmanager"
 	exchange "github.com/ipfs/boxo/exchange"
-	"github.com/ipfs/boxo/retrieval"
 	cid "github.com/ipfs/go-cid"
 	peer "github.com/libp2p/go-libp2p/core/peer"
 	"go.opentelemetry.io/otel/attribute"
@@ -20,8 +19,6 @@ import (
 )
 
 // Session represents a bitswap session managed by the SessionManager.
-// Sessions have their own lifecycle independent of request contexts and
-// must be explicitly closed via the Close() method when no longer needed.
 type Session interface {
 	exchange.Fetcher
 	ID() uint64
@@ -31,6 +28,7 @@ type Session interface {
 
 // SessionFactory generates a new session for the SessionManager to track.
 type SessionFactory func(
+	ctx context.Context,
 	sm bssession.SessionManager,
 	id uint64,
 	sprm bssession.SessionPeerManager,
@@ -40,8 +38,7 @@ type SessionFactory func(
 	notif notifications.PubSub,
 	provSearchDelay time.Duration,
 	rebroadcastDelay time.Duration,
-	self peer.ID,
-	retrievalState *retrieval.State) Session
+	self peer.ID) Session
 
 // PeerManagerFactory generates a new peer manager for a session.
 type PeerManagerFactory func(id uint64) bssession.SessionPeerManager
@@ -85,23 +82,16 @@ func New(sessionFactory SessionFactory, sessionInterestManager *bssim.SessionInt
 
 // NewSession initializes a session and adds to the session manager.
 //
-// The session is created with its own internal context and lifecycle management.
-// The retrievalState parameter, if provided, will be attached to the session's
-// internal context to enable diagnostic tracking of the retrieval process. This
-// includes tracking provider discovery attempts, peer connections, and data
-// retrieval phases, which is particularly useful for debugging timeout errors.
-//
-// The returned Session must be closed via its Close() method when no longer needed.
-// Note: When sessions are created via Client.NewSession(ctx), automatic cleanup
-// via context.AfterFunc is provided.
-func (sm *SessionManager) NewSession(provSearchDelay, rebroadcastDelay time.Duration, retrievalState *retrieval.State) Session {
+// The returned Session must be closed via its Close() method, or by canceling
+// the context, when no longer needed.
+func (sm *SessionManager) NewSession(ctx context.Context, provSearchDelay, rebroadcastDelay time.Duration) Session {
 	id := sm.GetNextSessionID()
 
-	_, span := internal.StartSpan(context.Background(), "SessionManager.NewSession", trace.WithAttributes(attribute.String("ID", strconv.FormatUint(id, 10))))
+	ctx, span := internal.StartSpan(ctx, "SessionManager.NewSession", trace.WithAttributes(attribute.String("ID", strconv.FormatUint(id, 10))))
 	defer span.End()
 
 	pm := sm.peerManagerFactory(id)
-	session := sm.sessionFactory(sm, id, pm, sm.sessionInterestManager, sm.peerManager, sm.blockPresenceManager, sm.notif, provSearchDelay, rebroadcastDelay, sm.self, retrievalState)
+	session := sm.sessionFactory(ctx, sm, id, pm, sm.sessionInterestManager, sm.peerManager, sm.blockPresenceManager, sm.notif, provSearchDelay, rebroadcastDelay, sm.self)
 
 	sm.sessLk.Lock()
 	if sm.sessions != nil { // check if SessionManager was shutdown



Test patch:
diff --git a/bitswap/client/internal/session/session_test.go b/bitswap/client/internal/session/session_test.go
index 9e8d820f..02c78a01 100644
--- a/bitswap/client/internal/session/session_test.go
+++ b/bitswap/client/internal/session/session_test.go
@@ -155,6 +155,9 @@ func (pm *fakePeerManager) BroadcastWantHaves(cids []cid.Cid) {
 func (pm *fakePeerManager) SendCancels(cancels []cid.Cid) {}
 
 func TestSessionGetBlocks(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -164,7 +167,7 @@ func TestSessionGetBlocks(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil, nil)
+	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil)
 	defer session.Close()
 	blks := random.BlocksOfSize(broadcastLiveWantsLimit*2, blockSize)
 	var cids []cid.Cid
@@ -172,7 +175,7 @@ func TestSessionGetBlocks(t *testing.T) {
 		cids = append(cids, block.Cid())
 	}
 
-	_, err := session.GetBlocks(context.Background(), cids)
+	_, err := session.GetBlocks(ctx, cids)
 	require.NoError(t, err, "error getting blocks")
 
 	// Wait for initial want request
@@ -234,6 +237,9 @@ func TestSessionGetBlocks(t *testing.T) {
 }
 
 func TestSessionFindMorePeers(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -243,7 +249,7 @@ func TestSessionFindMorePeers(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil, nil)
+	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil)
 	defer session.Close()
 	session.SetBaseTickDelay(200 * time.Microsecond)
 	blks := random.BlocksOfSize(broadcastLiveWantsLimit*2, blockSize)
@@ -252,9 +258,6 @@ func TestSessionFindMorePeers(t *testing.T) {
 		cids = append(cids, block.Cid())
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
-
 	_, err := session.GetBlocks(ctx, cids)
 	require.NoError(t, err, "error getting blocks")
 
@@ -306,6 +309,8 @@ func TestSessionFindMorePeers(t *testing.T) {
 }
 
 func TestSessionOnPeersExhausted(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -316,7 +321,7 @@ func TestSessionOnPeersExhausted(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil, nil)
+	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil)
 	defer session.Close()
 	blks := random.BlocksOfSize(broadcastLiveWantsLimit+5, blockSize)
 	var cids []cid.Cid
@@ -324,9 +329,6 @@ func TestSessionOnPeersExhausted(t *testing.T) {
 		cids = append(cids, block.Cid())
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
-	defer cancel()
-
 	_, err := session.GetBlocks(ctx, cids)
 	require.NoError(t, err, "error getting blocks")
 
@@ -347,6 +349,8 @@ func TestSessionOnPeersExhausted(t *testing.T) {
 }
 
 func TestSessionFailingToGetFirstBlock(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -356,7 +360,7 @@ func TestSessionFailingToGetFirstBlock(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, 10*time.Millisecond, 100*time.Millisecond, "", nil, nil)
+	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, 10*time.Millisecond, 100*time.Millisecond, "", nil)
 	defer session.Close()
 	blks := random.BlocksOfSize(4, blockSize)
 	var cids []cid.Cid
@@ -364,9 +368,6 @@ func TestSessionFailingToGetFirstBlock(t *testing.T) {
 		cids = append(cids, block.Cid())
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
-	defer cancel()
-
 	_, err := session.GetBlocks(ctx, cids)
 	require.NoError(t, err, "error getting blocks")
 
@@ -461,7 +462,9 @@ func TestSessionCtxCancelClosesGetBlocksChannel(t *testing.T) {
 	sm := newMockSessionMgr(sim)
 
 	// Create a new session with its own context
-	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil, nil)
+	sessctx, sesscancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer sesscancel()
+	session := New(sessctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil)
 	defer session.Close()
 
 	timerCtx, timerCancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
@@ -475,8 +478,8 @@ func TestSessionCtxCancelClosesGetBlocksChannel(t *testing.T) {
 	getBlocksCh, err := session.GetBlocks(getctx, []cid.Cid{blks[0].Cid()})
 	require.NoError(t, err, "error getting blocks")
 
-	// Cancel the session
-	session.Close()
+	// Cancel the session context
+	sesscancel()
 
 	// Expect the GetBlocks() channel to be closed
 	select {
@@ -493,6 +496,8 @@ func TestSessionCtxCancelClosesGetBlocksChannel(t *testing.T) {
 }
 
 func TestSessionOnShutdownCalled(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -504,7 +509,7 @@ func TestSessionOnShutdownCalled(t *testing.T) {
 	sm := newMockSessionMgr(sim)
 
 	// Create a new session
-	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil, nil)
+	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil)
 
 	// Shutdown the session
 	session.Close()
@@ -516,6 +521,8 @@ func TestSessionOnShutdownCalled(t *testing.T) {
 }
 
 func TestSessionReceiveMessageAfterClose(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -526,13 +533,13 @@ func TestSessionReceiveMessageAfterClose(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil, nil)
+	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, time.Minute, "", nil)
 	defer session.Close()
 
 	blks := random.BlocksOfSize(2, blockSize)
 	cids := []cid.Cid{blks[0].Cid(), blks[1].Cid()}
 
-	_, err := session.GetBlocks(context.Background(), cids)
+	_, err := session.GetBlocks(ctx, cids)
 	require.NoError(t, err, "error getting blocks")
 
 	// Wait for initial want request
diff --git a/bitswap/client/internal/sessionmanager/sessionmanager_test.go b/bitswap/client/internal/sessionmanager/sessionmanager_test.go
index 7aa58b50..7ae191fc 100644
--- a/bitswap/client/internal/sessionmanager/sessionmanager_test.go
+++ b/bitswap/client/internal/sessionmanager/sessionmanager_test.go
@@ -12,7 +12,6 @@ import (
 	bspm "github.com/ipfs/boxo/bitswap/client/internal/peermanager"
 	bssession "github.com/ipfs/boxo/bitswap/client/internal/session"
 	bssim "github.com/ipfs/boxo/bitswap/client/internal/sessioninterestmanager"
-	"github.com/ipfs/boxo/retrieval"
 	blocks "github.com/ipfs/go-block-format"
 	cid "github.com/ipfs/go-cid"
 	peer "github.com/libp2p/go-libp2p/core/peer"
@@ -84,6 +83,7 @@ func (fpm *fakePeerManager) cancelled() []cid.Cid {
 }
 
 func sessionFactory(
+	ctx context.Context,
 	sm bssession.SessionManager,
 	id uint64,
 	sprm bssession.SessionPeerManager,
@@ -94,7 +94,6 @@ func sessionFactory(
 	provSearchDelay time.Duration,
 	rebroadcastDelay time.Duration,
 	self peer.ID,
-	retrievalState *retrieval.State,
 ) Session {
 	fs := &fakeSession{
 		id:    id,
@@ -102,6 +101,10 @@ func sessionFactory(
 		sm:    sm,
 		notif: notif,
 	}
+	go func() {
+		<-ctx.Done()
+		sm.RemoveSession(fs.id)
+	}()
 	return fs
 }
 
@@ -122,11 +125,11 @@ func TestReceiveFrom(t *testing.T) {
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 
-	firstSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	firstSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer firstSession.Close()
-	secondSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	secondSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer secondSession.Close()
-	thirdSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	thirdSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer thirdSession.Close()
 
 	sim.RecordSessionInterest(firstSession.ID(), []cid.Cid{block.Cid()})
@@ -169,11 +172,11 @@ func TestReceiveBlocksWhenManagerShutdown(t *testing.T) {
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 
-	firstSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	firstSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer firstSession.Close()
-	secondSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	secondSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer secondSession.Close()
-	thirdSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	thirdSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer thirdSession.Close()
 
 	sim.RecordSessionInterest(firstSession.ID(), []cid.Cid{block.Cid()})
@@ -206,11 +209,11 @@ func TestReceiveBlocksWhenSessionContextCancelled(t *testing.T) {
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 
-	firstSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	firstSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer firstSession.Close()
-	secondSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	secondSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer secondSession.Close()
-	thirdSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	thirdSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer thirdSession.Close()
 
 	sim.RecordSessionInterest(firstSession.ID(), []cid.Cid{block.Cid()})
@@ -243,7 +246,7 @@ func TestShutdown(t *testing.T) {
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 	cids := []cid.Cid{block.Cid()}
-	firstSession := sm.NewSession(time.Second, time.Minute, nil).(*fakeSession)
+	firstSession := sm.NewSession(ctx, time.Second, time.Minute).(*fakeSession)
 	defer firstSession.Close()
 	sim.RecordSessionInterest(firstSession.ID(), cids)
 	sm.ReceiveFrom(ctx, p, []cid.Cid{}, []cid.Cid{}, cids)
diff --git a/namesys/republisher/repub_test.go b/namesys/republisher/repub_test.go
index ee7fa611..04e79283 100644
--- a/namesys/republisher/repub_test.go
+++ b/namesys/republisher/repub_test.go
@@ -92,9 +92,9 @@ func TestRepublish(t *testing.T) {
 	// Retry in case the record expires before we can fetch it. This can
 	// happen when running the test on a slow machine.
 	var expiration time.Time
-	timeout := time.Second
+	timeout := 2 * time.Second
 	for {
-		expiration = time.Now().Add(time.Second)
+		expiration = time.Now().Add(2 * time.Second)
 		err := rp.Publish(ctx, publisher.privKey, p, namesys.PublishWithEOL(expiration))
 		require.NoError(t, err)