uses SafeAsyncIter in "listBlocks" and in "getBlockExpirations"

Author: marcinczenko

codex/blockexchange/engine/advertiser.nim

@@ -7,6 +7,8 @@
 ## This file may not be copied, modified, or distributed except according to
 ## those terms.
 
+{.push raises: [].}
+
 import pkg/chronos
 import pkg/libp2p/cid
 import pkg/libp2p/multicodec
@@ -81,16 +83,12 @@ proc advertiseBlock(b: Advertiser, cid: Cid) {.async: (raises: [CancelledError])
 proc advertiseLocalStoreLoop(b: Advertiser) {.async: (raises: []).} =
   try:
     while b.advertiserRunning:
-      try:
-        if cids =? await b.localStore.listBlocks(blockType = BlockType.Manifest):
-          trace "Advertiser begins iterating blocks..."
-          for c in cids:
-            if cid =? await c:
-              await b.advertiseBlock(cid)
-          trace "Advertiser iterating blocks finished."
-      except CatchableError as e:
-        error "Error in advertise local store loop", error = e.msgDetail
-        raiseAssert("Unexpected exception in advertiseLocalStoreLoop")
+      if cidsIter =? await b.localStore.listBlocks(blockType = BlockType.Manifest):
+        trace "Advertiser begins iterating blocks..."
+        for c in cidsIter:
+          if cid =? (await cast[Future[?!Cid].Raising([CancelledError])](c)):
+            await b.advertiseBlock(cid)
+        trace "Advertiser iterating blocks finished."
 
       await sleepAsync(b.advertiseLocalStoreLoopSleep)
   except CancelledError:

codex/node.nim

@@ -480,12 +480,12 @@ proc store*(
   return manifestBlk.cid.success
 
 proc iterateManifests*(self: CodexNodeRef, onManifest: OnManifest) {.async.} =
-  without cids =? await self.networkStore.listBlocks(BlockType.Manifest):
+  without cidsIter =? await self.networkStore.listBlocks(BlockType.Manifest):
     warn "Failed to listBlocks"
     return
 
-  for c in cids:
-    if cid =? await c:
+  for c in cidsIter:
+    if cid =? (await cast[Future[?!Cid].Raising([CancelledError])](c)):
       without blk =? await self.networkStore.getBlock(cid):
         warn "Failed to get manifest block by cid", cid
         return

codex/stores/blockstore.nim

@@ -7,10 +7,7 @@
 ## This file may not be copied, modified, or distributed except according to
 ## those terms.
 
-import pkg/upraises
-
-push:
-  {.upraises: [].}
+{.push raises: [].}
 
 import pkg/chronos
 import pkg/libp2p
@@ -152,7 +149,7 @@ method hasBlock*(
 
 method listBlocks*(
     self: BlockStore, blockType = BlockType.Manifest
-): Future[?!AsyncIter[?Cid]] {.base, async: (raises: [CancelledError]), gcsafe.} =
+): Future[?!SafeAsyncIter[Cid]] {.base, async: (raises: [CancelledError]), gcsafe.} =
   ## Get the list of blocks in the BlockStore. This is an intensive operation
   ##
 

codex/stores/cachestore.nim

@@ -7,10 +7,7 @@
 ## This file may not be copied, modified, or distributed except according to
 ## those terms.
 
-import pkg/upraises
-
-push:
-  {.upraises: [].}
+{.push raises: [].}
 
 import std/options
 
@@ -142,7 +139,7 @@ func cids(self: CacheStore): (iterator (): Cid {.gcsafe.}) =
 
 method listBlocks*(
     self: CacheStore, blockType = BlockType.Manifest
-): Future[?!AsyncIter[?Cid]] {.async: (raises: [CancelledError]).} =
+): Future[?!SafeAsyncIter[Cid]] {.async: (raises: [CancelledError]).} =
   ## Get the list of blocks in the BlockStore. This is an intensive operation
   ##
 
@@ -151,36 +148,33 @@ method listBlocks*(
   proc isFinished(): bool =
     return finished(cids)
 
-  proc genNext(): Future[Cid] {.async.} =
-    cids()
-
-  try:
-    let iter = await (
-      AsyncIter[Cid].new(genNext, isFinished).filter(
-        proc(cid: Cid): Future[bool] {.async.} =
-          without isManifest =? cid.isManifest, err:
-            trace "Error checking if cid is a manifest", err = err.msg
-            return false
-
-          case blockType
-          of BlockType.Both:
-            return true
-          of BlockType.Manifest:
-            return isManifest
-          of BlockType.Block:
-            return not isManifest
-      )
-    )
-
-    return success(
-      map[Cid, ?Cid](
-        iter,
-        proc(cid: Cid): Future[?Cid] {.async.} =
-          some(cid),
-      )
+  proc genNext(): Future[?!Cid] {.async: (raises: [CancelledError]).} =
+    try:
+      let cid = cids()
+      success(cid)
+    except Exception as err:
+      failure(err.msg)
+
+  let iter = await (
+    SafeAsyncIter[Cid].new(genNext, isFinished).filter(
+      proc(cid: ?!Cid): Future[bool] {.async: (raises: [CancelledError]).} =
+        without cid =? cid, err:
+          trace "Cannot get Cid from the iterator", err = err.msg
+          return false
+        without isManifest =? cid.isManifest, err:
+          trace "Error checking if cid is a manifest", err = err.msg
+          return false
+
+        case blockType
+        of BlockType.Both:
+          return true
+        of BlockType.Manifest:
+          return isManifest
+        of BlockType.Block:
+          return not isManifest
     )
-  except CatchableError as err:
-    raiseAssert err.msg
+  )
+  success(iter)
 
 func putBlockSync(self: CacheStore, blk: Block): bool =
   let blkSize = blk.data.len.NBytes # in bytes

codex/stores/maintenance.nim

@@ -10,13 +10,15 @@
 ## Store maintenance module
 ## Looks for and removes expired blocks from blockstores.
 
+{.push raises: [].}
+
 import pkg/chronos
 import pkg/questionable
 import pkg/questionable/results
 
 import ./repostore
 import ../utils/timer
-import ../utils/asynciter
+import ../utils/safeasynciter
 import ../clock
 import ../logutils
 import ../systemclock
@@ -81,29 +83,28 @@ proc runBlockCheck(
 
   var numberReceived = 0
   for beFut in iter:
-    try:
-      let be = await beFut
-      inc numberReceived
-      await self.processBlockExpiration(be)
-      await sleepAsync(1.millis) # cooperative scheduling
-    except CatchableError as err:
-      raiseAssert err.msg
+    let beRes = await cast[Future[?!BlockExpiration].Raising([CancelledError])](beFut)
+    without be =? beRes, err:
+      trace "Unable to obtain blockExpiration from iterator"
+      continue
+    inc numberReceived
+    await self.processBlockExpiration(be)
+    await sleepAsync(1.millis) # cooperative scheduling
 
   # If we received fewer blockExpirations from the iterator than we asked for,
   # We're at the end of the dataset and should start from 0 next time.
   if numberReceived < self.numberOfBlocksPerInterval:
     self.offset = 0
 
 proc start*(self: BlockMaintainer) =
-  proc onTimer(): Future[void] {.async.} =
+  proc onTimer(): Future[void] {.async: (raises: []).} =
     try:
       await self.runBlockCheck()
-    except CancelledError as error:
-      raise error
-    except CatchableError as exc:
-      error "Unexpected exception in BlockMaintainer.onTimer(): ", msg = exc.msg
+    except CancelledError as err:
+      trace "Running block check in block maintenance timer callback cancelled: ",
+        err = err.msg
 
   self.timer.start(onTimer, self.interval)
 
-proc stop*(self: BlockMaintainer): Future[void] {.async.} =
+proc stop*(self: BlockMaintainer): Future[void] {.async: (raises: []).} =
   await self.timer.stop()

codex/stores/networkstore.nim

@@ -19,7 +19,7 @@ import ../blockexchange
 import ../logutils
 import ../merkletree
 import ../utils/asyncheapqueue
-import ../utils/asynciter
+import ../utils/safeasynciter
 import ./blockstore
 
 export blockstore, blockexchange, asyncheapqueue
@@ -124,7 +124,7 @@ method ensureExpiry*(
 
 method listBlocks*(
     self: NetworkStore, blockType = BlockType.Manifest
-): Future[?!AsyncIter[?Cid]] {.async: (raw: true, raises: [CancelledError]).} =
+): Future[?!SafeAsyncIter[Cid]] {.async: (raw: true, raises: [CancelledError]).} =
   self.localStore.listBlocks(blockType)
 
 method delBlock*(

codex/stores/queryiterhelper.nim

@@ -5,6 +5,9 @@ import pkg/chronicles
 import pkg/datastore/typedds
 
 import ../utils/asynciter
+import ../utils/safeasynciter
+
+{.push raises: [].}
 
 type KeyVal*[T] = tuple[key: Key, value: T]
 
@@ -42,6 +45,40 @@ proc toAsyncIter*[T](
 
   AsyncIter[?!QueryResponse[T]].new(genNext, isFinished).success
 
+proc toSafeAsyncIter*[T](
+    queryIter: QueryIter[T], finishOnErr: bool = true
+): Future[?!SafeAsyncIter[QueryResponse[T]]] {.async: (raises: [CancelledError]).} =
+  ## Converts `QueryIter[T]` to `SafeAsyncIter[QueryResponse[T]]` and automatically
+  ## runs dispose whenever `QueryIter` finishes or whenever an error occurs (only
+  ## if the flag finishOnErr is set to true)
+  ##
+
+  if queryIter.finished:
+    trace "Disposing iterator"
+    if error =? (await queryIter.dispose()).errorOption:
+      return failure(error)
+    return success(SafeAsyncIter[QueryResponse[T]].empty())
+
+  var errOccurred = false
+
+  proc genNext(): Future[?!QueryResponse[T]] {.async: (raises: [CancelledError]).} =
+    let queryResOrErr = await queryIter.next()
+
+    if queryResOrErr.isErr:
+      errOccurred = true
+
+    if queryIter.finished or (errOccurred and finishOnErr):
+      trace "Disposing iterator"
+      if error =? (await queryIter.dispose()).errorOption:
+        return failure(error)
+
+    return queryResOrErr
+
+  proc isFinished(): bool =
+    queryIter.finished
+
+  SafeAsyncIter[QueryResponse[T]].new(genNext, isFinished).success
+
 proc filterSuccess*[T](
     iter: AsyncIter[?!QueryResponse[T]]
 ): Future[AsyncIter[tuple[key: Key, value: T]]] {.async: (raises: [CancelledError]).} =
@@ -62,7 +99,30 @@ proc filterSuccess*[T](
 
     (key: key, value: value).some
 
-  try:
-    await mapFilter[?!QueryResponse[T], KeyVal[T]](iter, mapping)
-  except CatchableError as err:
-    raiseAssert err.msg
+  await mapFilter[?!QueryResponse[T], KeyVal[T]](iter, mapping)
+
+proc filterSuccess*[T](
+    iter: SafeAsyncIter[QueryResponse[T]]
+): Future[SafeAsyncIter[tuple[key: Key, value: T]]] {.
+    async: (raises: [CancelledError])
+.} =
+  ## Filters out any items that are not success
+
+  proc mapping(
+      resOrErr: ?!QueryResponse[T]
+  ): Future[Option[?!KeyVal[T]]] {.async: (raises: [CancelledError]).} =
+    without res =? resOrErr, error:
+      error "Error occurred when getting QueryResponse", msg = error.msg
+      return Result[KeyVal[T], ref CatchableError].none
+
+    without key =? res.key:
+      warn "No key for a QueryResponse"
+      return Result[KeyVal[T], ref CatchableError].none
+
+    without value =? res.value, error:
+      error "Error occurred when getting a value from QueryResponse", msg = error.msg
+      return Result[KeyVal[T], ref CatchableError].none
+
+    some(success((key: key, value: value)))
+
+  await mapFilter[QueryResponse[T], KeyVal[T]](iter, mapping)