envio 3.3.0-alpha.8 → 3.3.0-rc.0

package/src/PgStorage.res.mjs

@@ -493,15 +493,18 @@ function makeSchemaTableNamesQuery(pgSchema) {
   return `SELECT table_name FROM information_schema.tables WHERE table_schema = '` + pgSchema + `';`;
 }
 
-let cacheTablePrefixLength = Internal.cacheTablePrefix.length;
-
 function makeSchemaCacheTableInfoQuery(pgSchema) {
-  return `SELECT 
+  return `SELECT
     t.table_name,
     ` + getCacheRowCountFnName + `(t.table_name) as count
    FROM information_schema.tables t
-   WHERE t.table_schema = '` + pgSchema + `' 
-   AND t.table_name LIKE '` + Internal.cacheTablePrefix + `%';`;
+   WHERE t.table_schema = '` + pgSchema + `'
+   AND t.table_name ~ '^envio_([0-9]+_)?effect_.+'
+   AND (
+     SELECT array_agg(c.column_name::text ORDER BY c.column_name::text)
+     FROM information_schema.columns c
+     WHERE c.table_schema = t.table_schema AND c.table_name = t.table_name
+   ) = ARRAY['id', 'output'];`;
 }
 
 let psqlExecState = {
@@ -766,7 +769,7 @@ async function writeBatch(sql, batch, pgSchema, rollback, isInReorgThreshold, co
           }
           throw exn;
         }),
-        Promise.all(updatedEffectsCache.map(param => setEffectCacheOrThrow(param.effect, param.items, param.shouldInitialize)))
+        Promise.all(updatedEffectsCache.map(param => setEffectCacheOrThrow(param.table, param.itemSchema, param.items, param.shouldInitialize)))
       ]);
       let specificError$1 = specificError.contents;
       if (specificError$1 === undefined) {
@@ -789,7 +792,7 @@ async function writeBatch(sql, batch, pgSchema, rollback, isInReorgThreshold, co
   }
 }
 
