From ce59dbd4a21ffe7d0eb9e1fc81c2d9d7cbc0fb30 Mon Sep 17 00:00:00 2001
From: Tanguy <tanguy@status.im>
Date: Wed, 20 Apr 2022 14:28:11 +0200
Subject: [PATCH] Revert "Cleanup engine and rework discovery (#87)"

This reverts commit 4740ffc144916b8a47a1c386e2e8eda728355ce9.
---
 config.nims                                   |   4 +
 dagger/blockexchange/engine.nim               | 430 ++++++++++--------
 dagger/blockexchange/network.nim              |  39 +-
 dagger/blockexchange/networkpeer.nim          |   2 +-
 dagger/blockexchange/peercontext.nim          |  12 +-
 dagger/blockexchange/pendingblocks.nim        |  36 +-
 dagger/dagger.nim                             |   2 +-
 dagger/discovery.nim                          |  36 +-
 dagger/erasure/erasure.nim                    |   2 +-
 dagger/node.nim                               |  30 +-
 dagger/rest/api.nim                           |   7 +-
 dagger/stores/blockstore.nim                  |   3 +-
 dagger/stores/cachestore.nim                  |   5 +-
 dagger/stores/fsstore.nim                     |  26 +-
 dagger/stores/networkstore.nim                |   6 +-
 dagger/utils.nim                              |   4 -
 tests/config.nims                             |   1 +
 .../blockexc/discovery/testdiscovery.nim      | 192 --------
 tests/dagger/blockexc/testblockexc.nim        | 119 +++--
 tests/dagger/blockexc/testengine.nim          |   5 +-
 tests/dagger/helpers/nodeutils.nim            |   8 +-
 .../discovery => }/mockdiscovery.nim          |  27 +-
 tests/dagger/stores/testcachestore.nim        |  12 -
 tests/dagger/stores/testfsstore.nim           |   6 +-
 tests/dagger/testblockexc.nim                 |   1 -
 25 files changed, 417 insertions(+), 598 deletions(-)
 delete mode 100644 dagger/utils.nim
 create mode 100644 tests/config.nims
 delete mode 100644 tests/dagger/blockexc/discovery/testdiscovery.nim
 rename tests/dagger/{blockexc/discovery => }/mockdiscovery.nim (60%)

diff --git a/config.nims b/config.nims
index a906bd99..35023252 100644
--- a/config.nims
+++ b/config.nims
@@ -68,6 +68,10 @@ switch("warning", "ObservableStores:off")
 switch("warning", "LockLevel:off")
 
 switch("define", "libp2p_pki_schemes=secp256k1")
+#TODO this infects everything in this folder, ideally it would only
+# apply to dagger.nim, but since dagger.nims is used for other purpose
+# we can't use it. And dagger.cfg doesn't work
+switch("define", "chronicles_sinks=textlines[dynamic],json[dynamic]")
 
 # begin Nimble config (version 1)
 when system.fileExists("nimble.paths"):
diff --git a/dagger/blockexchange/engine.nim b/dagger/blockexchange/engine.nim
index 4a86c63a..7be6167e 100644
--- a/dagger/blockexchange/engine.nim
+++ b/dagger/blockexchange/engine.nim
@@ -7,8 +7,7 @@
 ## This file may not be copied, modified, or distributed except according to
 ## those terms.
 
-import std/sequtils
-import std/sets
+import std/[sequtils, sets, tables, sugar]
 
 import pkg/chronos
 import pkg/chronicles
@@ -16,7 +15,7 @@ import pkg/libp2p
 
 import ../stores/blockstore
 import ../blocktype as bt
-import ../utils
+import ../utils/asyncheapqueue
 import ../discovery
 
 import ./protobuf/blockexc
@@ -33,19 +32,31 @@ logScope:
   topics = "dagger blockexc engine"
 
 const
+  DefaultBlockTimeout* = 5.minutes
   DefaultMaxPeersPerRequest* = 10
   DefaultTaskQueueSize = 100
   DefaultConcurrentTasks = 10
   DefaultMaxRetries = 3
-  DefaultConcurrentDiscRequests = 10
-  DefaultConcurrentAdvertRequests = 10
-  DefaultDiscoveryTimeout = 1.minutes
-  DefaultMaxQueriedBlocksCache = 1000
+
+  # Current advertisement is meant to be more efficient than
+  # correct, so blocks could be advertised more slowly than that
+  # Put some margin
+  BlockAdvertisementFrequency = 30.minutes
 
 type
   TaskHandler* = proc(task: BlockExcPeerCtx): Future[void] {.gcsafe.}
   TaskScheduler* = proc(task: BlockExcPeerCtx): bool {.gcsafe.}
 
+  BlockDiscovery* = ref object
+    discoveredProvider: AsyncEvent
+    discoveryLoop: Future[void]
+    toDiscover: Cid
+    treatedPeer: HashSet[PeerId]
+    inflightIWant: HashSet[PeerId]
+    gotIWantResponse: AsyncEvent
+    provides: seq[PeerId]
+    lastDhtQuery: Moment
+
   BlockExcEngine* = ref object of RootObj
     localStore*: BlockStore                       # where we localStore blocks for this instance
     network*: BlockExcNetwork                     # network interface
@@ -59,15 +70,12 @@ type
     peersPerRequest: int                          # max number of peers to request from
     wallet*: WalletRef                            # nitro wallet for micropayments
     pricing*: ?Pricing                            # optional bandwidth pricing
-    discovery*: Discovery                         # Discovery interface
-    concurrentAdvReqs: int                        # Concurrent advertise requests
-    advertiseLoop*: Future[void]                  # Advertise loop task handle
-    advertiseQueue*: AsyncQueue[Cid]              # Advertise queue
-    advertiseTasks*: seq[Future[void]]            # Advertise tasks
-    concurrentDiscReqs: int                       # Concurrent discovery requests
-    discoveryLoop*: Future[void]                  # Discovery loop task handle
-    discoveryTasks*: seq[Future[void]]            # Discovery tasks
-    discoveryQueue*: AsyncQueue[Cid]              # Discovery queue
+    advertisedBlocks: seq[Cid]
+    advertisedIndex: int
+    advertisementFrequency: Duration
+    runningDiscoveries*: Table[Cid, BlockDiscovery]
+    blockAdded: AsyncEvent
+    discovery*: Discovery
 
   Pricing* = object
     address*: EthAddress
@@ -92,76 +100,7 @@ proc scheduleTask(b: BlockExcEngine, task: BlockExcPeerCtx): bool {.gcsafe} =
   b.taskQueue.pushOrUpdateNoWait(task).isOk()
 
 proc blockexcTaskRunner(b: BlockExcEngine): Future[void] {.gcsafe.}
-
-proc discoveryLoopRunner(b: BlockExcEngine) {.async.} =
-  while b.blockexcRunning:
-    for cid in toSeq(b.pendingBlocks.wantList):
-      try:
-        await b.discoveryQueue.put(cid)
-      except CatchableError as exc:
-        trace "Exception in discovery loop", exc = exc.msg
-
-    trace "About to sleep, number of wanted blocks", wanted = b.pendingBlocks.len
-    await sleepAsync(30.seconds)
-
-proc advertiseLoopRunner*(b: BlockExcEngine) {.async.} =
-  proc onBlock(cid: Cid) {.async.} =
-    try:
-      await b.advertiseQueue.put(cid)
-    except CatchableError as exc:
-      trace "Exception listing blocks", exc = exc.msg
-
-  while b.blockexcRunning:
-    await b.localStore.listBlocks(onBlock)
-    await sleepAsync(30.seconds)
-
-  trace "Exiting advertise task loop"
-
-proc advertiseTaskRunner(b: BlockExcEngine) {.async.} =
-  ## Run advertise tasks
-  ##
-
-  while b.blockexcRunning:
-    try:
-      let cid = await b.advertiseQueue.get()
-      await b.discovery.provideBlock(cid)
-    except CatchableError as exc:
-      trace "Exception in advertise task runner", exc = exc.msg
-
-  trace "Exiting advertise task runner"
-
-proc discoveryTaskRunner(b: BlockExcEngine) {.async.} =
-  ## Run discovery tasks
-  ##
-
-  while b.blockexcRunning:
-    try:
-      let
-        cid = await b.discoveryQueue.get()
-        providers = await b.discovery
-          .findBlockProviders(cid)
-          .wait(DefaultDiscoveryTimeout)
-
-      await allFuturesThrowing(
-        allFinished(providers.mapIt( b.network.dialPeer(it.data) )))
-    except CatchableError as exc:
-      trace "Exception in discovery task runner", exc = exc.msg
-
-  trace "Exiting discovery task runner"
-
-proc queueFindBlocksReq(b: BlockExcEngine, cids: seq[Cid]) {.async.} =
-  try:
-    for cid in cids:
-      await b.discoveryQueue.put(cid)
-  except CatchableError as exc:
-    trace "Exception queueing discovery request", exc = exc.msg
-
-proc queueProvideBlocksReq(b: BlockExcEngine, cids: seq[Cid]) {.async.} =
-  try:
-    for cid in cids:
-      await b.advertiseQueue.put(cid)
-  except CatchableError as exc:
-    trace "Exception queueing discovery request", exc = exc.msg
+proc advertiseLoop(b: BlockExcEngine): Future[void] {.gcsafe.}
 
 proc start*(b: BlockExcEngine) {.async.} =
   ## Start the blockexc task
