Move a lot of logic to the Batch module

Author: DZakh

codegenerator/cli/npm/envio/src/Batch.res

@@ -1,16 +1,24 @@
-type progressedChain = {
-  chainId: int,
+open Belt
+
+type chainAfterBatch = {
   batchSize: int,
   progressBlockNumber: int,
   totalEventsProcessed: int,
+  fetchState: FetchState.t,
+  dcsToStore: option<array<FetchState.indexingContract>>,
+  isProgressAtHeadWhenBatchCreated: bool,
+}
+
+type chainBeforeBatch = {
+  fetchState: FetchState.t,
+  progressBlockNumber: int,
+  sourceBlockNumber: int,
+  totalEventsProcessed: int,
 }
 
 type t = {
   items: array<Internal.item>,
-  progressedChains: array<progressedChain>,
-  updatedFetchStates: ChainMap.t<FetchState.t>,
-  dcsToStoreByChainId: dict<array<FetchState.indexingContract>>,
-  creationTimeMs: int,
+  progressedChainsById: dict<chainAfterBatch>,
 }
 
 /**
@@ -21,7 +29,7 @@ let getOrderedNextChain = (fetchStates: ChainMap.t<FetchState.t>, ~batchSizePerC
   let earliestChainTimestamp = ref(0)
   let chainKeys = fetchStates->ChainMap.keys
   for idx in 0 to chainKeys->Array.length - 1 {
-    let chain = chainKeys->Array.get(idx)
+    let chain = chainKeys->Array.getUnsafe(idx)
     let fetchState = fetchStates->ChainMap.get(chain)
     if fetchState->FetchState.isActivelyIndexing {
       let timestamp = fetchState->FetchState.getTimestampAt(
@@ -74,14 +82,116 @@ let hasMultichainReadyItem = (
   }
 }
 
-let prepareOrderedBatch = (
-  ~batchSizeTarget,
-  ~fetchStates: ChainMap.t<FetchState.t>,
-  ~mutBatchSizePerChain: dict<int>,
-) => {
+let getProgressedChainsById = {
+  let getChainAfterBatchIfProgressed = (
+    ~chainBeforeBatch: chainBeforeBatch,
+    ~updatedFetchState,
+    ~batchSize,
+    ~dcsToStore,
+  ) => {
+    let nextProgressBlockNumber = updatedFetchState->FetchState.getProgressBlockNumber
+
+    // The check is sufficient, since we guarantee to include a full block in a batch
+    // Also, this might be true even if batchSize is 0,
+    // eg when indexing at the head and chain doesn't have items in a block
+    if chainBeforeBatch.progressBlockNumber < nextProgressBlockNumber {
+      Some(
+        (
+          {
+            batchSize,
+            progressBlockNumber: nextProgressBlockNumber,
+            totalEventsProcessed: chainBeforeBatch.totalEventsProcessed + batchSize,
+            dcsToStore,
+            fetchState: updatedFetchState,
+            isProgressAtHeadWhenBatchCreated: nextProgressBlockNumber >=
+            chainBeforeBatch.sourceBlockNumber,
+          }: chainAfterBatch
+        ),
+      )
+    } else {
+      None
+    }
+  }
+
+  (~chainsBeforeBatch: ChainMap.t<chainBeforeBatch>, ~batchSizePerChain: dict<int>) => {
+    let progressedChainsById = Js.Dict.empty()
+
+    // Needed to:
+    // - Recalculate the computed queue sizes
+    // - Accumulate registered dynamic contracts to store in the db
+    // - Trigger onBlock pointer update
+    chainsBeforeBatch
+    ->ChainMap.values
+    ->Array.forEachU(chainBeforeBatch => {
+      let fetchState = chainBeforeBatch.fetchState
+      switch switch batchSizePerChain->Utils.Dict.dangerouslyGetNonOption(
+        fetchState.chainId->Int.toString,
+      ) {
+      | Some(batchSize) =>
+        let leftItems = fetchState.buffer->Js.Array2.sliceFrom(batchSize)
+        switch fetchState.dcsToStore {
+        | [] =>
+          getChainAfterBatchIfProgressed(
+            ~chainBeforeBatch,
+            ~batchSize,
+            ~dcsToStore=None,
+            ~updatedFetchState=fetchState->FetchState.updateInternal(~mutItems=leftItems),
+          )
+
+        | dcs => {
+            let leftDcsToStore = []
+            let batchDcs = []
+            let updatedFetchState =
+              fetchState->FetchState.updateInternal(~mutItems=leftItems, ~dcsToStore=leftDcsToStore)
+            let nextProgressBlockNumber = updatedFetchState->FetchState.getProgressBlockNumber
+
+            dcs->Array.forEach(dc => {
+              // Important: This should be a registering block number.
+              // This works for now since dc.startBlock is a registering block number.
+              if dc.startBlock <= nextProgressBlockNumber {
+                batchDcs->Array.push(dc)
+              } else {
+                // Mutate the array we passed to the updateInternal beforehand
+                leftDcsToStore->Array.push(dc)
+              }
+            })
+
+            getChainAfterBatchIfProgressed(
+              ~chainBeforeBatch,
+              ~batchSize,
+              ~dcsToStore=Some(batchDcs),
+              ~updatedFetchState,
+            )
+          }
+        }
+      // Skip not affected chains
+      | None =>
+        getChainAfterBatchIfProgressed(
+          ~chainBeforeBatch,
+          ~batchSize=0,
+          ~dcsToStore=None,
+          ~updatedFetchState=chainBeforeBatch.fetchState,
+        )
+      } {
+      | Some(progressedChain) =>
+        progressedChainsById->Utils.Dict.setByInt(
+          chainBeforeBatch.fetchState.chainId,
+          progressedChain,
+        )
+      | None => ()
+      }
+    })
+
+    progressedChainsById
+  }
+}
+
+let prepareOrderedBatch = (~chainsBeforeBatch: ChainMap.t<chainBeforeBatch>, ~batchSizeTarget) => {
   let batchSize = ref(0)
   let isFinished = ref(false)
+  let mutBatchSizePerChain = Js.Dict.empty()
   let items = []
+  let fetchStates = chainsBeforeBatch->ChainMap.map(chainBeforeBatch => chainBeforeBatch.fetchState)
 
   while batchSize.contents < batchSizeTarget && !isFinished.contents {
     switch fetchStates->getOrderedNextChain(~batchSizePerChain=mutBatchSizePerChain) {
@@ -116,23 +226,30 @@ let prepareOrderedBatch = (
     }
   }
 
-  items
+  {
+    items,
+    progressedChainsById: getProgressedChainsById(
+      ~chainsBeforeBatch,
+      ~batchSizePerChain=mutBatchSizePerChain,
+    ),
+  }
 }
 
 let prepareUnorderedBatch = (
+  ~chainsBeforeBatch: ChainMap.t<chainBeforeBatch>,
   ~batchSizeTarget,
-  ~fetchStates: ChainMap.t<FetchState.t>,
-  ~mutBatchSizePerChain: dict<int>,
 ) => {
   let preparedFetchStates =
-    fetchStates
+    chainsBeforeBatch
     ->ChainMap.values
+    ->Js.Array2.map(chainBeforeBatch => chainBeforeBatch.fetchState)
     ->FetchState.filterAndSortForUnorderedBatch(~batchSizeTarget)
 
   let chainIdx = ref(0)
   let preparedNumber = preparedFetchStates->Array.length
   let batchSize = ref(0)
 
+  let mutBatchSizePerChain = Js.Dict.empty()
   let items = []
 
   // Accumulate items for all actively indexing chains
@@ -156,5 +273,28 @@ let prepareUnorderedBatch = (
     chainIdx := chainIdx.contents + 1
   }
 
-  items
-}
+  {
+    items,
+    progressedChainsById: getProgressedChainsById(
+      ~chainsBeforeBatch,
+      ~batchSizePerChain=mutBatchSizePerChain,
+    ),
+  }
+}
+
+let make = (
+  ~chainsBeforeBatch: ChainMap.t<chainBeforeBatch>,
+  ~multichain: InternalConfig.multichain,
+  ~batchSizeTarget,
+) => {
+  if (
+    switch multichain {
+    | Unordered => true
+    | Ordered => chainsBeforeBatch->ChainMap.size === 1
+    }
+  ) {
+    prepareUnorderedBatch(~chainsBeforeBatch, ~batchSizeTarget)
+  } else {
+    prepareOrderedBatch(~chainsBeforeBatch, ~batchSizeTarget)
+  }
+}

codegenerator/cli/npm/envio/src/FetchState.res

@@ -92,15 +92,15 @@ let mergeIntoPartition = (p: partition, ~target: partition, ~maxAddrInPartition)
 
       let allowedAddressesNumber = ref(maxAddrInPartition)
 
-      target.addressesByContractName->Utils.Dict.forEachWithKey((contractName, addresses) => {
+      target.addressesByContractName->Utils.Dict.forEachWithKey((addresses, contractName) => {
         allowedAddressesNumber := allowedAddressesNumber.contents - addresses->Array.length
         mergedAddresses->Js.Dict.set(contractName, addresses)
       })
 
       // Start with putting all addresses to the merging dict
       // And if they exceed the limit, start removing from the merging dict
       // and putting into the rest dict
-      p.addressesByContractName->Utils.Dict.forEachWithKey((contractName, addresses) => {
+      p.addressesByContractName->Utils.Dict.forEachWithKey((addresses, contractName) => {
         allowedAddressesNumber := allowedAddressesNumber.contents - addresses->Array.length
         switch mergedAddresses->Utils.Dict.dangerouslyGetNonOption(contractName) {
         | Some(targetAddresses) =>
@@ -112,7 +112,7 @@ let mergeIntoPartition = (p: partition, ~target: partition, ~maxAddrInPartition)
       let rest = if allowedAddressesNumber.contents < 0 {
         let restAddresses = Js.Dict.empty()
 
-        mergedAddresses->Utils.Dict.forEachWithKey((contractName, addresses) => {
+        mergedAddresses->Utils.Dict.forEachWithKey((addresses, contractName) => {
           if allowedAddressesNumber.contents === 0 {
             ()
           } else if addresses->Array.length <= -allowedAddressesNumber.contents {
@@ -1153,7 +1153,7 @@ let rollbackPartition = (
     })
   | {addressesByContractName} =>
     let rollbackedAddressesByContractName = Js.Dict.empty()
-    addressesByContractName->Utils.Dict.forEachWithKey((contractName, addresses) => {
+    addressesByContractName->Utils.Dict.forEachWithKey((addresses, contractName) => {
       let keptAddresses =
         addresses->Array.keep(address => !(addressesToRemove->Utils.Set.has(address)))
       if keptAddresses->Array.length > 0 {

codegenerator/cli/npm/envio/src/PgStorage.res

@@ -325,7 +325,7 @@ let chunkArray = (arr: array<'a>, ~chunkSize) => {
 let removeInvalidUtf8InPlace = entities =>
   entities->Js.Array2.forEach(item => {
     let dict = item->(Utils.magic: 'a => dict<unknown>)
-    dict->Utils.Dict.forEachWithKey((key, value) => {
+    dict->Utils.Dict.forEachWithKey((value, key) => {
       if value->Js.typeof === "string" {
         let value = value->(Utils.magic: unknown => string)
         // We mutate here, since we don't care

codegenerator/cli/npm/envio/src/Utils.res

@@ -97,30 +97,66 @@ module Dict = {
   @val
   external mergeInPlace: (dict<'a>, dict<'a>) => dict<'a> = "Object.assign"
 
-  let map = (dict, fn) => {
-    let newDict = Js.Dict.empty()
-    let keys = dict->Js.Dict.keys
-    for idx in 0 to keys->Js.Array2.length - 1 {
-      let key = keys->Js.Array2.unsafe_get(idx)
-      newDict->Js.Dict.set(key, fn(dict->Js.Dict.unsafeGet(key)))
+  // Use %raw to support for..in which is a ~10% faster than .forEach
+  let mapValues: (dict<'a>, 'a => 'b) => dict<'b> = %raw(`(dict, f) => {
+    var target = {}, i;
+    for (i in dict) {
+      target[i] = f(dict[i]);
     }
-    newDict
-  }
+    return target;
+  }`)
 
-  let forEach = (dict, fn) => {
-    let keys = dict->Js.Dict.keys
-    for idx in 0 to keys->Js.Array2.length - 1 {
-      fn(dict->Js.Dict.unsafeGet(keys->Js.Array2.unsafe_get(idx)))
+  // Use %raw to support for..in which is a ~10% faster than .forEach
+  let filterMapValues: (dict<'a>, 'a => option<'b>) => dict<'b> = %raw(`(dict, f) => {
+    var target = {}, i, v;
+    for (i in dict) {
+      v = f(dict[i]);
+      if (v !== undefined) {
+        target[i] = v;
+      }
     }
-  }
+    return target;
+  }`)
 
-  let forEachWithKey = (dict, fn) => {
-    let keys = dict->Js.Dict.keys
-    for idx in 0 to keys->Js.Array2.length - 1 {
-      let key = keys->Js.Array2.unsafe_get(idx)
-      fn(key, dict->Js.Dict.unsafeGet(key))
+  // Use %raw to support for..in which is a ~10% faster than .forEach
+  let mapValuesToArray: (dict<'a>, 'a => 'b) => array<'b> = %raw(`(dict, f) => {
+    var target = [], i;
+    for (i in dict) {
+      target.push(f(dict[i]));
     }
-  }
+    return target;
+  }`)
+
+  // Use %raw to support for..in which is a ~10% faster than .forEach
+  let forEach: (dict<'a>, 'a => unit) => unit = %raw(`(dict, f) => {
+    for (var i in dict) {
+      f(dict[i]);
+    }
+  }`)
+
+  // Use %raw to support for..in which is a ~10% faster than .forEach
+  let forEachWithKey: (dict<'a>, ('a, string) => unit) => unit = %raw(`(dict, f) => {
+    for (var i in dict) {
+      f(dict[i], i);
+    }
+  }`)
+
+  // Use %raw to support for..in which is a ~10% faster than Object.keys
+  let size: dict<'a> => int = %raw(`(dict) => {
+    var size = 0, i;
+    for (i in dict) {
+      size++;
+    }
+    return size;
+  }`)
+
+  // Use %raw to support for..in which is a 2x faster than Object.keys
+  let isEmpty: dict<'a> => bool = %raw(`(dict) => {
+    for (var _ in dict) {
+      return false
+    }
+    return true
+  }`)
 
   let deleteInPlace: (dict<'a>, string) => unit = %raw(`(dict, key) => {
       delete dict[key];
@@ -135,8 +171,6 @@ module Dict = {
 
   let shallowCopy: dict<'a> => dict<'a> = %raw(`(dict) => ({...dict})`)
 
-  let size = dict => dict->Js.Dict.keys->Js.Array2.length
-
   @set_index
   external setByInt: (dict<'a>, int, 'a) => unit = ""