Repository: ipfs/boxo

Issue:
bitswap/client: sessions should close via a close method, not a context
Bitswap sessions should close using a `Close` method, and not a context. A context should be used for request cancelation, not for session lifetime.

Discussion:

Solution patch:
diff --git a/bitswap/client/client.go b/bitswap/client/client.go
index e77e47de..306e74b6 100644
--- a/bitswap/client/client.go
+++ b/bitswap/client/client.go
@@ -300,7 +300,6 @@ func New(parent context.Context, network bsnet.BitSwapNetwork, providerFinder ro
 	}
 
 	sessionFactory := func(
-		sessctx context.Context,
 		sessmgr bssession.SessionManager,
 		id uint64,
 		spm bssession.SessionPeerManager,
@@ -311,6 +310,7 @@ func New(parent context.Context, network bsnet.BitSwapNetwork, providerFinder ro
 		provSearchDelay time.Duration,
 		rebroadcastDelay delay.D,
 		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
@@ -322,13 +322,13 @@ func New(parent context.Context, network bsnet.BitSwapNetwork, providerFinder ro
 		} else if providerFinder != nil {
 			sessionProvFinder = providerFinder
 		}
-		return bssession.New(sessctx, sessmgr, id, spm, sessionProvFinder, sim, pm, bpm, notif, provSearchDelay, rebroadcastDelay, self, bs.havesReceivedGauge)
+		return bssession.New(sessmgr, id, spm, sessionProvFinder, sim, pm, bpm, notif, provSearchDelay, rebroadcastDelay, self, bs.havesReceivedGauge, retrievalState)
 	}
-	sessionPeerManagerFactory := func(ctx context.Context, id uint64) bssession.SessionPeerManager {
+	sessionPeerManagerFactory := func(id uint64) bssession.SessionPeerManager {
 		return bsspm.New(id, network)
 	}
 	notif := notifications.New(bs.traceBlock)
-	sm = bssm.New(ctx, sessionFactory, sim, sessionPeerManagerFactory, bpm, pm, notif, network.Self())
+	sm = bssm.New(sessionFactory, sim, sessionPeerManagerFactory, bpm, pm, notif, network.Self())
 
 	bs.sm = sm
 	bs.notif = notif
@@ -439,22 +439,19 @@ func (bs *Client) GetBlock(ctx context.Context, k cid.Cid) (blocks.Block, error)
 //
 // When the provided context is canceled, GetBlocks will stop returning blocks
 // and the returned channel will be closed.
+//
+// This method creates a temporary internal session that is automatically cleaned up
+// when the operation completes. Any retrieval.State attached to the context will be
+// preserved and used for diagnostic tracking throughout the retrieval process.
 func (bs *Client) GetBlocks(ctx context.Context, keys []cid.Cid) (<-chan blocks.Block, error) {
 	ctx, span := internal.StartSpan(ctx, "GetBlocks", trace.WithAttributes(attribute.Int("NumKeys", len(keys))))
 	defer span.End()
 
-	// 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)
+	session := bs.sm.NewSession(bs.provSearchDelay, bs.rebroadcastDelay, retrieval.StateFromContext(ctx))
 
 	blocksChan, err := session.GetBlocks(ctx, keys)
 	if err != nil {
-		cancelSession()
+		session.Close()
 		return nil, err
 	}
 
@@ -462,7 +459,7 @@ func (bs *Client) GetBlocks(ctx context.Context, keys []cid.Cid) (<-chan blocks.
 	go func() {
 		defer func() {
 			close(out)
-			cancelSession()
+			session.Close()
 		}()
 
 		ctxDone := ctx.Done()
@@ -709,8 +706,21 @@ func (bs *Client) IsOnline() bool {
 // method, but the session will use the fact that the requests are related to
 // be more efficient in its requests to peers. If you are using a session
 // from blockservice, it will create a bitswap session automatically.
+//
+// The context is used for:
+//   - Automatic session cleanup: when the context is canceled, the session will be closed
+//   - Retrieval state tracking: any retrieval.State attached to the context will be
+//     preserved in the session for diagnostic tracking of provider discovery and data retrieval
+//
+// The session's lifetime is controlled by the context and requires that the
+// context be canceled for the session to be closed.
 func (bs *Client) NewSession(ctx context.Context) exchange.Fetcher {
 	ctx, span := internal.StartSpan(ctx, "NewSession")
 	defer span.End()
-	return bs.sm.NewSession(ctx, bs.provSearchDelay, bs.rebroadcastDelay)
+
+	session := bs.sm.NewSession(bs.provSearchDelay, bs.rebroadcastDelay, retrieval.StateFromContext(ctx))
+	context.AfterFunc(ctx, func() {
+		session.Close()
+	})
+	return session
 }
diff --git a/bitswap/client/internal/peermanager/peermanager.go b/bitswap/client/internal/peermanager/peermanager.go
index bfe83fc3..940d6f60 100644
--- a/bitswap/client/internal/peermanager/peermanager.go
+++ b/bitswap/client/internal/peermanager/peermanager.go
@@ -140,7 +140,7 @@ func (pm *PeerManager) ResponseReceived(p peer.ID, ks []cid.Cid) {
 // to discover seeds).
 // For each peer it filters out want-haves that have previously been sent to
 // the peer.
-func (pm *PeerManager) BroadcastWantHaves(ctx context.Context, wantHaves []cid.Cid) {
+func (pm *PeerManager) BroadcastWantHaves(wantHaves []cid.Cid) {
 	pm.pqLk.Lock()
 	defer pm.pqLk.Unlock()
 
@@ -149,7 +149,7 @@ func (pm *PeerManager) BroadcastWantHaves(ctx context.Context, wantHaves []cid.C
 
 // SendWants sends the given want-blocks and want-haves to the given peer.
 // It filters out wants that have previously been sent to the peer.
-func (pm *PeerManager) SendWants(ctx context.Context, p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) bool {
+func (pm *PeerManager) SendWants(p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) bool {
 	pm.pqLk.Lock()
 	defer pm.pqLk.Unlock()
 
@@ -162,7 +162,7 @@ func (pm *PeerManager) SendWants(ctx context.Context, p peer.ID, wantBlocks []ci
 
 // SendCancels sends cancels for the given keys to all peers who had previously
 // received a want for those keys.
-func (pm *PeerManager) SendCancels(ctx context.Context, cancelKs []cid.Cid) {
+func (pm *PeerManager) SendCancels(cancelKs []cid.Cid) {
 	pm.pqLk.Lock()
 	defer pm.pqLk.Unlock()
 
diff --git a/bitswap/client/internal/session/session.go b/bitswap/client/internal/session/session.go
index 878111ec..ce17ba14 100644
--- a/bitswap/client/internal/session/session.go
+++ b/bitswap/client/internal/session/session.go
@@ -10,6 +10,7 @@ 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"
 	delay "github.com/ipfs/go-ipfs-delay"
@@ -25,51 +26,54 @@ const (
 	broadcastLiveWantsLimit = 64
 )
 
-// PeerManager keeps track of which sessions are interested in which peers
-// and takes care of sending wants for the sessions
+// PeerManager keeps track of which sessions are interested in which peers and
+// takes care of sending wants for the sessions.
 type PeerManager interface {
-	// RegisterSession tells the PeerManager that the session is interested
-	// in a peer's connection state
+	// RegisterSession tells the PeerManager that the session is interested in
+	// a peer's connection state.
 	RegisterSession(peer.ID, bspm.Session)
 	// UnregisterSession tells the PeerManager that the session is no longer
-	// interested in a peer's connection state
+	// interested in a peer's connection state.
 	UnregisterSession(uint64)
-	// SendWants tells the PeerManager to send wants to the given peer
-	SendWants(ctx context.Context, peerId peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) bool
+	// SendWants tells the PeerManager to send wants to the given peer.
+	SendWants(peerId peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) bool
 	// BroadcastWantHaves sends want-haves to all connected peers (used for
-	// session discovery)
-	BroadcastWantHaves(context.Context, []cid.Cid)
-	// SendCancels tells the PeerManager to send cancels to all peers
-	SendCancels(context.Context, []cid.Cid)
+	// session discovery).
+	BroadcastWantHaves([]cid.Cid)
+	// SendCancels tells the PeerManager to send cancels to all peers.
+	SendCancels([]cid.Cid)
 }
 
-// SessionManager manages all the sessions
+// SessionManager manages all the sessions,
 type SessionManager interface {
-	// Remove a session (called when the session shuts down)
+	// RemoveSession removes a session (called when the session shuts down).
 	RemoveSession(sesid uint64)
-	// Cancel wants (called when a call to GetBlocks() is cancelled)
+	// CancelSessionWants cancels the specified session's wants (called when a
+	// call to GetBlocks() is cancelled).
 	CancelSessionWants(sid uint64, wants []cid.Cid)
 }
 
 // SessionPeerManager keeps track of peers in the session
 type SessionPeerManager interface {
-	// PeersDiscovered indicates if any peers have been discovered yet
+	// PeersDiscovered indicates if any peers have been discovered yet,
 	PeersDiscovered() bool
-	// Shutdown the SessionPeerManager
+	// Shutdown stops the SessionPeerManager.
 	Shutdown()
-	// Adds a peer to the session, returning true if the peer is new
+	// AddPeer adds a peer to the session, returning true if the peer is new.
 	AddPeer(peer.ID) bool
-	// Removes a peer from the session, returning true if the peer existed
+	// RemovePeer removes a peer from the session, returning true if the peer
+	// existed.
 	RemovePeer(peer.ID) bool
-	// All peers in the session
+	// Peers returns all peers in the session.
 	Peers() []peer.ID
-	// Whether there are any peers in the session
+	// HasPeers returns true if there are any peers in the session.
 	HasPeers() bool
-	// Protect connection from being pruned by the connection manager
+	// ProtectConnection prevents a connection from being pruned by the
+	// connection manager.
 	ProtectConnection(peer.ID)
 }
 
-// opType is the kind of operation that is being processed by the event loop
+// opType is the kind of operation that is being processed by the event loop.
 type opType int
 
 const (
@@ -93,10 +97,15 @@ type op struct {
 // Session holds state for an individual bitswap transfer operation.
 // This allows bitswap to make smarter decisions about who to send wantlist
 // info to, and who to request blocks from.
+//
+// Sessions manage their own lifecycle independently of any request contexts.
+// Each session maintains an internal context for its operations and must be
+// explicitly closed via the Close() method when no longer needed. The session's
+// internal context is used for provider discovery and other long-running operations.
 type Session struct {
 	// dependencies
 	ctx            context.Context
-	shutdown       func()
+	cancel         context.CancelFunc
 	sm             SessionManager
 	pm             PeerManager
 	sprm           SessionPeerManager
@@ -126,10 +135,20 @@ type Session struct {
 	self peer.ID
 }
 
-// New creates a new bitswap session whose lifetime is bounded by the
-// given context.
+// New creates a new bitswap session.
+//
+// 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 retrievalState parameter, if provided, enables diagnostic tracking of
+// the retrieval process. It is attached to the session's internal context and
+// used to track provider discovery, connection attempts, and data retrieval
+// phases. This is particularly useful for debugging timeout errors and
+// understanding retrieval performance.
 func New(
-	ctx context.Context,
 	sm SessionManager,
 	id uint64,
 	sprm SessionPeerManager,
@@ -142,13 +161,20 @@ func New(
 	periodicSearchDelay delay.D,
 	self peer.ID,
 	havesReceivedGauge bspm.Gauge,
+	retrievalState *retrieval.State,
 ) *Session {
-	ctx, cancel := context.WithCancel(ctx)
+	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)
+	}
+
 	s := &Session{
 		sw:                  newSessionWants(broadcastLiveWantsLimit),
 		tickDelayReqs:       make(chan time.Duration),
 		ctx:                 ctx,
-		shutdown:            cancel,
+		cancel:              cancel,
 		sm:                  sm,
 		pm:                  pm,
 		sprm:                sprm,
@@ -174,8 +200,16 @@ func (s *Session) ID() uint64 {
 	return s.id
 }
 
-func (s *Session) Shutdown() {
-	s.shutdown()
+// 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.
+func (s *Session) Close() {
+	s.cancel()
 }
 
 // ReceiveFrom receives incoming blocks from the given peer.
@@ -189,14 +223,14 @@ func (s *Session) ReceiveFrom(from peer.ID, ks []cid.Cid, haves []cid.Cid, dontH
 	dontHaves = interestedRes[2]
 	s.logReceiveFrom(from, ks, haves, dontHaves)
 
-	// Inform the session want sender that a message has been received
+	// Inform the session want sender that a message has been received.
 	s.sws.Update(from, ks, haves, dontHaves)
 
 	if len(ks) == 0 {
 		return
 	}
 
-	// Inform the session that blocks have been received
+	// Inform the session that blocks have been received.
 	select {
 	case s.incoming <- op{op: opReceive, keys: ks}:
 	case <-s.ctx.Done():
@@ -204,7 +238,6 @@ func (s *Session) ReceiveFrom(from peer.ID, ks []cid.Cid, haves []cid.Cid, dontH
 }
 
 func (s *Session) logReceiveFrom(from peer.ID, interestedKs []cid.Cid, haves []cid.Cid, dontHaves []cid.Cid) {
-	// Save some CPU cycles if log level is higher than debug
 	if !log.LevelEnabled(logging.LevelDebug) {
 		return
 	}
@@ -271,9 +304,10 @@ func (s *Session) onPeersExhausted(ks []cid.Cid) {
 	s.nonBlockingEnqueue(op{op: opBroadcast, keys: ks})
 }
 
-// We don't want to block the sessionWantSender if the incoming channel
-// is full. So if we can't immediately send on the incoming channel spin
-// it off into a go-routine.
+// nonBlockingEnqueue enqueues an operation on the incoming queue without
+// blocking. We do not want to block the sessionWantSender if the incoming
+// channel is full. So if we cannot immediately send on the incoming channel
+// spin it off into a goroutine.
 func (s *Session) nonBlockingEnqueue(o op) {
 	select {
 	case s.incoming <- o:
@@ -287,8 +321,8 @@ func (s *Session) nonBlockingEnqueue(o op) {
 	}
 }
 
-// Session run loop -- everything in this function should not be called
-// outside of this loop
+// run is the session run loop. Everything in this function should not be
+// called outside of this loop.
 func (s *Session) run(ctx context.Context) {
 	go s.sws.Run()
 
@@ -345,30 +379,31 @@ func (s *Session) run(ctx context.Context) {
 	}
 }
 
-// Called when the session hasn't received any blocks for some time, or when
-// all peers in the session have sent DONT_HAVE for a particular set of CIDs.
-// Send want-haves to all connected peers, and search for new peers with the CID.
+// broadcast is called when the session has not received any blocks for some
+// time, or when all peers in the session have sent DONT_HAVE for a particular
+// set of CIDs. Send want-haves to all connected peers, and search for new
+// peers with the CID.
 func (s *Session) broadcast(ctx context.Context, wants []cid.Cid) {
-	// If this broadcast is because of an idle timeout (we haven't received
-	// any blocks for a while) then broadcast all pending wants
+	// If this broadcast is because of an idle timeout (we haven't received any
+	// blocks for a while) then broadcast all pending wants.
 	if wants == nil {
 		wants = s.sw.PrepareBroadcast()
 	}
 
-	// Broadcast a want-have for the live wants to everyone we're connected to
+	// Broadcast a want-have for the live wants to connected peers.
 	s.broadcastWantHaves(ctx, wants)
 
-	// do not find providers on consecutive ticks
-	// -- just rely on periodic search widening
+	// Do not find providers on consecutive ticks. Instead rely on periodic
+	// search widening.
 	if len(wants) > 0 && (s.consecutiveTicks == 0) {
-		// Search for providers who have the first want in the list.
-		// Typically if the provider has the first block they will have
-		// the rest of the blocks also.
+		// Search for providers who have the first want in the list. Typically
+		// if the provider has the first block they will have the rest of the
+		// blocks also.
 		log.Debugw("FindMorePeers", "session", s.id, "cid", wants[0], "pending", len(wants))
 		s.findMorePeers(ctx, wants[0])
 	}
 
-	// If we have live wants record a consecutive tick
+	// If there are live wants record a consecutive tick.
 	if s.sw.HasLiveWants() {
 		s.resetIdleTick() // call before incrementing s.consecutiveTicks
 		s.consecutiveTicks++
@@ -393,7 +428,7 @@ func (s *Session) handlePeriodicSearch(ctx context.Context) {
 }
 
 // findMorePeers attempts to find more peers for a session by searching for
-// providers for the given Cid
+// providers for the given CID.
 func (s *Session) findMorePeers(ctx context.Context, c cid.Cid) {
 	// noop when provider finder is disabled
 	if s.providerFinder == nil {
@@ -403,8 +438,8 @@ func (s *Session) findMorePeers(ctx context.Context, c cid.Cid) {
 		ctx, span := internal.StartSpan(ctx, "Session.FindMorePeers")
 		defer span.End()
 		for p := range s.providerFinder.FindProvidersAsync(ctx, k, 0) {
-			// When a provider indicates that it has a cid, it's equivalent to
-			// the providing peer sending a HAVE
+			// When a provider indicates that it has a cid, it is equivalent to
+			// the providing peer sending a HAVE.
 			log.Infow("Found peer for CID", "peer", p, "cid", k)
 			span.AddEvent("FoundPeer")
 			s.sws.Update(p.ID, nil, []cid.Cid{c}, nil)
@@ -412,7 +447,7 @@ func (s *Session) findMorePeers(ctx context.Context, c cid.Cid) {
 	}(c)
 }
 
-// handleShutdown is called when the session shuts down
+// handleShutdown is called when the session shuts down.
 func (s *Session) handleShutdown() {
 	// Stop the idle timer
 	s.idleTick.Stop()
@@ -426,10 +461,10 @@ func (s *Session) handleShutdown() {
 	s.sm.RemoveSession(s.id)
 }
 
-// handleReceive is called when the session receives blocks from a peer
+// handleReceive is called when the session receives blocks from a peer.
 func (s *Session) handleReceive(ks []cid.Cid) {
 	// Record which blocks have been received and figure out the total latency
-	// for fetching the blocks
+	// for fetching the blocks.
 	wanted, totalLatency := s.sw.BlocksReceived(ks)
 	if len(wanted) == 0 {
 		return
@@ -439,13 +474,13 @@ func (s *Session) handleReceive(ks []cid.Cid) {
 	s.latencyTrkr.receiveUpdate(len(wanted), totalLatency)
 
 	// Inform the SessionManager that this session is no longer expecting to
-	// receive the wanted keys, since we now have them
+	// receive the wanted keys, since we now have them,
 	s.sm.CancelSessionWants(s.id, wanted)
 
 	s.idleTick.Stop()
 
-	// We've received new wanted blocks, so reset the number of ticks
-	// that have occurred since the last new block
+	// We have received new wanted blocks, so reset the number of ticks that
+	// have occurred since the last new block,
 	s.consecutiveTicks = 0
 
 	// Reset rebroadcast timer if there are still outstanding wants.
@@ -479,12 +514,14 @@ func (s *Session) wantBlocks(ctx context.Context, newks []cid.Cid) {
 	}
 }
 
-// Send want-haves to all connected peers
+// broadcastWantHaves sends want-haves to all connected peers.
 func (s *Session) broadcastWantHaves(ctx context.Context, wants []cid.Cid) {
 	log.Debugw("broadcastWantHaves", "session", s.id, "cids", wants)
-	s.pm.BroadcastWantHaves(ctx, wants)
+	s.pm.BroadcastWantHaves(wants)
 }
 
+// resetIdleTick sets the idle tick time based on average latency.
+//
 // The session will broadcast if it has outstanding wants and doesn't receive
 // any blocks for some time.
 // The length of time is calculated
diff --git a/bitswap/client/internal/session/sessionwantsender.go b/bitswap/client/internal/session/sessionwantsender.go
index 3ad5967d..63544522 100644
--- a/bitswap/client/internal/session/sessionwantsender.go
+++ b/bitswap/client/internal/session/sessionwantsender.go
@@ -579,7 +579,7 @@ func (sws *sessionWantSender) sendWants(sends allWants) {
 		// precedence over want-haves.
 		wblks := snd.wantBlocks.Keys()
 		whaves := snd.wantHaves.Keys()
-		snd.sent = sws.pm.SendWants(sws.ctx, p, wblks, whaves)
+		snd.sent = sws.pm.SendWants(p, wblks, whaves)
 		if !snd.sent {
 			// Do not update state if the wants not sent.
 			continue
diff --git a/bitswap/client/internal/sessionmanager/sessionmanager.go b/bitswap/client/internal/sessionmanager/sessionmanager.go
index 0cfa7291..452e1616 100644
--- a/bitswap/client/internal/sessionmanager/sessionmanager.go
+++ b/bitswap/client/internal/sessionmanager/sessionmanager.go
@@ -12,6 +12,7 @@ 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"
 	delay "github.com/ipfs/go-ipfs-delay"
 	peer "github.com/libp2p/go-libp2p/core/peer"
@@ -19,17 +20,18 @@ import (
 	"go.opentelemetry.io/otel/trace"
 )
 
-// Session is a session that is managed by the session manager
+// 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
 	ReceiveFrom(peer.ID, []cid.Cid, []cid.Cid, []cid.Cid)
-	Shutdown()
+	Close()
 }
 
 // SessionFactory generates a new session for the SessionManager to track.
 type SessionFactory func(
-	ctx context.Context,
 	sm bssession.SessionManager,
 	id uint64,
 	sprm bssession.SessionPeerManager,
@@ -39,15 +41,15 @@ type SessionFactory func(
 	notif notifications.PubSub,
 	provSearchDelay time.Duration,
 	rebroadcastDelay delay.D,
-	self peer.ID) Session
+	self peer.ID,
+	retrievalState *retrieval.State) Session
 
 // PeerManagerFactory generates a new peer manager for a session.
-type PeerManagerFactory func(ctx context.Context, id uint64) bssession.SessionPeerManager
+type PeerManagerFactory func(id uint64) bssession.SessionPeerManager
 
 // SessionManager is responsible for creating, managing, and dispatching to
 // sessions.
 type SessionManager struct {
-	ctx                    context.Context
 	sessionFactory         SessionFactory
 	sessionInterestManager *bssim.SessionInterestManager
 	peerManagerFactory     PeerManagerFactory
@@ -67,11 +69,10 @@ type SessionManager struct {
 }
 
 // New creates a new SessionManager.
-func New(ctx context.Context, sessionFactory SessionFactory, sessionInterestManager *bssim.SessionInterestManager, peerManagerFactory PeerManagerFactory,
+func New(sessionFactory SessionFactory, sessionInterestManager *bssim.SessionInterestManager, peerManagerFactory PeerManagerFactory,
 	blockPresenceManager *bsbpm.BlockPresenceManager, peerManager bssession.PeerManager, notif notifications.PubSub, self peer.ID,
 ) *SessionManager {
 	return &SessionManager{
-		ctx:                    ctx,
 		sessionFactory:         sessionFactory,
 		sessionInterestManager: sessionInterestManager,
 		peerManagerFactory:     peerManagerFactory,
@@ -83,19 +84,25 @@ func New(ctx context.Context, sessionFactory SessionFactory, sessionInterestMana
 	}
 }
 
-// NewSession initializes a session with the given context, and adds to the
-// session manager.
-func (sm *SessionManager) NewSession(ctx context.Context,
-	provSearchDelay time.Duration,
-	rebroadcastDelay delay.D,
-) exchange.Fetcher {
+// 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 time.Duration, rebroadcastDelay delay.D, retrievalState *retrieval.State) Session {
 	id := sm.GetNextSessionID()
 
-	ctx, span := internal.StartSpan(ctx, "SessionManager.NewSession", trace.WithAttributes(attribute.String("ID", strconv.FormatUint(id, 10))))
+	_, span := internal.StartSpan(context.Background(), "SessionManager.NewSession", trace.WithAttributes(attribute.String("ID", strconv.FormatUint(id, 10))))
 	defer span.End()
 
-	pm := sm.peerManagerFactory(ctx, id)
-	session := sm.sessionFactory(ctx, sm, id, pm, sm.sessionInterestManager, sm.peerManager, sm.blockPresenceManager, sm.notif, provSearchDelay, rebroadcastDelay, sm.self)
+	pm := sm.peerManagerFactory(id)
+	session := sm.sessionFactory(sm, id, pm, sm.sessionInterestManager, sm.peerManager, sm.blockPresenceManager, sm.notif, provSearchDelay, rebroadcastDelay, sm.self, retrievalState)
 
 	sm.sessLk.Lock()
 	if sm.sessions != nil { // check if SessionManager was shutdown
@@ -109,11 +116,7 @@ func (sm *SessionManager) NewSession(ctx context.Context,
 func (sm *SessionManager) Shutdown() {
 	sm.sessLk.Lock()
 
-	sessions := make([]Session, 0, len(sm.sessions))
-	for _, ses := range sm.sessions {
-		sessions = append(sessions, ses)
-	}
-
+	sessions := sm.sessions
 	// Ensure that if Shutdown() is called twice we only shut down
 	// the sessions once
 	sm.sessions = nil
@@ -121,7 +124,7 @@ func (sm *SessionManager) Shutdown() {
 	sm.sessLk.Unlock()
 
 	for _, ses := range sessions {
-		ses.Shutdown()
+		ses.Close()
 	}
 }
 
@@ -198,5 +201,5 @@ func (sm *SessionManager) cancelWants(wants []cid.Cid) {
 	// Send CANCEL to all peers for blocks that no session is interested in
 	// anymore.
 	// Note: use bitswap context because session context may already be Done.
-	sm.peerManager.SendCancels(sm.ctx, wants)
+	sm.peerManager.SendCancels(wants)
 }



Test patch:
diff --git a/bitswap/client/internal/peermanager/peermanager_test.go b/bitswap/client/internal/peermanager/peermanager_test.go
index e85a3ad6..d35942e7 100644
--- a/bitswap/client/internal/peermanager/peermanager_test.go
+++ b/bitswap/client/internal/peermanager/peermanager_test.go
@@ -136,7 +136,7 @@ func TestBroadcastOnConnect(t *testing.T) {
 	peerManager := New(ctx, peerQueueFactory, bcastAlways)
 
 	cids := random.Cids(2)
-	peerManager.BroadcastWantHaves(ctx, cids)
+	peerManager.BroadcastWantHaves(cids)
 
 	// Connect with two broadcast wants for first peer
 	peerManager.Connected(peer1)
@@ -159,7 +159,7 @@ func TestBroadcastWantHaves(t *testing.T) {
 	cids := random.Cids(3)
 
 	// Broadcast the first two.
-	peerManager.BroadcastWantHaves(ctx, cids[:2])
+	peerManager.BroadcastWantHaves(cids[:2])
 
 	// First peer should get them.
 	peerManager.Connected(peer1)
@@ -174,7 +174,7 @@ func TestBroadcastWantHaves(t *testing.T) {
 
 	// Send a broadcast to all peers, including cid that was already sent to
 	// first peer
-	peerManager.BroadcastWantHaves(ctx, []cid.Cid{cids[0], cids[2]})
+	peerManager.BroadcastWantHaves([]cid.Cid{cids[0], cids[2]})
 	collected = collectMessages(msgs, 2*time.Millisecond)
 
 	// One of the want-haves was already sent to peer1
@@ -199,7 +199,7 @@ func TestSendWants(t *testing.T) {
 	cids := random.Cids(4)
 
 	peerManager.Connected(peer1)
-	peerManager.SendWants(ctx, peer1, []cid.Cid{cids[0]}, []cid.Cid{cids[2]})
+	peerManager.SendWants(peer1, []cid.Cid{cids[0]}, []cid.Cid{cids[2]})
 	collected := collectMessages(msgs, 2*time.Millisecond)
 
 	if len(collected[peer1].wantHaves) != 1 {
@@ -209,7 +209,7 @@ func TestSendWants(t *testing.T) {
 		t.Fatal("Expected want-block to be sent to peer")
 	}
 
-	peerManager.SendWants(ctx, peer1, []cid.Cid{cids[0], cids[1]}, []cid.Cid{cids[2], cids[3]})
+	peerManager.SendWants(peer1, []cid.Cid{cids[0], cids[1]}, []cid.Cid{cids[2], cids[3]})
 	collected = collectMessages(msgs, 2*time.Millisecond)
 
 	// First want-have and want-block should be filtered (because they were
@@ -237,13 +237,13 @@ func TestSendCancels(t *testing.T) {
 	peerManager.Connected(peer2)
 
 	// Send 2 want-blocks and 1 want-have to peer1
-	peerManager.SendWants(ctx, peer1, []cid.Cid{cids[0], cids[1]}, []cid.Cid{cids[2]})
+	peerManager.SendWants(peer1, []cid.Cid{cids[0], cids[1]}, []cid.Cid{cids[2]})
 
 	// Clear messages
 	collectMessages(msgs, 2*time.Millisecond)
 
 	// Send cancels for 1 want-block and 1 want-have
-	peerManager.SendCancels(ctx, []cid.Cid{cids[0], cids[2]})
+	peerManager.SendCancels([]cid.Cid{cids[0], cids[2]})
 	collected := collectMessages(msgs, 2*time.Millisecond)
 
 	if _, ok := collected[peer2]; ok {
@@ -254,7 +254,7 @@ func TestSendCancels(t *testing.T) {
 	}
 
 	// Send cancels for all cids
-	peerManager.SendCancels(ctx, cids)
+	peerManager.SendCancels(cids)
 	collected = collectMessages(msgs, 2*time.Millisecond)
 
 	if _, ok := collected[peer2]; ok {
@@ -370,17 +370,17 @@ func BenchmarkPeerManager(b *testing.B) {
 		r := rand.Intn(8)
 		if r == 0 {
 			wants := random.Cids(10)
-			peerManager.SendWants(ctx, peers[i], wants[:2], wants[2:])
+			peerManager.SendWants(peers[i], wants[:2], wants[2:])
 			wanted = append(wanted, wants...)
 		} else if r == 1 {
 			wants := random.Cids(30)
-			peerManager.BroadcastWantHaves(ctx, wants)
+			peerManager.BroadcastWantHaves(wants)
 			wanted = append(wanted, wants...)
 		} else {
 			limit := len(wanted) / 10
 			cancel := wanted[:limit]
 			wanted = wanted[limit:]
-			peerManager.SendCancels(ctx, cancel)
+			peerManager.SendCancels(cancel)
 		}
 	}
 }
diff --git a/bitswap/client/internal/session/session_test.go b/bitswap/client/internal/session/session_test.go
index 49f2c08a..61d57bb0 100644
--- a/bitswap/client/internal/session/session_test.go
+++ b/bitswap/client/internal/session/session_test.go
@@ -146,20 +146,16 @@ func newFakePeerManager() *fakePeerManager {
 
 func (pm *fakePeerManager) RegisterSession(peer.ID, bspm.Session) {}
 func (pm *fakePeerManager) UnregisterSession(uint64)              {}
-func (pm *fakePeerManager) SendWants(context.Context, peer.ID, []cid.Cid, []cid.Cid) bool {
+func (pm *fakePeerManager) SendWants(peer.ID, []cid.Cid, []cid.Cid) bool {
 	return true
 }
 
-func (pm *fakePeerManager) BroadcastWantHaves(ctx context.Context, cids []cid.Cid) {
-	select {
-	case pm.wantReqs <- wantReq{cids}:
-	case <-ctx.Done():
-	}
+func (pm *fakePeerManager) BroadcastWantHaves(cids []cid.Cid) {
+	pm.wantReqs <- wantReq{cids}
 }
-func (pm *fakePeerManager) SendCancels(ctx context.Context, cancels []cid.Cid) {}
+func (pm *fakePeerManager) SendCancels(cancels []cid.Cid) {}
 
 func TestSessionGetBlocks(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -169,14 +165,15 @@ func TestSessionGetBlocks(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil)
+	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil, nil)
+	defer session.Close()
 	blks := random.BlocksOfSize(broadcastLiveWantsLimit*2, blockSize)
 	var cids []cid.Cid
 	for _, block := range blks {
 		cids = append(cids, block.Cid())
 	}
 
-	_, err := session.GetBlocks(ctx, cids)
+	_, err := session.GetBlocks(context.Background(), cids)
 	require.NoError(t, err, "error getting blocks")
 
 	// Wait for initial want request
@@ -229,7 +226,7 @@ func TestSessionGetBlocks(t *testing.T) {
 	require.Len(t, wanted, len(blks)-1, "session wants incorrect number of blocks")
 
 	// Shut down session
-	cancel()
+	session.Close()
 
 	time.Sleep(10 * time.Millisecond)
 
@@ -238,8 +235,6 @@ func TestSessionGetBlocks(t *testing.T) {
 }
 
 func TestSessionFindMorePeers(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 900*time.Millisecond)
-	defer cancel()
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -249,13 +244,18 @@ func TestSessionFindMorePeers(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil)
+	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil, nil)
+	defer session.Close()
 	session.SetBaseTickDelay(200 * time.Microsecond)
 	blks := random.BlocksOfSize(broadcastLiveWantsLimit*2, blockSize)
 	var cids []cid.Cid
 	for _, block := range blks {
 		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")
 
@@ -307,8 +307,6 @@ func TestSessionFindMorePeers(t *testing.T) {
 }
 
 func TestSessionOnPeersExhausted(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
-	defer cancel()
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -319,12 +317,17 @@ func TestSessionOnPeersExhausted(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil)
+	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil, nil)
+	defer session.Close()
 	blks := random.BlocksOfSize(broadcastLiveWantsLimit+5, blockSize)
 	var cids []cid.Cid
 	for _, block := range blks {
 		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")
 
@@ -345,8 +348,6 @@ 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,12 +357,17 @@ func TestSessionFailingToGetFirstBlock(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, 10*time.Millisecond, delay.Fixed(100*time.Millisecond), "", nil)
+	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, 10*time.Millisecond, delay.Fixed(100*time.Millisecond), "", nil, nil)
+	defer session.Close()
 	blks := random.BlocksOfSize(4, blockSize)
 	var cids []cid.Cid
 	for _, block := range blks {
 		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")
 
@@ -456,8 +462,8 @@ func TestSessionCtxCancelClosesGetBlocksChannel(t *testing.T) {
 	sm := newMockSessionMgr(sim)
 
 	// Create a new session with its own context
-	sessctx, sesscancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
-	session := New(sessctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil)
+	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil, nil)
+	defer session.Close()
 
 	timerCtx, timerCancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
 	defer timerCancel()
@@ -470,8 +476,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 context
-	sesscancel()
+	// Cancel the session
+	session.Close()
 
 	// Expect the GetBlocks() channel to be closed
 	select {
@@ -498,13 +504,11 @@ func TestSessionOnShutdownCalled(t *testing.T) {
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
 
-	// Create a new session with its own context
-	sessctx, sesscancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
-	defer sesscancel()
-	session := New(sessctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil)
+	// Create a new session
+	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil, nil)
 
 	// Shutdown the session
-	session.Shutdown()
+	session.Close()
 
 	time.Sleep(10 * time.Millisecond)
 
@@ -512,8 +516,7 @@ func TestSessionOnShutdownCalled(t *testing.T) {
 	require.True(t, sm.removeSessionCalled(), "expected onShutdown to be called")
 }
 
-func TestSessionReceiveMessageAfterCtxCancel(t *testing.T) {
-	ctx, cancelCtx := context.WithTimeout(context.Background(), 20*time.Millisecond)
+func TestSessionReceiveMessageAfterClose(t *testing.T) {
 	fpm := newFakePeerManager()
 	fspm := newFakeSessionPeerManager()
 	fpf := newFakeProviderFinder()
@@ -524,18 +527,20 @@ func TestSessionReceiveMessageAfterCtxCancel(t *testing.T) {
 	defer notif.Shutdown()
 	id := random.SequenceNext()
 	sm := newMockSessionMgr(sim)
-	session := New(ctx, sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil)
+	session := New(sm, id, fspm, fpf, sim, fpm, bpm, notif, time.Second, delay.Fixed(time.Minute), "", nil, nil)
+	defer session.Close()
+
 	blks := random.BlocksOfSize(2, blockSize)
 	cids := []cid.Cid{blks[0].Cid(), blks[1].Cid()}
 
-	_, err := session.GetBlocks(ctx, cids)
+	_, err := session.GetBlocks(context.Background(), cids)
 	require.NoError(t, err, "error getting blocks")
 
 	// Wait for initial want request
 	<-fpm.wantReqs
 
 	// Shut down session
-	cancelCtx()
+	session.Close()
 
 	// Simulate receiving block for a CID
 	peer := random.Peers(1)[0]
diff --git a/bitswap/client/internal/session/sessionwantsender_test.go b/bitswap/client/internal/session/sessionwantsender_test.go
index 2a786d39..73cc1c2e 100644
--- a/bitswap/client/internal/session/sessionwantsender_test.go
+++ b/bitswap/client/internal/session/sessionwantsender_test.go
@@ -1,7 +1,6 @@
 package session
 
 import (
-	"context"
 	"sync"
 	"testing"
 	"time"
@@ -79,11 +78,11 @@ func (pm *mockPeerManager) has(p peer.ID, sid uint64) bool {
 	return false
 }
 
-func (*mockPeerManager) UnregisterSession(uint64)                      {}
-func (*mockPeerManager) BroadcastWantHaves(context.Context, []cid.Cid) {}
-func (*mockPeerManager) SendCancels(context.Context, []cid.Cid)        {}
+func (*mockPeerManager) UnregisterSession(uint64)     {}
+func (*mockPeerManager) BroadcastWantHaves([]cid.Cid) {}
+func (*mockPeerManager) SendCancels([]cid.Cid)        {}
 
-func (pm *mockPeerManager) SendWants(ctx context.Context, p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) bool {
+func (pm *mockPeerManager) SendWants(p peer.ID, wantBlocks []cid.Cid, wantHaves []cid.Cid) bool {
 	pm.lk.Lock()
 	defer pm.lk.Unlock()
 
diff --git a/bitswap/client/internal/sessionmanager/sessionmanager_test.go b/bitswap/client/internal/sessionmanager/sessionmanager_test.go
index 7c66c0db..2f968534 100644
--- a/bitswap/client/internal/sessionmanager/sessionmanager_test.go
+++ b/bitswap/client/internal/sessionmanager/sessionmanager_test.go
@@ -12,6 +12,7 @@ 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"
 	delay "github.com/ipfs/go-ipfs-delay"
@@ -48,7 +49,7 @@ func (fs *fakeSession) ReceiveFrom(p peer.ID, ks []cid.Cid, wantBlocks []cid.Cid
 	fs.sm.CancelSessionWants(fs.id, ks)
 }
 
-func (fs *fakeSession) Shutdown() {
+func (fs *fakeSession) Close() {
 	fs.sm.RemoveSession(fs.id)
 }
 
@@ -67,11 +68,11 @@ type fakePeerManager struct {
 	cancels []cid.Cid
 }
 
-func (*fakePeerManager) RegisterSession(peer.ID, bspm.Session)                         {}
-func (*fakePeerManager) UnregisterSession(uint64)                                      {}
-func (*fakePeerManager) SendWants(context.Context, peer.ID, []cid.Cid, []cid.Cid) bool { return true }
-func (*fakePeerManager) BroadcastWantHaves(context.Context, []cid.Cid)                 {}
-func (fpm *fakePeerManager) SendCancels(ctx context.Context, cancels []cid.Cid) {
+func (*fakePeerManager) RegisterSession(peer.ID, bspm.Session)        {}
+func (*fakePeerManager) UnregisterSession(uint64)                     {}
+func (*fakePeerManager) SendWants(peer.ID, []cid.Cid, []cid.Cid) bool { return true }
+func (*fakePeerManager) BroadcastWantHaves([]cid.Cid)                 {}
+func (fpm *fakePeerManager) SendCancels(cancels []cid.Cid) {
 	fpm.lk.Lock()
 	defer fpm.lk.Unlock()
 	fpm.cancels = append(fpm.cancels, cancels...)
@@ -83,7 +84,7 @@ func (fpm *fakePeerManager) cancelled() []cid.Cid {
 	return fpm.cancels
 }
 
-func sessionFactory(ctx context.Context,
+func sessionFactory(
 	sm bssession.SessionManager,
 	id uint64,
 	sprm bssession.SessionPeerManager,
@@ -94,6 +95,7 @@ func sessionFactory(ctx context.Context,
 	provSearchDelay time.Duration,
 	rebroadcastDelay delay.D,
 	self peer.ID,
+	retrievalState *retrieval.State,
 ) Session {
 	fs := &fakeSession{
 		id:    id,
@@ -101,34 +103,32 @@ func sessionFactory(ctx context.Context,
 		sm:    sm,
 		notif: notif,
 	}
-	go func() {
-		<-ctx.Done()
-		sm.RemoveSession(fs.id)
-	}()
 	return fs
 }
 
-func peerManagerFactory(ctx context.Context, id uint64) bssession.SessionPeerManager {
+func peerManagerFactory(id uint64) bssession.SessionPeerManager {
 	return &fakeSesPeerManager{}
 }
 
 func TestReceiveFrom(t *testing.T) {
-	ctx := context.Background()
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	notif := notifications.New(false)
 	defer notif.Shutdown()
 	sim := bssim.New()
 	bpm := bsbpm.New()
 	pm := &fakePeerManager{}
-	sm := New(ctx, sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
+	sm := New(sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
 
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 
-	firstSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
-	secondSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
-	thirdSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
+	firstSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer firstSession.Close()
+	secondSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer secondSession.Close()
+	thirdSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer thirdSession.Close()
 
 	sim.RecordSessionInterest(firstSession.ID(), []cid.Cid{block.Cid()})
 	sim.RecordSessionInterest(thirdSession.ID(), []cid.Cid{block.Cid()})
@@ -158,22 +158,24 @@ func TestReceiveFrom(t *testing.T) {
 }
 
 func TestReceiveBlocksWhenManagerShutdown(t *testing.T) {
-	ctx := context.Background()
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	notif := notifications.New(false)
 	defer notif.Shutdown()
 	sim := bssim.New()
 	bpm := bsbpm.New()
 	pm := &fakePeerManager{}
-	sm := New(ctx, sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
+	sm := New(sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
 
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 
-	firstSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
-	secondSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
-	thirdSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
+	firstSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer firstSession.Close()
+	secondSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer secondSession.Close()
+	thirdSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer thirdSession.Close()
 
 	sim.RecordSessionInterest(firstSession.ID(), []cid.Cid{block.Cid()})
 	sim.RecordSessionInterest(secondSession.ID(), []cid.Cid{block.Cid()})
@@ -200,21 +202,23 @@ func TestReceiveBlocksWhenSessionContextCancelled(t *testing.T) {
 	sim := bssim.New()
 	bpm := bsbpm.New()
 	pm := &fakePeerManager{}
-	sm := New(ctx, sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
+	sm := New(sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
 
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 
-	firstSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
-	sessionCtx, sessionCancel := context.WithCancel(ctx)
-	secondSession := sm.NewSession(sessionCtx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
-	thirdSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
+	firstSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer firstSession.Close()
+	secondSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer secondSession.Close()
+	thirdSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer thirdSession.Close()
 
 	sim.RecordSessionInterest(firstSession.ID(), []cid.Cid{block.Cid()})
 	sim.RecordSessionInterest(secondSession.ID(), []cid.Cid{block.Cid()})
 	sim.RecordSessionInterest(thirdSession.ID(), []cid.Cid{block.Cid()})
 
-	sessionCancel()
+	secondSession.Close()
 
 	// wait for sessions to get removed
 	time.Sleep(10 * time.Millisecond)
@@ -228,20 +232,20 @@ func TestReceiveBlocksWhenSessionContextCancelled(t *testing.T) {
 }
 
 func TestShutdown(t *testing.T) {
-	ctx := context.Background()
-	ctx, cancel := context.WithCancel(ctx)
+	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 	notif := notifications.New(false)
 	defer notif.Shutdown()
 	sim := bssim.New()
 	bpm := bsbpm.New()
 	pm := &fakePeerManager{}
-	sm := New(ctx, sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
+	sm := New(sessionFactory, sim, peerManagerFactory, bpm, pm, notif, "")
 
 	p := peer.ID(strconv.Itoa(123))
 	block := blocks.NewBlock([]byte("block"))
 	cids := []cid.Cid{block.Cid()}
-	firstSession := sm.NewSession(ctx, time.Second, delay.Fixed(time.Minute)).(*fakeSession)
+	firstSession := sm.NewSession(time.Second, delay.Fixed(time.Minute), nil).(*fakeSession)
+	defer firstSession.Close()
 	sim.RecordSessionInterest(firstSession.ID(), cids)
 	sm.ReceiveFrom(ctx, p, []cid.Cid{}, []cid.Cid{}, cids)