diff --git a/dagger/bitswap/engine.nim b/dagger/bitswap/engine.nim index d1f6119e..ac862a52 100644 --- a/dagger/bitswap/engine.nim +++ b/dagger/bitswap/engine.nim @@ -16,6 +16,7 @@ import pkg/libp2p import pkg/libp2p/errors import ./protobuf/bitswap as pb +import ./protobuf/presence import ../blocktype as bt import ../stores/blockstore import ../utils/asyncheapqueue @@ -151,7 +152,7 @@ proc requestBlocks*( proc blockPresenceHandler*( b: BitswapEngine, peer: PeerID, - presence: seq[BlockPresence]) = + blocks: seq[BlockPresence]) = ## Handle block presence ## @@ -159,11 +160,9 @@ proc blockPresenceHandler*( if isNil(peerCtx): return - for blk in presence: - let cid = Cid.init(blk.cid).get() - if cid notin peerCtx.peerHave: - if blk.type == BlockPresenceType.presenceHave: - peerCtx.peerHave.add(cid) + for blk in blocks: + if presence =? Presence.init(blk): + peerCtx.updatePresence(presence) proc scheduleTasks(b: BitswapEngine, blocks: seq[bt.Block]) = trace "Schedule a task for new blocks" diff --git a/dagger/bitswap/peercontext.nim b/dagger/bitswap/peercontext.nim index b57a82e4..52e74cb8 100644 --- a/dagger/bitswap/peercontext.nim +++ b/dagger/bitswap/peercontext.nim @@ -1,10 +1,12 @@ import std/sequtils +import std/tables import pkg/libp2p import pkg/chronos import pkg/nitro import pkg/questionable import ./protobuf/bitswap import ./protobuf/payments +import ./protobuf/presence export payments export nitro @@ -12,6 +14,7 @@ export nitro type BitswapPeerCtx* = ref object of RootObj id*: PeerID + peerPrices*: Table[Cid, UInt256] # remote peer have list including price peerHave*: seq[Cid] # remote peers have lists peerWants*: seq[Entry] # remote peers want lists bytesSent*: int # bytes sent to remote @@ -33,3 +36,13 @@ proc debtRatio*(b: BitswapPeerCtx): float = proc `<`*(a, b: BitswapPeerCtx): bool = a.debtRatio < b.debtRatio +func updatePresence*(context: BitswapPeerCtx, presence: Presence) = + let cid = presence.cid + let price = presence.price + + if cid notin context.peerHave and presence.have: + context.peerHave.add(cid) + context.peerPrices[cid] = price + elif cid in context.peerHave and not presence.have: + context.peerHave.keepItIf(it != cid) + context.peerPrices.del(cid) diff --git a/tests/dagger/bitswap/testengine.nim b/tests/dagger/bitswap/testengine.nim index 14c3d1bf..94f98940 100644 --- a/tests/dagger/bitswap/testengine.nim +++ b/tests/dagger/bitswap/testengine.nim @@ -10,6 +10,7 @@ import pkg/dagger/p2p/rng import pkg/dagger/bitswap import pkg/dagger/bitswap/pendingblocks import pkg/dagger/bitswap/engine/payments +import pkg/dagger/bitswap/protobuf/presence import pkg/dagger/stores/memorystore import pkg/dagger/chunker import pkg/dagger/blocktype as bt @@ -166,15 +167,20 @@ suite "Bitswap engine handlers": await done.wait(100.millis) test "should handle block presence": + let price = UInt256.example engine.blockPresenceHandler( peerId, blocks.mapIt( - BlockPresence( - cid: it.cid.data.buffer, - `type`: BlockPresenceType.presenceHave - ))) + PresenceMessage.init( + Presence( + cid: it.cid, + have: true, + price: price + )))) - check peerCtx.peerHave == blocks.mapIt( it.cid ) + for cid in blocks.mapIt(it.cid): + check peerCtx.peerHave.contains(cid) + check peerCtx.peerPrices[cid] == price suite "Bitswap engine blocks":