Move maxPartitionConcurrency check after all the filtering is done

Author: DZakh

codegenerator/cli/templates/static/codegen/src/eventFetching/ChainManager.res

@@ -1,4 +1,5 @@
 open Belt
+
 type t = {
   chainFetchers: ChainMap.t<ChainFetcher.t>,
   //Holds arbitrary events that were added when a batch ended processing early

codegenerator/cli/templates/static/codegen/src/eventFetching/PartitionedFetchState.res

@@ -184,14 +184,9 @@ at the max queue size.
 */
 let getMostBehindPartitions = (
   {partitions}: t,
-  ~maxNumQueries,
   ~maxPerChainQueueSize,
   ~partitionsCurrentlyFetching,
 ) => {
-  let maxNumQueries = Pervasives.max(
-    maxNumQueries - partitionsCurrentlyFetching->Belt.Set.Int.size,
-    0,
-  )
   let numPartitions = partitions->Array.length
   let maxPartitionQueueSize = maxPerChainQueueSize / numPartitions
 
@@ -211,7 +206,6 @@ let getMostBehindPartitions = (
     FetchState.getLatestFullyFetchedBlock(a.fetchState).blockNumber -
     FetchState.getLatestFullyFetchedBlock(b.fetchState).blockNumber
   )
-  ->Js.Array.slice(~start=0, ~end_=maxNumQueries)
 }
 
 type nextQueries = WaitForNewBlock | NextQuery(array<FetchState.nextQuery>)
@@ -225,7 +219,6 @@ let getNextQueries = (self: t, ~maxPerChainQueueSize, ~partitionsCurrentlyFetchi
 
   self
   ->getMostBehindPartitions(
-    ~maxNumQueries=Env.maxPartitionConcurrency,
     ~maxPerChainQueueSize,
     ~partitionsCurrentlyFetching,
   )

codegenerator/cli/templates/static/codegen/src/globalState/GlobalState.res

@@ -857,7 +857,7 @@ let checkAndFetchForChain = (
 ) => async chain => {
   let chainFetcher = state.chainManager.chainFetchers->ChainMap.get(chain)
   if !isRollingBack(state) {
-    let {chainConfig: {chainWorker}, logger, currentBlockHeight} = chainFetcher
+    let {chainConfig: {chainWorker}, logger, currentBlockHeight, partitionsCurrentlyFetching} = chainFetcher
 
     let (nextQueries, updatedPartitions) =
       chainFetcher->ChainFetcher.getNextQueries(~maxPerChainQueueSize=state.maxPerChainQueueSize)
@@ -895,6 +895,11 @@ let checkAndFetchForChain = (
         dispatchAction(FinishWaitingForNewBlock({chain, currentBlockHeight}))
       }
     | (queries, _) =>
+      let maxNumQueries = Pervasives.max(
+        Env.maxPartitionConcurrency - partitionsCurrentlyFetching->Belt.Set.Int.size,
+        0,
+      )
+      let queries = queries->Js.Array2.slice(~start=0, ~end_=maxNumQueries) 
       let newPartitionsCurrentlyFetching =
         queries->Array.map(query => query.partitionId)->Set.Int.fromArray
       dispatchAction(SetCurrentlyFetchingBatch(chain, newPartitionsCurrentlyFetching))

scenarios/test_codegen/test/lib_tests/PartitionedFetchState_test.res

@@ -52,56 +52,47 @@ describe("PartitionedFetchState getMostBehindPartitions", () => {
     mockFetchState(~latestFetchedBlockNumber=3),
   }
   let partitionedFetchState = mockPartitionedFetchState(~partitions)
-  it(
-    "With multiple partitions always returns the most behind partitions up to the max concurrency level",
-    () => {
-      let maxNumQueries = 3
-      let partitionsCurrentlyFetching = Set.Int.empty
+  it("With multiple partitions always returns the most behind partitions first", () => {
+    let partitionsCurrentlyFetching = Set.Int.empty
 
-      let mostBehindPartitions =
-        partitionedFetchState->PartitionedFetchState.getMostBehindPartitions(
-          ~maxNumQueries,
-          ~maxPerChainQueueSize=10,
-          ~partitionsCurrentlyFetching,
-        )
-
-      Assert.equal(mostBehindPartitions->Array.length, maxNumQueries)
-
-      let partitionIds = mostBehindPartitions->Array.map(p => p.partitionId)
-      Assert.deepEqual(
-        partitionIds,
-        [2, 3, 4],
-        ~message="Should have returned the partitions with the lowest latestFetchedBlock",
+    let mostBehindPartitions =
+      partitionedFetchState->PartitionedFetchState.getMostBehindPartitions(
+        ~maxPerChainQueueSize=10,
+        ~partitionsCurrentlyFetching,
       )
-    },
-  )
+
+    let partitionIds = mostBehindPartitions->Array.map(p => p.partitionId)
+    Assert.deepEqual(
+      partitionIds,
+      [2, 3, 4, 0, 1],
+      ~message="Should have returned the partitions with the lowest latestFetchedBlock first",
+    )
+  })
 
   it("Will not return partitions that are currently fetching", () => {
-    let maxNumQueries = 3
     let partitionsCurrentlyFetching = Set.Int.fromArray([2, 3])
 
     let mostBehindPartitions =
       partitionedFetchState->PartitionedFetchState.getMostBehindPartitions(
-        ~maxNumQueries,
         ~maxPerChainQueueSize=10,
         ~partitionsCurrentlyFetching,
       )
 
     Assert.equal(
       mostBehindPartitions->Array.length,
-      maxNumQueries - partitionsCurrentlyFetching->Set.Int.size,
+      partitionedFetchState.partitions->Js.Array2.length -
+        partitionsCurrentlyFetching->Set.Int.size,
     )
 
     let partitionIds = mostBehindPartitions->Array.map(p => p.partitionId)
     Assert.deepEqual(
       partitionIds,
-      [4],
+      [4, 0, 1],
       ~message="Should have returned the partitions with the lowest latestFetchedBlock that are not currently fetching",
     )
   })
 
-  it("Should not return partition that is at max partition size", () => {
-    let maxNumQueries = 3
+  it("Should not return partition that is at max queue size", () => {
     let partitions = list{
       mockFetchState(~latestFetchedBlockNumber=4),
       mockFetchState(~latestFetchedBlockNumber=5),
@@ -119,7 +110,6 @@ describe("PartitionedFetchState getMostBehindPartitions", () => {
 
     let mostBehindPartitions =
       partitionedFetchState->PartitionedFetchState.getMostBehindPartitions(
-        ~maxNumQueries,
         ~maxPerChainQueueSize=10, //each partition should therefore have a max of 2 events
         ~partitionsCurrentlyFetching=Set.Int.empty,
       )
@@ -133,7 +123,6 @@ describe("PartitionedFetchState getMostBehindPartitions", () => {
   })
 
   it("if need be should return less than maxNum queries if all partitions at their max", () => {
-    let maxNumQueries = 3
     let partitions = list{
       mockFetchState(~latestFetchedBlockNumber=4),
       mockFetchState(~latestFetchedBlockNumber=5),
@@ -154,7 +143,6 @@ describe("PartitionedFetchState getMostBehindPartitions", () => {
 
     let mostBehindPartitions =
       partitionedFetchState->PartitionedFetchState.getMostBehindPartitions(
-        ~maxNumQueries,
         ~maxPerChainQueueSize=10, //each partition should therefore have a max of 2 events
         ~partitionsCurrentlyFetching=Set.Int.empty,
       )