@@ -177,14 +116,13 @@ proc start*(b: BlockExcEngine) {.async.} =
   for i in 0..<b.concurrentTasks:
     b.blockexcTasks.add(blockexcTaskRunner(b))
 
-  for i in 0..<b.concurrentAdvReqs:
-    b.advertiseTasks.add(advertiseTaskRunner(b))
+  info "Getting existing block list"
+  let blocks = await b.localStore.blockList()
+  b.advertisedBlocks = blocks
+  # We start faster to publish everything ASAP
+  b.advertisementFrequency = 5.seconds
 
-  for i in 0..<b.concurrentDiscReqs:
-    b.discoveryTasks.add(discoveryTaskRunner(b))
-
-  b.advertiseLoop = advertiseLoopRunner(b)
-  b.discoveryLoop = discoveryLoopRunner(b)
+  b.blockexcTasks.add(b.advertiseLoop())
 
 proc stop*(b: BlockExcEngine) {.async.} =
   ## Stop the blockexc blockexc
@@ -202,87 +140,156 @@ proc stop*(b: BlockExcEngine) {.async.} =
       await t.cancelAndWait()
       trace "Task stopped"
 
-  for t in b.advertiseTasks:
-    if not t.finished:
-      trace "Awaiting task to stop"
-      await t.cancelAndWait()
-      trace "Task stopped"
+  for _, bd in b.runningDiscoveries:
+    await bd.discoveryLoop.cancelAndWait()
 
-  for t in b.discoveryTasks:
-    if not t.finished:
-      trace "Awaiting task to stop"
-      await t.cancelAndWait()
-      trace "Task stopped"
-
-  if not b.advertiseLoop.isNil and not b.advertiseLoop.finished:
-    trace "Awaiting advertise loop to stop"
-    await b.advertiseLoop.cancelAndWait()
-    trace "Advertise loop stopped"
-
-  if not b.discoveryLoop.isNil and not b.discoveryLoop.finished:
-    trace "Awaiting discovery loop to stop"
-    await b.discoveryLoop.cancelAndWait()
-    trace "Discovery loop stopped"
+  b.runningDiscoveries.clear()
 
   trace "NetworkStore stopped"
 
+proc discoverOnDht(b: BlockExcEngine, bd: BlockDiscovery) {.async.} =
+  bd.lastDhtQuery = Moment.fromNow(10.hours)
+  defer: bd.lastDhtQuery = Moment.now()
+
+  let discoveredProviders = await b.discovery.findBlockProviders(bd.toDiscover)
+
+  for peer in discoveredProviders:
+    asyncSpawn b.network.dialPeer(peer.data)
+
+proc discoverLoop(b: BlockExcEngine, bd: BlockDiscovery) {.async.} =
+  # First, try connected peers
+  # After a percent of peers declined, or a timeout passed, query DHT
+  # rinse & repeat
+  #
+  # TODO add a global timeout
+
+  debug "starting block discovery", cid=bd.toDiscover
+
+  bd.gotIWantResponse.fire()
+  while true:
+    # wait for iwant replies
+    await bd.gotIWantResponse.wait()
+    bd.gotIWantResponse.clear()
+
+    var foundPeerNew = false
+    for p in b.peers:
+      if bd.toDiscover in p.peerHave and p.id notin bd.treatedPeer:
+        bd.provides.add(p.id)
+        bd.treatedPeer.incl(p.id)
+        bd.inflightIWant.excl(p.id)
+        foundPeerNew = true
+
+    if foundPeerNew:
+      bd.discoveredProvider.fire()
+      continue
+
+    trace "asking peers", cid=bd.toDiscover, peers=b.peers.len, treated=bd.treatedPeer.len, inflight=bd.inflightIWant.len
+    for p in b.peers:
+      if p.id notin bd.treatedPeer and p.id notin bd.inflightIWant:
+        # just send wants
+        bd.inflightIWant.incl(p.id)
+        b.network.request.sendWantList(
+          p.id,
+          @[bd.toDiscover],
+          wantType = WantType.wantHave,
+          sendDontHave = true)
+
+    if bd.inflightIWant.len < 3 and #TODO or a timeout
+      bd.lastDhtQuery < Moment.now() - 5.seconds:
+        #start query
+        asyncSpawn b.discoverOnDht(bd)
+
+proc discoverBlock*(b: BlockExcEngine, cid: Cid): BlockDiscovery =
+  if cid in b.runningDiscoveries:
+    return b.runningDiscoveries[cid]
+  else:
+    result = BlockDiscovery(
+      toDiscover: cid,
+      discoveredProvider: newAsyncEvent(),
+      gotIWantResponse: newAsyncEvent(),
+    )
+    result.discoveryLoop = b.discoverLoop(result)
+    b.runningDiscoveries[cid] = result
+    return result
+
+proc stopDiscovery(b: BlockExcEngine, cid: Cid) =
+  if cid in b.runningDiscoveries:
+    b.runningDiscoveries[cid].discoveryLoop.cancel()
+    b.runningDiscoveries.del(cid)
+
 proc requestBlock*(
   b: BlockExcEngine,
   cid: Cid,
-  timeout = DefaultBlockTimeout): Future[bt.Block] =
+  timeout = DefaultBlockTimeout): Future[bt.Block] {.async.} =
   ## Request a block from remotes
   ##
 
+  debug "requesting block", cid
+
+  # TODO
+  # we could optimize "groups of related chunks"
+  # be requesting multiple chunks, and running discovery
+  # less often
+
+  if cid in b.localStore:
+    return (await b.localStore.getBlock(cid)).get()
+
+  # be careful, don't give back control to main loop here
+  # otherwise, the block might slip in
+
+  if cid in b.pendingBlocks:
+    return await b.pendingBlocks.blocks[cid].wait(timeout)
+
+  # We are the first one to request this block, so we handle it
   let
-    blk = b.pendingBlocks.getWantHandle(cid, timeout)
+    timeoutFut = sleepAsync(timeout)
+    blk = b.pendingBlocks.addOrAwait(cid)
+    discovery = b.discoverBlock(cid)
 
-  if b.peers.len <= 0:
-    trace "No peers to request blocks from", cid = $cid
-    asyncSpawn b.queueFindBlocksReq(@[cid])
-    return blk
+  # Just take the first discovered peer
+  try:
+    await timeoutFut or blk or discovery.discoveredProvider.wait()
+    discovery.discoveredProvider.clear()
+  except CancelledError as exc:
+    #TODO also wrong, same issue as below
+    blk.cancel()
+    b.stopDiscovery(cid)
+    raise exc
 
-  var peers = b.peers
+  if timeoutFut.finished:
+    # TODO this is wrong, because other user may rely on us
+    # to handle this block. This proc should be asyncSpawned
+    #
+    # Other people may be using the discovery or blk
+    # so don't kill them
+    blk.cancel()
+    b.stopDiscovery(cid)
+    raise newException(AsyncTimeoutError, "")
 
-  # get the first peer with at least one (any)
-  # matching cid
-  var blockPeer: BlockExcPeerCtx
-  for p in peers:
-    if cid in p.peerHave:
-      blockPeer = p
-      break
+  if blk.finished:
+    # a peer sent us the block out of the blue, why not
+    b.stopDiscovery(cid)
+    return await blk
 
-  # didn't find any peer with matching cids
-  # use the first one in the sorted array
-  if isNil(blockPeer):
-    blockPeer = peers[0]
+  # We got a provider
+  # Currently, we just ask him for the block, and hope he gives it to us
+  #
+  # In reality, we could keep discovering until we find a suitable price, etc
+  b.stopDiscovery(cid)
+  timeoutFut.cancel()
 
-  peers.keepItIf(
-    it != blockPeer
-  )
+  assert discovery.provides.len > 0
 
+  debug "Requesting block from peer", providerCount = discovery.provides.len,
+    peer = discovery.provides[0], cid
   # request block
   b.network.request.sendWantList(
-    blockPeer.id,
+    discovery.provides[0],
     @[cid],
     wantType = WantType.wantBlock) # we want this remote to send us a block
 