-function makeGetRollbackRestoredEntitiesQuery(entityConfig, pgSchema) {
+function makeGetRollbackPreTargetRowsQuery(entityConfig, pgSchema) {
   let dataFieldNames = Stdlib_Array.filterMap(entityConfig.table.fields, fieldOrDerived => {
     if (fieldOrDerived.TAG === "Field") {
       return Table.getPgDbFieldName(fieldOrDerived._0);
@@ -797,31 +800,36 @@ function makeGetRollbackRestoredEntitiesQuery(entityConfig, pgSchema) {
   });
   let dataFieldsCommaSeparated = dataFieldNames.map(name => `"` + name + `"`).join(", ");
   let historyTableName = EntityHistory.historyTableName(entityConfig.name, entityConfig.index);
-  return `SELECT DISTINCT ON (` + Table.idFieldName + `) ` + dataFieldsCommaSeparated + `
-FROM "` + pgSchema + `"."` + historyTableName + `"
-WHERE "` + EntityHistory.checkpointIdFieldName + `" <= $1
-  AND EXISTS (
-    SELECT 1
-    FROM "` + pgSchema + `"."` + historyTableName + `" h
-    WHERE h.` + Table.idFieldName + ` = "` + historyTableName + `".` + Table.idFieldName + `
-      AND h."` + EntityHistory.checkpointIdFieldName + `" > $1
-  )
-ORDER BY ` + Table.idFieldName + `, "` + EntityHistory.checkpointIdFieldName + `" DESC`;
+  return `SELECT DISTINCT ON ("` + Table.idFieldName + `") ` + dataFieldsCommaSeparated + `, "` + EntityHistory.changeFieldName + `"
+  FROM "` + pgSchema + `"."` + historyTableName + `"
+  WHERE "` + EntityHistory.checkpointIdFieldName + `" <= $1
+    AND EXISTS (
+      SELECT 1
+      FROM "` + pgSchema + `"."` + historyTableName + `" h
+      WHERE h."` + Table.idFieldName + `" = "` + historyTableName + `"."` + Table.idFieldName + `"
+        AND h."` + EntityHistory.checkpointIdFieldName + `" > $1
+    )
+  ORDER BY "` + Table.idFieldName + `", "` + EntityHistory.checkpointIdFieldName + `" DESC`;
 }
 
 function makeGetRollbackRemovedIdsQuery(entityConfig, pgSchema) {
   let historyTableName = EntityHistory.historyTableName(entityConfig.name, entityConfig.index);
-  return `SELECT DISTINCT ` + Table.idFieldName + `
-FROM "` + pgSchema + `"."` + historyTableName + `"
-WHERE "` + EntityHistory.checkpointIdFieldName + `" > $1
-AND NOT EXISTS (
-  SELECT 1
-  FROM "` + pgSchema + `"."` + historyTableName + `" h
-  WHERE h.` + Table.idFieldName + ` = "` + historyTableName + `".` + Table.idFieldName + `
-    AND h."` + EntityHistory.checkpointIdFieldName + `" <= $1
-)`;
+  return `SELECT DISTINCT "` + Table.idFieldName + `"
+  FROM "` + pgSchema + `"."` + historyTableName + `"
+  WHERE "` + EntityHistory.checkpointIdFieldName + `" > $1
+    AND NOT EXISTS (
+      SELECT 1
+      FROM "` + pgSchema + `"."` + historyTableName + `" h
+      WHERE h."` + Table.idFieldName + `" = "` + historyTableName + `"."` + Table.idFieldName + `"
+        AND h."` + EntityHistory.checkpointIdFieldName + `" <= $1
+    )`;
 }
 
+let rollbackRowStateSchema = S$RescriptSchema.object(s => [
+  s.f(Table.idFieldName, S$RescriptSchema.string),
+  s.f(EntityHistory.changeFieldName, EntityHistory.RowAction.schema)
+]);
+
 function make(sql, pgHost, pgSchema, pgPort, pgUser, pgDatabase, pgPassword, isHasuraEnabled, sink, onInitialize, onNewTables) {
   let containerName = "envio-postgres";
   let psqlExecOptions_env = Object.fromEntries([
@@ -839,32 +847,60 @@ function make(sql, pgHost, pgSchema, pgPort, pgUser, pgDatabase, pgPassword, isH
   };
   let cacheDirPath = Path.resolve(".envio", "cache");
   let isInitialized = async () => Utils.$$Array.notEmpty(await sql.unsafe(`SELECT table_schema FROM information_schema.tables WHERE table_schema = '` + pgSchema + `' AND (table_name = '` + "event_sync_state" + `' OR table_name = '` + InternalTable.Chains.table.tableName + `');`));
+  let scanCacheDir = async () => {
+    let topEntries;
+    try {
+      topEntries = await Fs.promises.readdir(cacheDirPath);
+    } catch (exn) {
+      topEntries = [];
+    }
+    let result = [];
+    await Promise.all(topEntries.map(async entry => {
+      let entryPath = Path.join(cacheDirPath, entry);
+      let isDir = (await Fs.promises.stat(entryPath)).isDirectory();
+      if (isDir) {
+        let subEntries = await Fs.promises.readdir(entryPath);
+        let tsvs = subEntries.filter(sub => sub.endsWith(".tsv"));
+        let chainId = Internal.EffectCache.parseChainId(entry);
+        if (chainId !== undefined) {
+          tsvs.forEach(sub => {
+            let effectName = sub.slice(0, -4);
+            let table = Internal.makeCacheTable(effectName, chainId);
+            result.push([
+              table,
+              Path.join(entryPath, sub)
+            ]);
+          });
+          return;
+        } else if (Utils.$$Array.notEmpty(tsvs)) {
+          return Stdlib_JsError.throwWithMessage(`Invalid effect cache directory ".envio/cache/` + entry + `". Chain cache directories must be named by a numeric chain id (e.g. "1"). Found cache files: ` + tsvs.join(", ") + `.`);
+        } else {
+          return;
+        }
+      }
+      if (!entry.endsWith(".tsv")) {
+        return;
+      }
+      let effectName = entry.slice(0, -4);
+      let table = Internal.makeCacheTable(effectName, "crossChain");
+      result.push([
+        table,
+        entryPath
+      ]);
+    }));
+    return result;
+  };
   let restoreEffectCache = async withUpload => {
     if (withUpload) {
-      let nothingToUploadErrorMessage = "Nothing to upload.";
-      let match = await Promise.all([
-        Stdlib_Promise.$$catch(Fs.promises.readdir(cacheDirPath).then(e => ({
-          TAG: "Ok",
-          _0: e
-        })), param => Promise.resolve({
-          TAG: "Error",
-          _0: nothingToUploadErrorMessage
-        })),
-        getConnectedPsqlExec(pgUser, pgHost, pgDatabase, pgPort, containerName)
-      ]);
-      let exit = 0;
-      let message;
-      let entries = match[0];
-      if (entries.TAG === "Ok") {
-        let psqlExec = match[1];
+      let entries = await scanCacheDir();
+      if (entries.length !== 0) {
+        let psqlExec = await getConnectedPsqlExec(pgUser, pgHost, pgDatabase, pgPort, containerName);
         if (psqlExec.TAG === "Ok") {
           let psqlExec$1 = psqlExec._0;
-          let cacheFiles = entries._0.filter(entry => entry.endsWith(".tsv"));
-          await Promise.all(cacheFiles.map(entry => {
-            let effectName = entry.slice(0, -4);
-            let table = Internal.makeCacheTable(effectName);
+          await Promise.all(entries.map(param => {
+            let inputFile = param[1];
+            let table = param[0];
             return sql.unsafe(makeCreateTableQuery(table, pgSchema, false)).then(() => {
-              let inputFile = Path.join(cacheDirPath, entry);
               let command = psqlExec$1 + ` -c 'COPY "` + pgSchema + `"."` + table.tableName + `" FROM STDIN WITH (FORMAT text, HEADER);' < ` + inputFile;
               return new Promise((resolve, reject) => {
                 Child_process.exec(command, psqlExecOptions, (error, stdout, param) => {
@@ -879,19 +915,10 @@ function make(sql, pgHost, pgSchema, pgPort, pgUser, pgDatabase, pgPassword, isH
           }));
           Logging.info("Successfully uploaded cache.");
         } else {
-          message = match[1]._0;
-          exit = 1;
+          Logging.error(`Failed to upload cache, continuing without it. ` + psqlExec._0);
         }
       } else {
-        message = entries._0;
-        exit = 1;
-      }
-      if (exit === 1) {
-        if (message === nothingToUploadErrorMessage) {
-          Logging.info("No cache found to upload.");
-        } else {
-          Logging.error(`Failed to upload cache, continuing without it. ` + message);
-        }
+        Logging.info("No cache found to upload.");
       }
     }
     let cacheTableInfo = await sql.unsafe(makeSchemaCacheTableInfoQuery(pgSchema));
@@ -900,11 +927,17 @@ function make(sql, pgHost, pgSchema, pgPort, pgUser, pgDatabase, pgPassword, isH
     }
     let cache = {};
     cacheTableInfo.forEach(param => {
-      let effectName = param.table_name.slice(cacheTablePrefixLength);
-      cache[effectName] = {
-        effectName: effectName,
-        count: param.count
-      };
+      let tableName = param.table_name;
+      let match = Internal.EffectCache.fromTableName(tableName);
+      if (match !== undefined) {
+        cache[tableName] = {
+          effectName: match[0],
+          scope: match[1],
+          tableName: tableName,
+          count: param.count
+        };
+        return;
+      }
     });
     return cache;
   };
@@ -998,16 +1031,14 @@ SELECT id, chain_id, -1, -1, contract_name FROM unnest($1::text[],$2::int[],$3::
     }
   };
   let setOrThrow$1 = (items, table, itemSchema) => setOrThrow(sql, items, table, itemSchema, pgSchema);
-  let setEffectCacheOrThrow = async (effect, items, initialize) => {
-    let match = effect.storageMeta;
-    let table = match.table;
+  let setEffectCacheOrThrow = async (table, itemSchema, items, initialize) => {
     if (initialize) {
       await sql.unsafe(makeCreateTableQuery(table, pgSchema, false));
       if (onNewTables !== undefined) {
         await onNewTables([table.tableName]);
       }
     }
-    return await setOrThrow$1(items, table, match.itemSchema);
+    return await setOrThrow$1(items, table, itemSchema);
   };
   let dumpEffectCache = async () => {
     try {
@@ -1030,10 +1061,16 @@ SELECT id, chain_id, -1, -1, contract_name FROM unnest($1::text[],$2::int[],$3::
       Logging.info(`Dumping cache: ` + cacheTableInfo.map(param => param.table_name + " (" + param.count.toString() + " rows)").join(", "));
       let promises = cacheTableInfo.map(async param => {
         let tableName = param.table_name;
-        let cacheName = tableName.slice(cacheTablePrefixLength);
-        let outputFile = Path.join(cacheDirPath, cacheName + ".tsv");
-        let command = psqlExec$1 + ` -c 'COPY "` + pgSchema + `"."` + tableName + `" TO STDOUT WITH (FORMAT text, HEADER);' > ` + outputFile;
-        return new Promise((resolve, reject) => {
+        let match = Internal.EffectCache.fromTableName(tableName);
+        if (match === undefined) {
+          return "";
+        }
+        let outputPath = Path.join(cacheDirPath, Internal.EffectCache.toCachePath(match[0], match[1]));
+        await Fs.promises.mkdir(Path.dirname(outputPath), {
+          recursive: true
+        });
+        let command = psqlExec$1 + ` -c 'COPY "` + pgSchema + `"."` + tableName + `" TO STDOUT WITH (FORMAT text, HEADER);' > ` + outputPath;
+        return await new Promise((resolve, reject) => {
           Child_process.exec(command, psqlExecOptions, (error, stdout, param) => {
             if (error === null) {
               return resolve(stdout);
@@ -1097,10 +1134,26 @@ SELECT id, chain_id, -1, -1, contract_name FROM unnest($1::text[],$2::int[],$3::
   let pruneStaleEntityHistory = (entityName, entityIndex, safeCheckpointId) => EntityHistory.pruneStaleEntityHistory(sql, entityName, entityIndex, pgSchema, safeCheckpointId);
   let getRollbackTargetCheckpoint = (reorgChainId, lastKnownValidBlockNumber) => InternalTable.Checkpoints.getRollbackTargetCheckpoint(sql, pgSchema, reorgChainId, lastKnownValidBlockNumber);
   let getRollbackProgressDiff = rollbackTargetCheckpointId => InternalTable.Checkpoints.getRollbackProgressDiff(sql, pgSchema, rollbackTargetCheckpointId);
-  let getRollbackData = async (entityConfig, rollbackTargetCheckpointId) => await Promise.all([
-    sql.unsafe(makeGetRollbackRemovedIdsQuery(entityConfig, pgSchema), [rollbackTargetCheckpointId.toString()], {prepare: true}),
-    sql.unsafe(makeGetRollbackRestoredEntitiesQuery(entityConfig, pgSchema), [rollbackTargetCheckpointId.toString()], {prepare: true})
-  ]);
+  let getRollbackData = async (entityConfig, rollbackTargetCheckpointId) => {
+    let match = await Promise.all([
+      sql.unsafe(makeGetRollbackRemovedIdsQuery(entityConfig, pgSchema), [rollbackTargetCheckpointId.toString()], {prepare: true}),
+      sql.unsafe(makeGetRollbackPreTargetRowsQuery(entityConfig, pgSchema), [rollbackTargetCheckpointId.toString()], {prepare: true})
+    ]);
+    let removedIds = match[0].map(row => row.id);
+    let restoredEntitiesResult = [];
+    match[1].forEach(row => {
+      let match = S$RescriptSchema.parseOrThrow(row, rollbackRowStateSchema);
+      if (match[1] === "SET") {
+        restoredEntitiesResult.push(row);
+        return;
+      }
+      removedIds.push(match[0]);
+    });
+    return [
+      removedIds,
+      restoredEntitiesResult
+    ];
+  };
   let writeBatchMethod = async (batch, rollback, isInReorgThreshold, config, allEntities, updatedEffectsCache, updatedEntities, chainMetaData) => {
     let pgUpdates = [];
     let chUpdates = [];
@@ -1227,15 +1280,15 @@ export {
   setQueryCache,
   setOrThrow,
   makeSchemaTableNamesQuery,
-  cacheTablePrefixLength,
   makeSchemaCacheTableInfoQuery,
   getConnectedPsqlExec,
   deleteByIdsOrThrow,
   makeInsertDeleteUpdatesQuery,
   executeSet,
   writeBatch,
-  makeGetRollbackRestoredEntitiesQuery,
+  makeGetRollbackPreTargetRowsQuery,
   makeGetRollbackRemovedIdsQuery,
+  rollbackRowStateSchema,
   make,
   makeStorageFromEnv,
   makePersistenceFromConfig,

package/src/Prometheus.res

@@ -606,29 +606,39 @@ let effectLabelsSchema = S.object(s => {
   s.field("effect", S.string)
 })
 
+// For metrics whose backing state lives per scope ("crossChain" or a chain
+// id) — without the label, scopes of the same effect would clobber each
+// other's gauge value.
+let effectScopeLabelsSchema = S.schema(s =>
+  {
+    "effect": s.matches(S.string),
+    "scope": s.matches(S.string),
+  }
+)
+
 module EffectCalls = {
   let timeCounter = SafeCounter.makeOrThrow(
     ~name="envio_effect_call_seconds",
     ~help="Processing time taken to call the Effect function.",
-    ~labelSchema=effectLabelsSchema,
+    ~labelSchema=effectScopeLabelsSchema,
   )
 
   let sumTimeCounter = SafeCounter.makeOrThrow(
     ~name="envio_effect_call_seconds_total",
     ~help="Cumulative time spent calling the Effect function during the indexing process.",
-    ~labelSchema=effectLabelsSchema,
+    ~labelSchema=effectScopeLabelsSchema,
   )
 
   let totalCallsCount = SafeCounter.makeOrThrow(
     ~name="envio_effect_call_total",
     ~help="Cumulative number of resolved Effect function calls during the indexing process.",
-    ~labelSchema=effectLabelsSchema,
+    ~labelSchema=effectScopeLabelsSchema,
   )
 
   let activeCallsCount = SafeGauge.makeOrThrow(
     ~name="envio_effect_active_calls",
     ~help="The number of Effect function calls that are currently running.",
-    ~labelSchema=effectLabelsSchema,
+    ~labelSchema=effectScopeLabelsSchema,
   )
 }
 
@@ -636,11 +646,11 @@ module EffectCacheCount = {
   let gauge = SafeGauge.makeOrThrow(
     ~name="envio_effect_cache",
     ~help="The number of items in the effect cache.",
-    ~labelSchema=effectLabelsSchema,
+    ~labelSchema=effectScopeLabelsSchema,
   )
 
-  let set = (~count, ~effectName) => {
-    gauge->SafeGauge.handleInt(~labels=effectName, ~value=count)
+  let set = (~count, ~effectName, ~scope) => {
+    gauge->SafeGauge.handleInt(~labels={"effect": effectName, "scope": scope}, ~value=count)
   }
 }
 
@@ -660,7 +670,7 @@ module EffectQueueCount = {
   let gauge = SafeGauge.makeOrThrow(
     ~name="envio_effect_queue",
     ~help="The number of effect calls waiting in the rate limit queue.",
-    ~labelSchema=effectLabelsSchema,
+    ~labelSchema=effectScopeLabelsSchema,
   )
 
   let timeCounter = SafeCounter.makeOrThrow(
@@ -669,8 +679,8 @@ module EffectQueueCount = {
     ~labelSchema=effectLabelsSchema,
   )
 
-  let set = (~count, ~effectName) => {
-    gauge->SafeGauge.handleInt(~labels=effectName, ~value=count)
+  let set = (~count, ~effectName, ~scope) => {
+    gauge->SafeGauge.handleInt(~labels={"effect": effectName, "scope": scope}, ~value=count)
   }
 }
 

package/src/Prometheus.res.mjs

@@ -651,13 +651,18 @@ let ProgressLatency = {
 
 let effectLabelsSchema = S$RescriptSchema.object(s => s.f("effect", S$RescriptSchema.string));
 
-let timeCounter$5 = makeOrThrow("envio_effect_call_seconds", "Processing time taken to call the Effect function.", effectLabelsSchema);
+let effectScopeLabelsSchema = S$RescriptSchema.schema(s => ({
+  effect: s.m(S$RescriptSchema.string),
+  scope: s.m(S$RescriptSchema.string)
+}));
+
+let timeCounter$5 = makeOrThrow("envio_effect_call_seconds", "Processing time taken to call the Effect function.", effectScopeLabelsSchema);
 
-let sumTimeCounter$1 = makeOrThrow("envio_effect_call_seconds_total", "Cumulative time spent calling the Effect function during the indexing process.", effectLabelsSchema);
+let sumTimeCounter$1 = makeOrThrow("envio_effect_call_seconds_total", "Cumulative time spent calling the Effect function during the indexing process.", effectScopeLabelsSchema);
 
-let totalCallsCount = makeOrThrow("envio_effect_call_total", "Cumulative number of resolved Effect function calls during the indexing process.", effectLabelsSchema);
+let totalCallsCount = makeOrThrow("envio_effect_call_total", "Cumulative number of resolved Effect function calls during the indexing process.", effectScopeLabelsSchema);
 
-let activeCallsCount = makeOrThrow$1("envio_effect_active_calls", "The number of Effect function calls that are currently running.", effectLabelsSchema);
+let activeCallsCount = makeOrThrow$1("envio_effect_active_calls", "The number of Effect function calls that are currently running.", effectScopeLabelsSchema);
 
 let EffectCalls = {
   timeCounter: timeCounter$5,
@@ -666,10 +671,13 @@ let EffectCalls = {
   activeCallsCount: activeCallsCount
 };
 
-let gauge$19 = makeOrThrow$1("envio_effect_cache", "The number of items in the effect cache.", effectLabelsSchema);
+let gauge$19 = makeOrThrow$1("envio_effect_cache", "The number of items in the effect cache.", effectScopeLabelsSchema);
 
-function set$19(count, effectName) {
-  handleInt$1(gauge$19, effectName, count);
+function set$19(count, effectName, scope) {
+  handleInt$1(gauge$19, {
+    effect: effectName,
+    scope: scope
+  }, count);
 }
 
 let EffectCacheCount = {
@@ -688,12 +696,15 @@ let EffectCacheInvalidationsCount = {
   increment: increment$7
 };
 
-let gauge$20 = makeOrThrow$1("envio_effect_queue", "The number of effect calls waiting in the rate limit queue.", effectLabelsSchema);
+let gauge$20 = makeOrThrow$1("envio_effect_queue", "The number of effect calls waiting in the rate limit queue.", effectScopeLabelsSchema);
 
 let timeCounter$6 = makeOrThrow("envio_effect_queue_wait_seconds", "The time spent waiting in the rate limit queue.", effectLabelsSchema);
 
-function set$20(count, effectName) {
-  handleInt$1(gauge$20, effectName, count);
+function set$20(count, effectName, scope) {
+  handleInt$1(gauge$20, {
+    effect: effectName,
+    scope: scope
+  }, count);
 }
 
 let EffectQueueCount = {
@@ -825,6 +836,7 @@ export {
   ProgressEventsCount,
   ProgressLatency,
   effectLabelsSchema,
+  effectScopeLabelsSchema,
   EffectCalls,
   EffectCacheCount,
   EffectCacheInvalidationsCount,

package/src/PruneStaleHistory.res

@@ -1,45 +1,156 @@
-// Throttled prune of stale entity-history rows below the safe checkpoint.
-
-let runPrune = async (state: IndexerState.t) => {
-  switch state->IndexerState.getSafeCheckpointId {
-  | None => ()
-  | Some(safeCheckpointId) =>
-    let persistence = state->IndexerState.persistence
-    await persistence.storage.pruneStaleCheckpoints(~safeCheckpointId)
-
-    for idx in 0 to persistence.allEntities->Array.length - 1 {
-      if idx !== 0 {
-        // Add some delay between entities
-        // To unblock the pg client if it's needed for something else
-        await Utils.delay(1000)
+// Prune of stale entity-history rows below the safe checkpoint.
+//
+// Pruning must never run concurrently with a write to the same entity's history
+// table: the prune's anchor deletion relies on "no history after the safe
+// checkpoint", which a concurrently committing batch can falsify, losing the
+// anchor and breaking a later rollback. The write loop enforces the safety by
+// running the concurrent group only for entities absent from the batch being
+// written (and awaiting it before the next write starts), and the forced group
+// alone after the write.
+
+let maxEntitiesPerWrite = 5
+let forcedIntervalMultiplier = 5.
+
+type targets = {
+  safeCheckpointId: Internal.checkpointId,
+  concurrent: array<Internal.entityConfig>,
+  forced: array<Internal.entityConfig>,
+}
+
+let selectFrom = (
+  ~allEntities: array<Internal.entityConfig>,
+  ~lastPrunedAtMillis: dict<float>,
+  ~writtenEntityNames: Utils.Set.t<string>,
+  ~isRollback,
+  ~nowMillis,
+  ~intervalMillis,
+  ~safeCheckpointId,
+) => {
+  let byOldestPrune = ((a, _), (b, _)) => a -. b
+  let toEntities = candidates => candidates->Array.map(((_, entityConfig)) => entityConfig)
+
+  let concurrentCandidates = []
+  let forcedCandidates = []
+  allEntities->Array.forEach(entityConfig => {
+    if entityConfig.storage.postgres {
+      let lastPrunedAt =
+        lastPrunedAtMillis
+        ->Utils.Dict.dangerouslyGetNonOption(entityConfig.name)
+        ->Option.getOr(0.)
+      if (
+        !isRollback &&
+        !(writtenEntityNames->Utils.Set.has(entityConfig.name)) &&
+        nowMillis -. lastPrunedAt >= intervalMillis
+      ) {
+        concurrentCandidates->Array.push((lastPrunedAt, entityConfig))
+      } else if nowMillis -. lastPrunedAt >= intervalMillis *. forcedIntervalMultiplier {
+        forcedCandidates->Array.push((lastPrunedAt, entityConfig))
       }
-      let entityConfig = persistence.allEntities->Array.getUnsafe(idx)
-      let timeRef = Performance.now()
-      try {
-        let () = await persistence.storage.pruneStaleEntityHistory(
-          ~entityName=entityConfig.name,
-          ~entityIndex=entityConfig.index,
+    }
+  })
+
+  let sortedConcurrent = concurrentCandidates->Array.toSorted(byOldestPrune)
+  // Concurrent candidates beyond the cap are not selected, so the starved
+  // ones among them still qualify for the forced group.
+  for idx in maxEntitiesPerWrite to sortedConcurrent->Array.length - 1 {
+    let (lastPrunedAt, _) = sortedConcurrent->Array.getUnsafe(idx)
+    if nowMillis -. lastPrunedAt >= intervalMillis *. forcedIntervalMultiplier {
+      forcedCandidates->Array.push(sortedConcurrent->Array.getUnsafe(idx))
+    }
+  }
+
+  {
+    safeCheckpointId,
+    concurrent: sortedConcurrent->Array.slice(~start=0, ~end=maxEntitiesPerWrite)->toEntities,
+    forced: forcedCandidates
+    ->Array.toSorted(byOldestPrune)
+    ->Array.slice(~start=0, ~end=maxEntitiesPerWrite)
+    ->toEntities,
+  }
+}
+
+let select = (state: IndexerState.t, ~writtenEntityNames, ~isRollback) => {
+  let config = state->IndexerState.config
+  if config->Config.shouldPruneHistory(~isInReorgThreshold=state->IndexerState.isInReorgThreshold) {
+    switch state->IndexerState.getSafeCheckpointId {
+    | None => None
+    | Some(safeCheckpointId) =>
+      Some(
+        selectFrom(
+          ~allEntities=(state->IndexerState.persistence).allEntities,
+          ~lastPrunedAtMillis=state->IndexerState.lastPrunedAtMillis,
+          ~writtenEntityNames,
+          ~isRollback,
+          ~nowMillis=Date.now(),
+          ~intervalMillis=Env.ThrottleWrites.pruneStaleDataIntervalMillis->Int.toFloat,
           ~safeCheckpointId,
-        )
-      } catch {
-      | exn =>
-        exn->ErrorHandling.mkLogAndRaise(
-          ~msg=`Failed to prune stale entity history`,
-          ~logger=Logging.createChild(
-            ~params={
-              "entityName": entityConfig.name,
-              "safeCheckpointId": safeCheckpointId,
-            },
-          ),
-        )
-      }
+        ),
+      )
+    }
+  } else {
+    None
+  }
+}
+
+let pruneEntities = async (state: IndexerState.t, ~entities, ~safeCheckpointId) => {
+  let persistence = state->IndexerState.persistence
+  for idx in 0 to entities->Array.length - 1 {
+    let entityConfig: Internal.entityConfig = entities->Array.getUnsafe(idx)
+    let timeRef = Performance.now()
+    // Recorded for failures too, so a failing prune retries on the same
+    // interval instead of on every write.
+    state->IndexerState.lastPrunedAtMillis->Dict.set(entityConfig.name, Date.now())
+    switch await persistence.storage.pruneStaleEntityHistory(
+      ~entityName=entityConfig.name,
+      ~entityIndex=entityConfig.index,
+      ~safeCheckpointId,
+    ) {
+    | () =>
       Prometheus.RollbackHistoryPrune.increment(
         ~timeSeconds=Performance.secondsSince(timeRef),
         ~entityName=entityConfig.name,
       )
+    | exception exn =>
+      // Pruning is cleanup; a failure must not fail the write loop.
+      Logging.createChild(
+        ~params={
+          "entityName": entityConfig.name,
+          "safeCheckpointId": safeCheckpointId,
+        },
+      )->Logging.childErrorWithExn(exn->Utils.prettifyExn, `Failed to prune stale entity history`)
     }
   }
 }
 
-let schedule = (state: IndexerState.t) =>
-  state->IndexerState.pruneStaleEntityHistoryThrottler->Throttler.schedule(() => runPrune(state))
+let pruneCheckpoints = async (state: IndexerState.t, ~safeCheckpointId) => {
+  switch await (state->IndexerState.persistence).storage.pruneStaleCheckpoints(~safeCheckpointId) {
+  | () => ()
+  | exception exn =>
+    Logging.createChild(~params={"safeCheckpointId": safeCheckpointId})->Logging.childErrorWithExn(
+      exn->Utils.prettifyExn,
+      `Failed to prune stale checkpoints`,
+    )
+  }
+}
+
+let runConcurrent = async (state: IndexerState.t, ~targets) => {
+  switch targets {
+  | Some({safeCheckpointId, concurrent}) if concurrent->Utils.Array.notEmpty =>
+    await pruneCheckpoints(state, ~safeCheckpointId)
+    await pruneEntities(state, ~entities=concurrent, ~safeCheckpointId)
+  | Some(_) | None => ()
+  }
+}
+
+let runForced = async (state: IndexerState.t, ~targets) => {
+  switch targets {
+  | Some({safeCheckpointId, concurrent, forced}) if forced->Utils.Array.notEmpty =>
+    // When nothing ran concurrently (eg a rollback write), checkpoint pruning
+    // lands here, after the write, so it never overlaps a rollback transaction.
+    if concurrent->Utils.Array.isEmpty {
+      await pruneCheckpoints(state, ~safeCheckpointId)
+    }
+    await pruneEntities(state, ~entities=forced, ~safeCheckpointId)
+  | Some(_) | None => ()
+  }
+}