diff --git a/bitswap.go b/bitswap.go index a2bd56ca..f2217b85 100644 --- a/bitswap.go +++ b/bitswap.go @@ -503,11 +503,17 @@ func (bs *Bitswap) Close() error { return bs.process.Close() } -// GetWantlist returns the current local wantlist. +// GetWantlist returns the current local wantlist (both want-blocks and +// want-haves). func (bs *Bitswap) GetWantlist() []cid.Cid { return bs.pm.CurrentWants() } +// GetWantBlocks returns the current list of want-blocks. +func (bs *Bitswap) GetWantBlocks() []cid.Cid { + return bs.pm.CurrentWantBlocks() +} + // GetWanthaves returns the current list of want-haves. func (bs *Bitswap) GetWantHaves() []cid.Cid { return bs.pm.CurrentWantHaves() diff --git a/internal/peermanager/peermanager.go b/internal/peermanager/peermanager.go index ab73fd96..726d4be7 100644 --- a/internal/peermanager/peermanager.go +++ b/internal/peermanager/peermanager.go @@ -170,11 +170,19 @@ func (pm *PeerManager) SendCancels(ctx context.Context, cancelKs []cid.Cid) { } } -// CurrentWants returns the list of pending want-blocks +// CurrentWants returns the list of pending wants (both want-haves and want-blocks). func (pm *PeerManager) CurrentWants() []cid.Cid { pm.pqLk.RLock() defer pm.pqLk.RUnlock() + return pm.pwm.GetWants() +} + +// CurrentWantBlocks returns the list of pending want-blocks +func (pm *PeerManager) CurrentWantBlocks() []cid.Cid { + pm.pqLk.RLock() + defer pm.pqLk.RUnlock() + return pm.pwm.GetWantBlocks() } diff --git a/internal/peermanager/peerwantmanager.go b/internal/peermanager/peerwantmanager.go index 2e8658bc..27e37ccd 100644 --- a/internal/peermanager/peerwantmanager.go +++ b/internal/peermanager/peerwantmanager.go @@ -189,6 +189,28 @@ func (pwm *peerWantManager) GetWantHaves() []cid.Cid { return res.Keys() } +// GetWants returns the set of all wants (both want-blocks and want-haves). +func (pwm *peerWantManager) GetWants() []cid.Cid { + res := cid.NewSet() + + // Iterate over all known peers + for _, pws := range pwm.peerWants { + // Iterate over all want-blocks + for _, c := range pws.wantBlocks.Keys() { + // Add the CID to the results + res.Add(c) + } + + // Iterate over all want-haves + for _, c := range pws.wantHaves.Keys() { + // Add the CID to the results + res.Add(c) + } + } + + return res.Keys() +} + func (pwm *peerWantManager) String() string { var b bytes.Buffer for p, ws := range pwm.peerWants {