-  if peers.len == 0:
-    trace "Not enough peers to send want list to", cid = $cid
-    asyncSpawn b.queueFindBlocksReq(@[cid])
-    return blk # no peers to send wants to
-
-  # filter out the peer we've already requested from
-  let stop = min(peers.high, b.peersPerRequest)
-  trace "Sending want list requests to remaining peers", count = stop + 1
-  for p in peers[0..stop]:
-    if cid notin p.peerHave:
-      # just send wants
-      b.network.request.sendWantList(
-        p.id,
-        @[cid],
-        wantType = WantType.wantHave) # we only want to know if the peer has the block
-
-  return blk
+  #TODO substract the discovery time
+  return await blk.wait(timeout)
 
 proc blockPresenceHandler*(
   b: BlockExcEngine,
@@ -291,25 +298,18 @@ proc blockPresenceHandler*(
   ## Handle block presence
   ##
 
-  trace "Received presence update for peer", peer
   let peerCtx = b.getPeerCtx(peer)
-  if isNil(peerCtx):
-    return
 
   for blk in blocks:
     if presence =? Presence.init(blk):
-      peerCtx.updatePresence(presence)
-
-  let
-    cids = toSeq(b.pendingBlocks.wantList).filterIt(
-      it in peerCtx.peerHave
-    )
-
-  if cids.len > 0:
-    b.network.request.sendWantList(
-      peerCtx.id,
-      cids,
-      wantType = WantType.wantBlock) # we want this remote to send us a block
+      if not isNil(peerCtx):
+        peerCtx.updatePresence(presence)
+      if presence.cid in b.runningDiscoveries:
+        let bd = b.runningDiscoveries[presence.cid]
+        if not presence.have:
+          bd.inflightIWant.excl(peer)
+          bd.treatedPeer.incl(peer)
+        bd.gotIWantResponse.fire()
 
 proc scheduleTasks(b: BlockExcEngine, blocks: seq[bt.Block]) =
   trace "Schedule a task for new blocks"
@@ -330,11 +330,21 @@ proc resolveBlocks*(b: BlockExcEngine, blocks: seq[bt.Block]) =
   ## and schedule any new task to be ran
   ##
 
-  trace "Resolving blocks", blocks = blocks.len
+  trace "Resolving blocks"
 
-  b.pendingBlocks.resolve(blocks)
-  b.scheduleTasks(blocks)
-  asyncSpawn b.queueProvideBlocksReq(blocks.mapIt( it.cid ))
+  var gotNewBlocks = false
+  for bl in blocks:
+    if bl.cid notin b.advertisedBlocks: #TODO that's very slow, maybe a ordered hashset instead
+      #TODO could do some smarter ordering here (insert it just before b.advertisedIndex, or similar)
+      b.advertisedBlocks.add(bl.cid)
+      asyncSpawn b.discovery.publishProvide(bl.cid)
+      gotNewBlocks = true
+
+  if gotNewBlocks:
+    b.pendingBlocks.resolve(blocks)
+    b.scheduleTasks(blocks)
+
+    b.blockAdded.fire()
 
 proc payForBlocks(engine: BlockExcEngine,
                   peer: BlockExcPeerCtx,
@@ -410,20 +420,14 @@ proc wantListHandler*(
   if not b.scheduleTask(peerCtx):
     trace "Unable to schedule task for peer", peer
 
-proc accountHandler*(
-  engine: BlockExcEngine,
-  peer: PeerID,
-  account: Account) {.async.} =
+proc accountHandler*(engine: BlockExcEngine, peer: PeerID, account: Account) {.async.} =
   let context = engine.getPeerCtx(peer)
   if context.isNil:
     return
 
   context.account = account.some
 
-proc paymentHandler*(
-  engine: BlockExcEngine,
-  peer: PeerId,
-  payment: SignedState) {.async.} =
+proc paymentHandler*(engine: BlockExcEngine, peer: PeerId, payment: SignedState) {.async.} =
   without context =? engine.getPeerCtx(peer).option and
           account =? context.account:
     return
@@ -446,8 +450,13 @@ proc setupPeer*(b: BlockExcEngine, peer: PeerID) =
     ))
 
   # broadcast our want list, the other peer will do the same
-  if b.pendingBlocks.len > 0:
-    b.network.request.sendWantList(peer, toSeq(b.pendingBlocks.wantList), full = true)
+  let wantList = collect(newSeqOfCap(b.runningDiscoveries.len)):
+    for cid, bd in b.runningDiscoveries:
+      bd.inflightIWant.incl(peer)
+      cid
+
+  if wantList.len > 0:
+    b.network.request.sendWantList(peer, wantList, full = true, sendDontHave = true)
 
   if address =? b.pricing.?address:
     b.network.request.sendAccount(peer, Account(address: address))
@@ -461,6 +470,31 @@ proc dropPeer*(b: BlockExcEngine, peer: PeerID) =
   # drop the peer from the peers table
   b.peers.keepItIf( it.id != peer )
 
+proc advertiseLoop(b: BlockExcEngine) {.async, gcsafe.} =
+  while true:
+    if b.advertisedIndex >= b.advertisedBlocks.len:
+      b.advertisedIndex = 0
+      b.advertisementFrequency = BlockAdvertisementFrequency
+
+    # check that we still have this block.
+    while
+      b.advertisedIndex < b.advertisedBlocks.len and
+      not(b.localStore.contains(b.advertisedBlocks[b.advertisedIndex])):
+        b.advertisedBlocks.delete(b.advertisedIndex)
+
+    #publish it
+    if b.advertisedIndex < b.advertisedBlocks.len:
+      asyncSpawn b.discovery.publishProvide(b.advertisedBlocks[b.advertisedIndex])
+
+    inc b.advertisedIndex
+    let toSleep =
+      if b.advertisedBlocks.len > 0:
+        b.advertisementFrequency div b.advertisedBlocks.len
+      else:
+        30.minutes
+    await sleepAsync(toSleep) or b.blockAdded.wait()
+    b.blockAdded.clear()
+
 proc taskHandler*(b: BlockExcEngine, task: BlockExcPeerCtx) {.gcsafe, async.} =
   trace "Handling task for peer", peer = task.id
 
@@ -482,7 +516,6 @@ proc taskHandler*(b: BlockExcEngine, task: BlockExcPeerCtx) {.gcsafe, async.} =
       .mapIt(!it.read)
 
     if blocks.len > 0:
-      trace "Sending blocks to peer", peer = task.id, blocks = blocks.len
       b.network.request.sendBlocks(
         task.id,
         blocks)
@@ -525,25 +558,19 @@ proc new*(
   discovery: Discovery,
   concurrentTasks = DefaultConcurrentTasks,
   maxRetries = DefaultMaxRetries,
-  peersPerRequest = DefaultMaxPeersPerRequest,
-  concurrentAdvReqs = DefaultConcurrentAdvertRequests,
-  concurrentDiscReqs = DefaultConcurrentDiscRequests): T =
+  peersPerRequest = DefaultMaxPeersPerRequest): T =
 
-  let
-    engine = BlockExcEngine(
-      localStore: localStore,
-      pendingBlocks: PendingBlocksManager.new(),
-      peersPerRequest: peersPerRequest,
-      network: network,
-      wallet: wallet,
-      concurrentTasks: concurrentTasks,
-      concurrentAdvReqs: concurrentAdvReqs,
-      concurrentDiscReqs: concurrentDiscReqs,
-      maxRetries: maxRetries,
-      taskQueue: newAsyncHeapQueue[BlockExcPeerCtx](DefaultTaskQueueSize),
-      discovery: discovery,
-      advertiseQueue: newAsyncQueue[Cid](DefaultTaskQueueSize),
-      discoveryQUeue: newAsyncQueue[Cid](DefaultTaskQueueSize))
+  let engine = BlockExcEngine(
+    localStore: localStore,
+    pendingBlocks: PendingBlocksManager.new(),
+    blockAdded: newAsyncEvent(),
+    peersPerRequest: peersPerRequest,
+    network: network,
+    wallet: wallet,
+    concurrentTasks: concurrentTasks,
+    maxRetries: maxRetries,
+    discovery: discovery,
+    taskQueue: newAsyncHeapQueue[BlockExcPeerCtx](DefaultTaskQueueSize))
 
   proc peerEventHandler(peerId: PeerID, event: PeerEvent) {.async.} =
     if event.kind == PeerEventKind.Joined:
@@ -581,6 +608,7 @@ proc new*(
     onBlocks: blocksHandler,
     onPresence: blockPresenceHandler,
     onAccount: accountHandler,
-    onPayment: paymentHandler)
+    onPayment: paymentHandler
+  )
 
   return engine
diff --git a/dagger/blockexchange/network.nim b/dagger/blockexchange/network.nim
index af212c26..3f8729c6 100644
--- a/dagger/blockexchange/network.nim
+++ b/dagger/blockexchange/network.nim
@@ -120,16 +120,14 @@ proc broadcastWantList*(
 
   trace "Sending want list to peer", peer = id, `type` = $wantType, len = cids.len
 
-  let
-    wantList = makeWantList(
-      cids,
-      priority,
-      cancel,
-      wantType,
-      full,
-      sendDontHave)
-  b.peers.withValue(id, peer):
-    peer[].broadcast(Message(wantlist: wantList))
+  let wantList = makeWantList(
+    cids,
+    priority,
+    cancel,
+    wantType,
+    full,
+    sendDontHave)
+  b.peers[id].broadcast(Message(wantlist: wantList))
 
 proc handleBlocks(
   b: BlockExcNetwork,
@@ -155,7 +153,9 @@ proc handleBlocks(
 
   b.handlers.onBlocks(peer.id, blks)
 
-template makeBlocks*(blocks: seq[bt.Block]): seq[pb.Block] =
+template makeBlocks*(
+  blocks: seq[bt.Block]):
+  seq[pb.Block] =
   var blks: seq[pb.Block]
   for blk in blocks:
     blks.add(pb.Block(
@@ -176,8 +176,7 @@ proc broadcastBlocks*(
     return
 
   trace "Sending blocks to peer", peer = id, len = blocks.len
-  b.peers.withValue(id, peer):
-    peer[].broadcast(pb.Message(payload: makeBlocks(blocks)))
+  b.peers[id].broadcast(pb.Message(payload: makeBlocks(blocks)))
 
 proc handleBlockPresence(
   b: BlockExcNetwork,
@@ -203,8 +202,7 @@ proc broadcastBlockPresence*(
     return
 
   trace "Sending presence to peer", peer = id
-  b.peers.withValue(id, peer):
-    peer[].broadcast(Message(blockPresences: @presence))
+  b.peers[id].broadcast(Message(blockPresences: presence))
 
 proc handleAccount(network: BlockExcNetwork,
                    peer: NetworkPeer,
@@ -220,8 +218,7 @@ proc broadcastAccount*(network: BlockExcNetwork,
     return
 
   let message = Message(account: AccountMessage.init(account))
-  network.peers.withValue(id, peer):
-    peer[].broadcast(message)
+  network.peers[id].broadcast(message)
 
 proc broadcastPayment*(network: BlockExcNetwork,
                        id: PeerId,
@@ -230,8 +227,7 @@ proc broadcastPayment*(network: BlockExcNetwork,
     return
 
   let message = Message(payment: StateChannelUpdate.init(payment))
-  network.peers.withValue(id, peer):
-    peer[].broadcast(message)
+  network.peers[id].broadcast(message)
 
 proc handlePayment(network: BlockExcNetwork,
                    peer: NetworkPeer,
@@ -265,7 +261,7 @@ proc getOrCreatePeer(b: BlockExcNetwork, peer: PeerID): NetworkPeer =
   ##
 
   if peer in b.peers:
-    return b.peers.getOrDefault(peer, nil)
+    return b.peers[peer]
 
   var getConn = proc(): Future[Connection] {.async.} =
     try:
@@ -367,7 +363,8 @@ proc new*(
     sendBlocks: sendBlocks,
     sendPresence: sendPresence,
     sendAccount: sendAccount,
-    sendPayment: sendPayment)
+    sendPayment: sendPayment
+  )
 
   b.init()
   return b
diff --git a/dagger/blockexchange/networkpeer.nim b/dagger/blockexchange/networkpeer.nim
index 2ef8d473..269291cc 100644
--- a/dagger/blockexchange/networkpeer.nim
+++ b/dagger/blockexchange/networkpeer.nim
@@ -17,7 +17,7 @@ import ./protobuf/blockexc
 logScope:
   topics = "dagger blockexc networkpeer"
 
-const MaxMessageSize = 100 * 1024 * 1024 # manifest files can be big
+const MaxMessageSize = 8 * 1024 * 1024
 
 type
   RPCHandler* = proc(peer: NetworkPeer, msg: Message): Future[void] {.gcsafe.}
diff --git a/dagger/blockexchange/peercontext.nim b/dagger/blockexchange/peercontext.nim
index 2fb24216..99c7b15f 100644
--- a/dagger/blockexchange/peercontext.nim
+++ b/dagger/blockexchange/peercontext.nim
@@ -13,12 +13,12 @@ export payments, nitro
 type
   BlockExcPeerCtx* = ref object of RootObj
     id*: PeerID
-    peerPrices*: Table[Cid, UInt256]  # remote peer have list including price
-    peerWants*: seq[Entry]            # remote peers want lists
-    exchanged*: int                   # times peer has exchanged with us
-    lastExchange*: Moment             # last time peer has exchanged with us
-    account*: ?Account                # ethereum account of this peer
-    paymentChannel*: ?ChannelId       # payment channel id
+    peerPrices*: Table[Cid, UInt256] # remote peer have list including price
+    peerWants*: seq[Entry]      # remote peers want lists
+    exchanged*: int             # times peer has exchanged with us
+    lastExchange*: Moment       # last time peer has exchanged with us
+    account*: ?Account          # ethereum account of this peer
+    paymentChannel*: ?ChannelId # payment channel id
 
 proc peerHave*(context: BlockExcPeerCtx): seq[Cid] =
   toSeq(context.peerPrices.keys)
diff --git a/dagger/blockexchange/pendingblocks.nim b/dagger/blockexchange/pendingblocks.nim
index 3d504f9b..3b426bb8 100644
--- a/dagger/blockexchange/pendingblocks.nim
+++ b/dagger/blockexchange/pendingblocks.nim
@@ -8,11 +8,6 @@
 ## those terms.
 
 import std/tables
-import std/sequtils
-
-import pkg/upraises
-
-push: {.upraises: [].}
 
 import pkg/questionable
 import pkg/chronicles
@@ -24,22 +19,18 @@ import ../blocktype
 logScope:
   topics = "dagger blockexc pendingblocks"
 
-const
-  DefaultBlockTimeout* = 10.minutes
-
 type
   PendingBlocksManager* = ref object of RootObj
     blocks*: Table[Cid, Future[Block]] # pending Block requests
 
-proc getWantHandle*(
+proc addOrAwait*(
   p: PendingBlocksManager,
-  cid: Cid,
-  timeout = DefaultBlockTimeout): Future[Block] {.async.} =
+  cid: Cid): Future[Block] {.async.} =
   ## Add an event for a block
   ##
 
   if cid notin p.blocks:
-     p.blocks[cid] = newFuture[Block]().wait(timeout)
+     p.blocks[cid] = newFuture[Block]()
      trace "Adding pending future for block", cid
 
   try:
@@ -61,11 +52,11 @@ proc resolve*(
   for blk in blocks:
     # resolve any pending blocks
     if blk.cid in p.blocks:
-      p.blocks.withValue(blk.cid, pending):
-        if not pending[].finished:
-          trace "Resolving block", cid = $blk.cid
-          pending[].complete(blk)
-          p.blocks.del(blk.cid)
+      let pending = p.blocks[blk.cid]
+      if not pending.finished:
+        trace "Resolving block", cid = $blk.cid
+        pending.complete(blk)
+        p.blocks.del(blk.cid)
 
 proc pending*(
   p: PendingBlocksManager,
@@ -75,17 +66,6 @@ proc contains*(
   p: PendingBlocksManager,
   cid: Cid): bool = p.pending(cid)
 
-iterator wantList*(p: PendingBlocksManager): Cid =
-  for k in p.blocks.keys:
-    yield k
-
-iterator wantHandles*(p: PendingBlocksManager): Future[Block] =
-  for v in p.blocks.values:
-    yield v
-
-func len*(p: PendingBlocksManager): int =
-  p.blocks.len
-
 func new*(T: type PendingBlocksManager): T =
   T(
     blocks: initTable[Cid, Future[Block]]()
diff --git a/dagger/dagger.nim b/dagger/dagger.nim
index 48ee6eac..db9a0222 100644
--- a/dagger/dagger.nim
+++ b/dagger/dagger.nim
@@ -123,7 +123,7 @@ proc new*(T: type DaggerServer, config: DaggerConf): T =
     erasure = Erasure.new(store, leoEncoderProvider, leoDecoderProvider)
     daggerNode = DaggerNodeRef.new(switch, store, engine, erasure, discovery)
     restServer = RestServerRef.new(
-      daggerNode.initRestApi(config),
+      daggerNode.initRestApi(),
       initTAddress("127.0.0.1" , config.apiPort),
       bufferSize = (1024 * 64),
       maxRequestBodySize = int.high)
diff --git a/dagger/discovery.nim b/dagger/discovery.nim
index 96c961a9..60d9f999 100644
--- a/dagger/discovery.nim
+++ b/dagger/discovery.nim
@@ -8,20 +8,18 @@
 ## those terms.
 
 import pkg/chronos
-import pkg/chronicles
 import pkg/libp2p
 import pkg/questionable
 import pkg/questionable/results
 import pkg/stew/shims/net
 import pkg/libp2pdht/discv5/protocol as discv5
 
-import ./rng
-import ./errors
+import rng
 
 export discv5
 
 type
-  Discovery* = ref object of RootObj
+  Discovery* = ref object
     protocol: discv5.Protocol
     localInfo: PeerInfo
 
@@ -57,33 +55,15 @@ proc toDiscoveryId*(cid: Cid): NodeId =
   ## To discovery id
   readUintBE[256](keccak256.digest(cid.data.buffer).data)
 
-method findBlockProviders*(
+proc findBlockProviders*(
   d: Discovery,
-  cid: Cid): Future[seq[SignedPeerRecord]] {.async, base.} =
-  ## Find block providers
-  ##
+  cid: Cid): Future[seq[SignedPeerRecord]] {.async.} =
+  return (await d.protocol.getProviders(cid.toDiscoveryId())).get()
 
-  trace "Finding providers for block", cid = $cid
-  without providers =?
-    (await d.protocol.getProviders(cid.toDiscoveryId())).mapFailure, error:
-    trace "Error finding providers for block", cid = $cid, error = error.msg
+proc publishProvide*(d: Discovery, cid: Cid) {.async.} =
+  let bid = cid.toDiscoveryId()
+  discard await d.protocol.addProvider(bid, d.localInfo.signedPeerRecord)
 
-  return providers
-
-method provideBlock*(d: Discovery, cid: Cid) {.async, base.} =
-  ## Provide a bock Cid
-  ##
-
-  trace "Providing block", cid = $cid
-  let
-    nodes = await d.protocol.addProvider(
-      cid.toDiscoveryId(),
-      d.localInfo.signedPeerRecord)
-
-  if nodes.len <= 0:
-    trace "Couldn't provide to any nodes!"
-
-  trace "Provided to nodes", nodes = nodes.len
 
 proc start*(d: Discovery) {.async.} =
   d.protocol.updateRecord(d.localInfo.signedPeerRecord).expect("updating SPR")
diff --git a/dagger/erasure/erasure.nim b/dagger/erasure/erasure.nim
index cf52ba94..bb8970c4 100644
--- a/dagger/erasure/erasure.nim
+++ b/dagger/erasure/erasure.nim
@@ -109,7 +109,7 @@ proc encode*(
       # TODO: this is a tight blocking loop so we sleep here to allow
       # other events to be processed, this should be addressed
       # by threading
-      await sleepAsync(100.millis)
+      await sleepAsync(10.millis)
 
       for j in 0..<blocks:
         let idx = blockIdx[j]
diff --git a/dagger/node.nim b/dagger/node.nim
index 329b14fb..eb849f33 100644
--- a/dagger/node.nim
+++ b/dagger/node.nim
@@ -44,17 +44,10 @@ type
     discovery*: Discovery
 
 proc start*(node: DaggerNodeRef) {.async.} =
-  if not node.switch.isNil:
-    await node.switch.start()
-
-  if not node.engine.isNil:
-    await node.engine.start()
-
-  if not node.erasure.isNil:
-    await node.erasure.start()
-
-  if not node.discovery.isNil:
-    await node.discovery.start()
+  await node.switch.start()
+  await node.engine.start()
+  await node.erasure.start()
+  await node.discovery.start()
 
   node.networkId = node.switch.peerInfo.peerId
   notice "Started dagger node", id = $node.networkId, addrs = node.switch.peerInfo.addrs
@@ -62,17 +55,10 @@ proc start*(node: DaggerNodeRef) {.async.} =
 proc stop*(node: DaggerNodeRef) {.async.} =
   trace "Stopping node"
 
-  if not node.engine.isNil:
-    await node.engine.stop()
-
-  if not node.switch.isNil:
-    await node.switch.stop()
-
-  if not node.erasure.isNil:
-    await node.erasure.stop()
-
-  if not node.discovery.isNil:
-    await node.discovery.stop()
+  await node.engine.stop()
+  await node.switch.stop()
+  await node.erasure.stop()
+  await node.discovery.stop()
 
 proc findPeer*(
   node: DaggerNodeRef,
diff --git a/dagger/rest/api.nim b/dagger/rest/api.nim
index f3442cdd..a6143688 100644
--- a/dagger/rest/api.nim
+++ b/dagger/rest/api.nim
@@ -21,13 +21,11 @@ import pkg/chronos
 import pkg/presto
 import pkg/libp2p
 import pkg/stew/base10
-import pkg/confutils
 
 import pkg/libp2p/routing_record
 
 import ../node
 import ../blocktype
-import ../conf
 
 proc validate(
   pattern: string,
@@ -85,7 +83,7 @@ proc decodeString(T: type bool, value: string): Result[T, cstring] =
 proc encodeString(value: bool): Result[string, cstring] =
   ok($value)
 
-proc initRestApi*(node: DaggerNodeRef, conf: DaggerConf): RestRouter =
+proc initRestApi*(node: DaggerNodeRef): RestRouter =
   var router = RestRouter.init(validate)
   router.api(
     MethodGet,
@@ -320,7 +318,6 @@ proc initRestApi*(node: DaggerNodeRef, conf: DaggerConf): RestRouter =
 
       return RestApiResponse.response(
         "Id: " & $node.switch.peerInfo.peerId &
-        "\nAddrs: \n" & addrs &
-        "\nRoot Dir: " & $conf.dataDir)
+        "\nAddrs: \n" & addrs & "\n")
 
   return router
diff --git a/dagger/stores/blockstore.nim b/dagger/stores/blockstore.nim
index ddcbc550..90a06173 100644
--- a/dagger/stores/blockstore.nim
+++ b/dagger/stores/blockstore.nim
@@ -20,7 +20,6 @@ import ../blocktype
 export blocktype, libp2p
 
 type
-  OnBlock* = proc(cid: Cid): Future[void] {.upraises: [], gcsafe.}
   BlockStore* = ref object of RootObj
 
 method getBlock*(
@@ -53,7 +52,7 @@ method hasBlock*(s: BlockStore, cid: Cid): bool {.base.} =
 
   return false
 
-method listBlocks*(s: BlockStore, onBlock: OnBlock): Future[void] {.base.} =
+method blockList*(s: BlockStore): Future[seq[Cid]] {.base.} =
   ## Get the list of blocks in the BlockStore. This is an intensive operation
   ##
 
diff --git a/dagger/stores/cachestore.nim b/dagger/stores/cachestore.nim
index 50cfaf9d..ddcd730f 100644
--- a/dagger/stores/cachestore.nim
+++ b/dagger/stores/cachestore.nim
@@ -68,9 +68,8 @@ method hasBlock*(self: CacheStore, cid: Cid): bool =
 
   cid in self.cache
 
-method listBlocks*(s: CacheStore, onBlock: OnBlock) {.async.} =
-  for cid in toSeq(s.cache.keys):
-    await onBlock(cid)
+method blockList*(s: CacheStore): Future[seq[Cid]] {.async.} =
+  return toSeq(s.cache.keys)
 
 func putBlockSync(self: CacheStore, blk: Block): bool =
 
diff --git a/dagger/stores/fsstore.nim b/dagger/stores/fsstore.nim
index b561250d..4870e3fd 100644
--- a/dagger/stores/fsstore.nim
+++ b/dagger/stores/fsstore.nim
@@ -90,7 +90,7 @@ method putBlock*(
     trace "Unable to store block", path, cid = blk.cid, error
     return false
 
-  if not (await self.cache.putBlock(blk)):
+  if await self.cache.putBlock(blk):
     trace "Unable to store block in cache", cid = blk.cid
 
   return true
@@ -113,8 +113,8 @@ method delBlock*(
     trace "Unable to delete block", path, cid, error
     return false
 
-  if not (await self.cache.delBlock(cid)):
-    trace "Unable to delete block from cache", cid
+  if await self.cache.delBlock(cid):
+    trace "Unable to store block in cache", cid
 
   return true
 
@@ -129,25 +129,21 @@ method hasBlock*(self: FSStore, cid: Cid): bool =
 
   self.blockPath(cid).isFile()
 
-method listBlocks*(self: FSStore, onBlock: OnBlock) {.async.} =
-  debug "Finding all blocks in store"
-  for (pkind, folderPath) in self.repoDir.walkDir():
+method blockList*(s: FSStore): Future[seq[Cid]] {.async.} =
+  ## Very expensive AND blocking!
+
+  debug "finding all blocks in store"
+  for (pkind, folderPath) in s.repoDir.walkDir():
     if pkind != pcDir: continue
     let baseName = basename(folderPath)
-    if baseName.len != self.postfixLen: continue
+    if baseName.len != s.postfixLen: continue
 
     for (fkind, filePath) in folderPath.walkDir(false):
       if fkind != pcFile: continue
       let cid = Cid.init(basename(filePath))
       if cid.isOk:
-        # getting a weird `Error: unhandled exception: index 1 not in 0 .. 0 [IndexError]`
-        # compilation error if using different syntax/construct bellow
-        try:
-          await onBlock(cid.get())
-        except CatchableError as exc:
-          trace "Couldn't get block", cid = $(cid.get())
-
-      await sleepAsync(100.millis) # avoid blocking
+        result.add(cid.get())
+  return result
 
 proc new*(
   T: type FSStore,
diff --git a/dagger/stores/networkstore.nim b/dagger/stores/networkstore.nim
index d9d27612..0d2320e0 100644
--- a/dagger/stores/networkstore.nim
+++ b/dagger/stores/networkstore.nim
@@ -42,10 +42,10 @@ method getBlock*(
   trace "Getting block", cid
   without var blk =? (await self.localStore.getBlock(cid)):
     trace "Couldn't get from local store", cid
-    try:
-      blk = await self.engine.requestBlock(cid)
+    blk = try:
+      await self.engine.requestBlock(cid)
     except CatchableError as exc:
-      trace "Exception requesting block", cid, exc = exc.msg
+      trace "Exception requestig block", cid, exc = exc.msg
       return failure(exc.msg)
 
   trace "Retrieved block from local store", cid
diff --git a/dagger/utils.nim b/dagger/utils.nim
deleted file mode 100644
index 2f481e62..00000000
--- a/dagger/utils.nim
+++ /dev/null
@@ -1,4 +0,0 @@
-import ./utils/asyncheapqueue
-import ./utils/fileutils
-
-export asyncheapqueue, fileutils
\ No newline at end of file
diff --git a/tests/config.nims b/tests/config.nims
new file mode 100644
index 00000000..c8062b36
--- /dev/null
+++ b/tests/config.nims
@@ -0,0 +1 @@
+patchFile("dagger", "discovery", "dagger/mockdiscovery")
diff --git a/tests/dagger/blockexc/discovery/testdiscovery.nim b/tests/dagger/blockexc/discovery/testdiscovery.nim
deleted file mode 100644
index bd596a56..00000000
--- a/tests/dagger/blockexc/discovery/testdiscovery.nim
+++ /dev/null
@@ -1,192 +0,0 @@
-import std/sequtils
-import std/sugar
-import std/algorithm
-import std/tables
-
-import pkg/asynctest
-import pkg/chronos
-import pkg/stew/byteutils
-
-import pkg/libp2p
-import pkg/libp2p/errors
-
-import pkg/dagger/rng
-import pkg/dagger/stores
-import pkg/dagger/blockexchange
-import pkg/dagger/chunker
-import pkg/dagger/blocktype as bt
-
-import ./mockdiscovery
-
-import ../../helpers
-import ../../examples
-
-suite "Block Advertising and Discovery":
-  let chunker = RandomChunker.new(Rng.instance(), size = 4096, chunkSize = 256)
-
-  var
-    switch: seq[Switch]
-    blockexc: seq[NetworkStore]
-    blocks: seq[bt.Block]
-
-  setup:
-    while true:
-      let chunk = await chunker.getBytes()
-      if chunk.len <= 0:
-        break
-
-      blocks.add(bt.Block.new(chunk).tryGet())
-
-  teardown:
-    switch = @[]
-    blockexc = @[]
-
-  test "Should discover want list":
-    let
-      s = newStandardSwitch(transportFlags = {ServerFlags.ReuseAddr})
-      discovery = MockDiscovery.new(s.peerInfo, 0.Port)
-      wallet = WalletRef.example
-      network = BlockExcNetwork.new(s)
-      localStore = CacheStore.new(blocks.mapIt( it ))
-      engine = BlockExcEngine.new(localStore, wallet, network, discovery)
-
-    s.mount(network)
-    switch.add(s)
-
-    await allFuturesThrowing(
-      switch.mapIt( it.start() )
-    )
-
-    let
-      pendingBlocks = blocks.mapIt(
-        engine.pendingBlocks.getWantHandle(it.cid)
-      )
-
-    await engine.start() # fire up discovery loop
-    discovery.findBlockProvidersHandler =
-      proc(d: MockDiscovery, cid: Cid): seq[SignedPeerRecord] =
-        engine.resolveBlocks(blocks.filterIt( it.cid == cid ))
-
-  test "Should advertise have blocks":
-    let
-      s = newStandardSwitch(transportFlags = {ServerFlags.ReuseAddr})
-      discovery = MockDiscovery.new(s.peerInfo, 0.Port)
-      wallet = WalletRef.example
-      network = BlockExcNetwork.new(s)
-      localStore = CacheStore.new(blocks.mapIt( it ))
-      engine = BlockExcEngine.new(localStore, wallet, network, discovery)
-
-    s.mount(network)
-    switch.add(s)
-
-    await allFuturesThrowing(
-      switch.mapIt( it.start() )
-    )
-
-    let
-      advertised = initTable.collect:
-        for b in blocks: {b.cid: newFuture[void]()}
-
-    discovery.publishProvideHandler = proc(d: MockDiscovery, cid: Cid) =
-      if cid in advertised and not advertised[cid].finished():
-        advertised[cid].complete()
-
-    await engine.start() # fire up advertise loop
-    await allFuturesThrowing(
-      allFinished(toSeq(advertised.values)))
-
-suite "E2E - Multiple Nodes Discovery":
-  let chunker = RandomChunker.new(Rng.instance(), size = 4096, chunkSize = 256)
-
-  var
-    switch: seq[Switch]
-    blockexc: seq[NetworkStore]
-    blocks: seq[bt.Block]
-
-  setup:
-    while true:
-      let chunk = await chunker.getBytes()
-      if chunk.len <= 0:
-        break
-
-      blocks.add(bt.Block.new(chunk).tryGet())
-
-    for _ in 0..<4:
-      let
-        s = newStandardSwitch(transportFlags = {ServerFlags.ReuseAddr})
-        discovery = MockDiscovery.new(s.peerInfo, 0.Port)
-        wallet = WalletRef.example
-        network = BlockExcNetwork.new(s)
-        localStore = CacheStore.new(blocks.mapIt( it ))
-        engine = BlockExcEngine.new(localStore, wallet, network, discovery)
-        networkStore = NetworkStore.new(engine, localStore)
-
-      s.mount(network)
-      switch.add(s)
-      blockexc.add(networkStore)
-
-  teardown:
-    switch = @[]
-    blockexc = @[]
-
-  test "Should not launch discovery request if we are already connected":
-    await allFuturesThrowing(
-      blockexc.mapIt( it.engine.start() ) &
-      switch.mapIt( it.start() )
-    )
-
-    await blockexc[0].engine.blocksHandler(switch[1].peerInfo.peerId, blocks)
-    MockDiscovery(blockexc[0].engine.discovery)
-      .findBlockProvidersHandler = proc(d: MockDiscovery, cid: Cid): seq[SignedPeerRecord] =
-        check false
-
-    await connectNodes(switch)
-    let blk = await blockexc[1].engine.requestBlock(blocks[0].cid)
-
-    await allFuturesThrowing(
-      blockexc.mapIt( it.engine.stop() ) &
-      switch.mapIt( it.stop() )
-    )
-
-  test "E2E - Should advertise and discover blocks":
-    # Distribute the blocks amongst 1..3
-    # Ask 0 to download everything without connecting him beforehand
-
-    var advertised: Table[Cid, SignedPeerRecord]
-
-    MockDiscovery(blockexc[1].engine.discovery)
-      .publishProvideHandler = proc(d: MockDiscovery, cid: Cid) =
-        advertised.add(cid, switch[1].peerInfo.signedPeerRecord)
-
-    MockDiscovery(blockexc[2].engine.discovery)
-      .publishProvideHandler = proc(d: MockDiscovery, cid: Cid) =
-        advertised.add(cid, switch[2].peerInfo.signedPeerRecord)
-
-    MockDiscovery(blockexc[3].engine.discovery)
-      .publishProvideHandler = proc(d: MockDiscovery, cid: Cid) =
-        advertised.add(cid, switch[3].peerInfo.signedPeerRecord)
-
-    await blockexc[1].engine.blocksHandler(switch[0].peerInfo.peerId, blocks[0..5])
-    await blockexc[2].engine.blocksHandler(switch[0].peerInfo.peerId, blocks[4..10])
-    await blockexc[3].engine.blocksHandler(switch[0].peerInfo.peerId, blocks[10..15])
-
-    MockDiscovery(blockexc[0].engine.discovery)
-      .findBlockProvidersHandler = proc(d: MockDiscovery, cid: Cid): seq[SignedPeerRecord] =
-        if cid in advertised:
-          result.add(advertised[cid])
-
-    let futs = collect(newSeq):
-      for b in blocks:
-        blockexc[0].engine.requestBlock(b.cid)
-
-    await allFuturesThrowing(
-      switch.mapIt( it.start() ) &
-      blockexc.mapIt( it.engine.start() )
-    )
-
-    await allFutures(futs)
-
-    await allFuturesThrowing(
-      blockexc.mapIt( it.engine.stop() ) &
-      switch.mapIt( it.stop() )
-    )
diff --git a/tests/dagger/blockexc/testblockexc.nim b/tests/dagger/blockexc/testblockexc.nim
index d80f643b..361c772f 100644
--- a/tests/dagger/blockexc/testblockexc.nim
+++ b/tests/dagger/blockexc/testblockexc.nim
@@ -1,4 +1,5 @@
 import std/sequtils
+import std/sugar
 import std/algorithm
 
 import pkg/asynctest
@@ -7,6 +8,7 @@ import pkg/stew/byteutils
 
 import pkg/libp2p
 import pkg/libp2p/errors
+import pkg/libp2pdht/discv5/protocol as discv5
 
 import pkg/dagger/rng
 import pkg/dagger/stores
@@ -36,7 +38,6 @@ suite "NetworkStore engine - 2 nodes":
     engine1, engine2: BlockExcEngine
     localStore1, localStore2: BlockStore
     discovery1, discovery2: Discovery
-    pendingBlocks1, pendingBlocks2: seq[Future[bt.Block]]
 
   setup:
     while true:
@@ -85,8 +86,8 @@ suite "NetworkStore engine - 2 nodes":
     )
 
     # initialize our want lists
-    pendingBlocks1 = blocks2.mapIt( blockexc1.engine.pendingBlocks.getWantHandle( it.cid ) )
-    pendingBlocks2 = blocks1.mapIt( blockexc2.engine.pendingBlocks.getWantHandle( it.cid ) )
+    for b in blocks2: discard blockexc1.engine.discoverBlock(b.cid)
+    for b in blocks1: discard blockexc2.engine.discoverBlock(b.cid)
 
     pricing1.address = wallet1.address
     pricing2.address = wallet2.address
@@ -97,7 +98,7 @@ suite "NetworkStore engine - 2 nodes":
       switch2.peerInfo.peerId,
       switch2.peerInfo.addrs)
 
-    await sleepAsync(1.seconds) # give some time to exchange lists
+    await sleepAsync(100.milliseconds) # give some time to exchange lists
     peerCtx2 = blockexc1.engine.getPeerCtx(peerId2)
     peerCtx1 = blockexc2.engine.getPeerCtx(peerId1)
 
@@ -112,18 +113,12 @@ suite "NetworkStore engine - 2 nodes":
     check not isNil(peerCtx1)
     check not isNil(peerCtx2)
 
-    await allFuturesThrowing(
-      allFinished(pendingBlocks1))
-
-    await allFuturesThrowing(
-      allFinished(pendingBlocks2))
-
     check:
       peerCtx1.peerHave.mapIt( $it ).sorted(cmp[string]) ==
-        pendingBlocks2.mapIt( $it.read.cid ).sorted(cmp[string])
+        toSeq(blockexc2.engine.runningDiscoveries.keys()).mapIt( $it ).sorted(cmp[string])
 
       peerCtx2.peerHave.mapIt( $it ).sorted(cmp[string]) ==
-        pendingBlocks1.mapIt( $it.read.cid ).sorted(cmp[string])
+        toSeq(blockexc1.engine.runningDiscoveries.keys()).mapIt( $it ).sorted(cmp[string])
 
   test "exchanges accounts on connect":
     check peerCtx1.account.?address == pricing1.address.some
@@ -180,8 +175,7 @@ suite "NetworkStore engine - 2 nodes":
     check wallet2.balance(channel, Asset) > 0
 
 suite "NetworkStore - multiple nodes":
-  let
-    chunker = RandomChunker.new(Rng.instance(), size = 4096, chunkSize = 256)
+  let chunker = RandomChunker.new(Rng.instance(), size = 4096, chunkSize = 256)
 
   var
     switch: seq[Switch]
@@ -219,9 +213,10 @@ suite "NetworkStore - multiple nodes":
       engine = downloader.engine
 
     # Add blocks from 1st peer to want list
-    let
-      pendingBlocks1 = blocks[0..3].mapIt( engine.pendingBlocks.getWantHandle( it.cid ) )
-      pendingBlocks2 = blocks[12..15].mapIt( engine.pendingBlocks.getWantHandle( it.cid ))
+    for b in blocks[0..3]:
+      discard engine.discoverBlock(b.cid)
+    for b in blocks[12..15]:
+      discard engine.discoverBlock(b.cid)
 
     await allFutures(
       blocks[0..3].mapIt( blockexc[0].engine.localStore.putBlock(it) ))
@@ -235,16 +230,12 @@ suite "NetworkStore - multiple nodes":
     await connectNodes(switch)
     await sleepAsync(1.seconds)
 
-    await allFuturesThrowing(
-      allFinished(pendingBlocks1),
-      allFinished(pendingBlocks2))
-
     check:
       engine.peers[0].peerHave.mapIt($it).sorted(cmp[string]) ==
-        blocks[0..3].mapIt( $(it.cid) ).sorted(cmp[string])
+        blocks[0..3].mapIt( it.cid ).mapIt($it).sorted(cmp[string])
 
       engine.peers[3].peerHave.mapIt($it).sorted(cmp[string]) ==
-        blocks[12..15].mapIt( $(it.cid) ).sorted(cmp[string])
+        blocks[12..15].mapIt( it.cid ).mapIt($it).sorted(cmp[string])
 
   test "should exchange blocks with multiple nodes":
     let
@@ -252,9 +243,10 @@ suite "NetworkStore - multiple nodes":
       engine = downloader.engine
 
     # Add blocks from 1st peer to want list
-    let
-      pendingBlocks1 = blocks[0..3].mapIt( engine.pendingBlocks.getWantHandle( it.cid ) )
-      pendingBlocks2 = blocks[12..15].mapIt( engine.pendingBlocks.getWantHandle( it.cid ))
+    for b in blocks[0..3]:
+      discard engine.discoverBlock(b.cid)
+    for b in blocks[12..15]:
+      discard engine.discoverBlock(b.cid)
 
     await allFutures(
       blocks[0..3].mapIt( blockexc[0].engine.localStore.putBlock(it) ))
@@ -268,9 +260,74 @@ suite "NetworkStore - multiple nodes":
     await connectNodes(switch)
     await sleepAsync(1.seconds)
 
-    await allFuturesThrowing(
-      allFinished(pendingBlocks1),
-      allFinished(pendingBlocks2))
+    let wantListBlocks = await allFinished(
+      blocks[0..3].mapIt( downloader.getBlock(it.cid) ))
+    check wantListBlocks.mapIt( !it.read ) == blocks[0..3]
 
-    check pendingBlocks1.mapIt( it.read ) == blocks[0..3]
-    check pendingBlocks2.mapIt( it.read ) == blocks[12..15]
+suite "NetworkStore - discovery":
+  let chunker = RandomChunker.new(Rng.instance(), size = 4096, chunkSize = 256)
+
+  var
+    switch: seq[Switch]
+    blockexc: seq[NetworkStore]
+    blocks: seq[bt.Block]
+
+  setup:
+    while true:
+      let chunk = await chunker.getBytes()
+      if chunk.len <= 0:
+        break
+
+      blocks.add(bt.Block.new(chunk).tryGet())
+
+    for e in generateNodes(4):
+      switch.add(e.switch)
+      blockexc.add(e.blockexc)
+      await e.blockexc.engine.start()
+
+    await allFuturesThrowing(
+      switch.mapIt( it.start() )
+    )
+
+  teardown:
+    await allFuturesThrowing(
+      switch.mapIt( it.stop() )
+    )
+
+    switch = @[]
+    blockexc = @[]
+
+  test "Shouldn't launch discovery request if we are already connected":
+    await blockexc[0].engine.blocksHandler(switch[1].peerInfo.peerId, blocks)
+    blockexc[0].engine.discovery.findBlockProviders_var = proc(d: Discovery, cid: Cid): seq[SignedPeerRecord] =
+      check false
+    await connectNodes(switch)
+    let blk = await blockexc[1].engine.requestBlock(blocks[0].cid)
+
+  test "E2E discovery":
+    # Distribute the blocks amongst 1..3
+    # Ask 0 to download everything without connecting him beforehand
+
+    var advertised: Table[Cid, SignedPeerRecord]
+
+    blockexc[1].engine.discovery.publishProvide_var = proc(d: Discovery, cid: Cid) =
+      advertised[cid] = switch[1].peerInfo.signedPeerRecord
+
+    blockexc[2].engine.discovery.publishProvide_var = proc(d: Discovery, cid: Cid) =
+      advertised[cid] = switch[2].peerInfo.signedPeerRecord
+
+    blockexc[3].engine.discovery.publishProvide_var = proc(d: Discovery, cid: Cid) =
+      advertised[cid] = switch[3].peerInfo.signedPeerRecord
+
+    await blockexc[1].engine.blocksHandler(switch[0].peerInfo.peerId, blocks[0..5])
+    await blockexc[2].engine.blocksHandler(switch[0].peerInfo.peerId, blocks[4..10])
+    await blockexc[3].engine.blocksHandler(switch[0].peerInfo.peerId, blocks[10..15])
+
+    blockexc[0].engine.discovery.findBlockProviders_var = proc(d: Discovery, cid: Cid): seq[SignedPeerRecord] =
+      if cid in advertised:
+        result.add(advertised[cid])
+
+    let futs = collect(newSeq):
+      for b in blocks:
+        blockexc[0].engine.requestBlock(b.cid)
+    await allFutures(futs)
diff --git a/tests/dagger/blockexc/testengine.nim b/tests/dagger/blockexc/testengine.nim
index d741a120..9940c9c6 100644
--- a/tests/dagger/blockexc/testengine.nim
+++ b/tests/dagger/blockexc/testengine.nim
@@ -66,9 +66,8 @@ suite "NetworkStore engine basic":
         wallet,
         network,
         discovery)
-
     for b in blocks:
-      discard engine.pendingBlocks.getWantHandle(b.cid)
+      discard engine.discoverBlock(b.cid)
     engine.setupPeer(peerId)
 
     await done
@@ -172,7 +171,7 @@ suite "NetworkStore engine handlers":
 
   test "stores blocks in local store":
     let pending = blocks.mapIt(
-      engine.pendingBlocks.getWantHandle( it.cid )
+      engine.pendingBlocks.addOrAwait( it.cid )
     )
 
     await engine.blocksHandler(peerId, blocks)
diff --git a/tests/dagger/helpers/nodeutils.nim b/tests/dagger/helpers/nodeutils.nim
index f190746f..e1ce5027 100644
--- a/tests/dagger/helpers/nodeutils.nim
+++ b/tests/dagger/helpers/nodeutils.nim
@@ -11,8 +11,10 @@ import ../examples
 
 proc generateNodes*(
   num: Natural,
-  blocks: openArray[bt.Block] = []):
-  seq[tuple[switch: Switch, blockexc: NetworkStore]] =
+  blocks: openArray[bt.Block] = [],
+  secureManagers: openarray[SecureProtocol] = [
+    SecureProtocol.Noise,
+  ]): seq[tuple[switch: Switch, blockexc: NetworkStore]] =
   for i in 0..<num:
     let
       switch = newStandardSwitch(transportFlags = {ServerFlags.ReuseAddr})
@@ -23,6 +25,8 @@ proc generateNodes*(
       engine = BlockExcEngine.new(localStore, wallet, network, discovery)
       networkStore = NetworkStore.new(engine, localStore)
 
+    switch.mount(network)
+
     switch.mount(network)
     result.add((switch, networkStore))
 
diff --git a/tests/dagger/blockexc/discovery/mockdiscovery.nim b/tests/dagger/mockdiscovery.nim
similarity index 60%
rename from tests/dagger/blockexc/discovery/mockdiscovery.nim
rename to tests/dagger/mockdiscovery.nim
index 137db9c8..6dd821cc 100644
--- a/tests/dagger/blockexc/discovery/mockdiscovery.nim
+++ b/tests/dagger/mockdiscovery.nim
@@ -12,15 +12,17 @@ import pkg/libp2p
 import pkg/questionable
 import pkg/questionable/results
 import pkg/stew/shims/net
-import pkg/dagger/discovery
+import pkg/libp2pdht/discv5/protocol as discv5
+
+export discv5
 
 type
-  MockDiscovery* = ref object of Discovery
-    findBlockProvidersHandler*: proc(d: MockDiscovery, cid: Cid): seq[SignedPeerRecord] {.gcsafe.}
-    publishProvideHandler*: proc(d: MockDiscovery, cid: Cid) {.gcsafe.}
+  Discovery* = ref object
+    findBlockProviders_var*: proc(d: Discovery, cid: Cid): seq[SignedPeerRecord] {.gcsafe.}
+    publishProvide_var*: proc(d: Discovery, cid: Cid) {.gcsafe.}
 
 proc new*(
-  T: type MockDiscovery,
+  T: type Discovery,
   localInfo: PeerInfo,
   discoveryPort: Port,
   bootstrapNodes = newSeq[SignedPeerRecord](),
@@ -33,16 +35,17 @@ proc findPeer*(
   peerId: PeerID): Future[?PeerRecord] {.async.} =
   return none(PeerRecord)
 
-method findBlockProviders*(
-  d: MockDiscovery,
+proc findBlockProviders*(
+  d: Discovery,
   cid: Cid): Future[seq[SignedPeerRecord]] {.async.} =
-  if isNil(d.findBlockProvidersHandler): return
+  if isNil(d.findBlockProviders_var): return
 
-  return d.findBlockProvidersHandler(d, cid)
+  return d.findBlockProviders_var(d, cid)
+
+proc publishProvide*(d: Discovery, cid: Cid) {.async.} =
+  if isNil(d.publishProvide_var): return
+  d.publishProvide_var(d, cid)
 
-method provideBlock*(d: MockDiscovery, cid: Cid) {.async.} =
-  if isNil(d.publishProvideHandler): return
-  d.publishProvideHandler(d, cid)
 
 proc start*(d: Discovery) {.async.} =
   discard
diff --git a/tests/dagger/stores/testcachestore.nim b/tests/dagger/stores/testcachestore.nim
index 3bd128dd..01cc18db 100644
--- a/tests/dagger/stores/testcachestore.nim
+++ b/tests/dagger/stores/testcachestore.nim
@@ -106,15 +106,3 @@ suite "Cache Store tests":
       await store.delBlock(newBlock2.cid)
       store.currentSize == 200
       newBlock2.cid notin store
-
-  test "listBlocks":
-    discard await store.putBlock(newBlock1)
-
-    var listed = false
-    await store.listBlocks(
-      proc(cid: Cid) {.gcsafe, async.} =
-        check cid in store
-        listed = true
-    )
-
-    check listed
diff --git a/tests/dagger/stores/testfsstore.nim b/tests/dagger/stores/testfsstore.nim
index 628cf42d..b2879591 100644
--- a/tests/dagger/stores/testfsstore.nim
+++ b/tests/dagger/stores/testfsstore.nim
@@ -52,13 +52,11 @@ suite "FS Store":
 
     check store.hasBlock(newBlock.cid)
 
-  test "listBlocks":
+  test "blockList":
     createDir(store.blockPath(newBlock.cid).parentDir)
     writeFile(store.blockPath(newBlock.cid), newBlock.data)
 
-    await store.listBlocks(
-      proc(cid: Cid) {.gcsafe, async.} =
-        check cid == newBlock.cid)
+    check (await store.blockList()) == @[newBlock.cid]
 
   test "fail hasBlock":
     check not store.hasBlock(newBlock.cid)
diff --git a/tests/dagger/testblockexc.nim b/tests/dagger/testblockexc.nim
index 6408c72f..8ba2e05b 100644
--- a/tests/dagger/testblockexc.nim
+++ b/tests/dagger/testblockexc.nim
@@ -4,6 +4,5 @@ import ./blockexc/protobuf/testpayments as testprotobufpayments
 import ./blockexc/protobuf/testpresence
 import ./blockexc/engine/testpayments as testenginepayments
 import ./blockexc/testblockexc
-import ./blockexc/discovery/testdiscovery
 
 {.warning[UnusedImport]: off.}