fix: discount re-orged microblock transactions from ft_balances table

migrations/1752852195051_recalculate-balances.js

@@ -0,0 +1,100 @@
+/* eslint-disable camelcase */
+
+exports.shorthands = undefined;
+
+exports.up = pgm => {
+  // Remove old balances.
+  pgm.sql(`TRUNCATE TABLE ft_balances`);
+
+  // Recalculate STX balances
+  pgm.sql(`
+    WITH all_balances AS (
+        SELECT sender AS address, -SUM(amount) AS balance_change
+        FROM stx_events
+        WHERE asset_event_type_id IN (1, 3) -- Transfers and Burns affect the sender's balance
+        AND canonical = true AND microblock_canonical = true
+        GROUP BY sender
+      UNION ALL
+        SELECT recipient AS address, SUM(amount) AS balance_change
+        FROM stx_events
+        WHERE asset_event_type_id IN (1, 2) -- Transfers and Mints affect the recipient's balance
+        AND canonical = true AND microblock_canonical = true
+        GROUP BY recipient
+    ),
+    net_balances AS (
+      SELECT address, SUM(balance_change) AS balance
+      FROM all_balances
+      GROUP BY address
+    ),
+    fees AS (
+      SELECT address, SUM(total_fees) AS total_fees
+      FROM (
+          SELECT sender_address AS address, SUM(fee_rate) AS total_fees
+          FROM txs
+          WHERE canonical = true AND microblock_canonical = true AND sponsored = false
+          GROUP BY sender_address
+        UNION ALL
+          SELECT sponsor_address AS address, SUM(fee_rate) AS total_fees
+          FROM txs
+          WHERE canonical = true AND microblock_canonical = true AND sponsored = true
+          GROUP BY sponsor_address
+      ) AS subquery
+      GROUP BY address
+    ),
+    rewards AS (
+      SELECT
+        recipient AS address,
+        SUM(
+          coinbase_amount + tx_fees_anchored + tx_fees_streamed_confirmed + tx_fees_streamed_produced
+        ) AS total_rewards
+      FROM miner_rewards
+      WHERE canonical = true
+      GROUP BY recipient
+    ),
+    all_addresses AS (
+      SELECT address FROM net_balances
+      UNION
+      SELECT address FROM fees
+      UNION
+      SELECT address FROM rewards
+    )
+    INSERT INTO ft_balances (address, balance, token)
+    SELECT
+      aa.address,
+      COALESCE(nb.balance, 0) - COALESCE(f.total_fees, 0) + COALESCE(r.total_rewards, 0) AS balance,
+      'stx' AS token
+    FROM all_addresses aa
+    LEFT JOIN net_balances nb ON aa.address = nb.address
+    LEFT JOIN fees f ON aa.address = f.address
+    LEFT JOIN rewards r ON aa.address = r.address
+  `);
+
+  // Recalculate FT balances
+  pgm.sql(`
+    WITH all_balances AS (
+        SELECT sender AS address, asset_identifier, -SUM(amount) AS balance_change
+        FROM ft_events
+        WHERE asset_event_type_id IN (1, 3) -- Transfers and Burns affect the sender's balance
+          AND canonical = true 
+          AND microblock_canonical = true
+        GROUP BY sender, asset_identifier
+      UNION ALL
+        SELECT recipient AS address, asset_identifier, SUM(amount) AS balance_change
+        FROM ft_events
+        WHERE asset_event_type_id IN (1, 2) -- Transfers and Mints affect the recipient's balance
+          AND canonical = true 
+          AND microblock_canonical = true
+        GROUP BY recipient, asset_identifier
+    ),
+    net_balances AS (
+      SELECT address, asset_identifier, SUM(balance_change) AS balance
+      FROM all_balances
+      GROUP BY address, asset_identifier
+    )
+    INSERT INTO ft_balances (address, balance, token)
+    SELECT address, balance, asset_identifier AS token
+    FROM net_balances
+  `);
+};
+
+exports.down = pgm => {};

src/datastore/common.ts

@@ -1109,6 +1109,7 @@ export interface DbPoxCycleSignerStacker {
 
 interface ReOrgEntities {
   blocks: number;
+  microblockHashes: string[];
   microblocks: number;
   minerRewards: number;
   txs: number;

src/datastore/helpers.ts

@@ -1313,6 +1313,7 @@ export function newReOrgUpdatedEntities(): ReOrgUpdatedEntities {
   return {
     markedCanonical: {
       blocks: 0,
+      microblockHashes: [],
       microblocks: 0,
       minerRewards: 0,
       txs: 0,
@@ -1333,6 +1334,7 @@ export function newReOrgUpdatedEntities(): ReOrgUpdatedEntities {
     },
     markedNonCanonical: {
       blocks: 0,
+      microblockHashes: [],
       microblocks: 0,
       minerRewards: 0,
       txs: 0,

src/datastore/pg-write-store.ts

@@ -207,7 +207,7 @@ export class PgWriteStore extends PgStore {
 
     await this.sqlWriteTransaction(async sql => {
       const chainTip = await this.getChainTip(sql);
-      await this.handleReorg(sql, data.block, chainTip.block_height);
+      const reorg = await this.handleReorg(sql, data.block, chainTip.block_height);
       const isCanonical = data.block.block_height > chainTip.block_height;
       if (!isCanonical) {
         markBlockUpdateDataAsNonCanonical(data);
@@ -292,6 +292,14 @@ export class PgWriteStore extends PgStore {
           // via microblocks.
           q.enqueue(() => this.updateStxBalances(sql, data.txs, data.minerRewards));
           q.enqueue(() => this.updateFtBalances(sql, data.txs));
+          // If this block re-orgs past microblocks, though, we must discount the balances generated
+          // by their txs which are now also reorged. We must do this here because the block re-org
+          // logic is decoupled from the microblock re-org logic so previous balance updates will
+          // not apply.
+          await this.updateFtBalancesFromMicroblockReOrg(sql, [
+            ...reorg.markedNonCanonical.microblockHashes,
+            ...reorg.markedCanonical.microblockHashes,
+          ]);
         }
         if (data.poxSetSigners && data.poxSetSigners.signers) {
           const poxSet = data.poxSetSigners;
@@ -1150,6 +1158,116 @@ export class PgWriteStore extends PgStore {
     }
   }
 
+  async updateFtBalancesFromMicroblockReOrg(sql: PgSqlClient, microblockHashes: string[]) {
+    if (microblockHashes.length === 0) return;
+    await sql`
+      WITH updated_txs AS (
+        SELECT tx_id, sender_address, nonce, sponsor_address, fee_rate, sponsored, canonical, microblock_canonical
+        FROM txs
+        WHERE microblock_hash IN ${sql(microblockHashes)}
+      ),
+      affected_addresses AS (
+          SELECT
+            sender_address AS address,
+            fee_rate AS fee_change,
+            canonical,
+            microblock_canonical,
+            sponsored
+          FROM updated_txs
+          WHERE sponsored = false
+        UNION ALL
+          SELECT
+            sponsor_address AS address,
+            fee_rate AS fee_change,
+            canonical,
+            microblock_canonical,
+            sponsored
+          FROM updated_txs
+          WHERE sponsored = true
+      ),
+      balances_update AS (
+        SELECT
+          a.address,
+          SUM(CASE WHEN a.canonical AND a.microblock_canonical THEN -a.fee_change ELSE a.fee_change END) AS balance_change
+        FROM affected_addresses a
+        GROUP BY a.address
+      )
+      INSERT INTO ft_balances (address, token, balance)
+      SELECT b.address, 'stx', b.balance_change
+      FROM balances_update b
+      ON CONFLICT (address, token)
+      DO UPDATE
+      SET balance = ft_balances.balance + EXCLUDED.balance
+      RETURNING ft_balances.address
+    `;
+    await sql`
+      WITH updated_events AS (
+        SELECT sender, recipient, amount, asset_event_type_id, asset_identifier, canonical, microblock_canonical
+        FROM ft_events
+        WHERE microblock_hash IN ${sql(microblockHashes)}
+      ),
+      event_changes AS (
+        SELECT address, asset_identifier, SUM(balance_change) AS balance_change
+        FROM (
+            SELECT sender AS address, asset_identifier,
+              SUM(CASE WHEN canonical AND microblock_canonical THEN -amount ELSE amount END) AS balance_change
+            FROM updated_events
+            WHERE asset_event_type_id IN (1, 3) -- Transfers and Burns affect the sender's balance
+            GROUP BY sender, asset_identifier
+          UNION ALL
+            SELECT recipient AS address, asset_identifier,
+              SUM(CASE WHEN canonical AND microblock_canonical THEN amount ELSE -amount END) AS balance_change
+            FROM updated_events
+            WHERE asset_event_type_id IN (1, 2) -- Transfers and Mints affect the recipient's balance
+            GROUP BY recipient, asset_identifier
+        ) AS subquery
+        GROUP BY address, asset_identifier
+      )
+      INSERT INTO ft_balances (address, token, balance)
+      SELECT ec.address, ec.asset_identifier, ec.balance_change
+      FROM event_changes ec
+      ON CONFLICT (address, token)
+      DO UPDATE
+      SET balance = ft_balances.balance + EXCLUDED.balance
+      RETURNING ft_balances.address
+    `;
+    await sql`
+      WITH updated_events AS (
+        SELECT sender, recipient, amount, asset_event_type_id, canonical, microblock_canonical
+        FROM stx_events
+        WHERE microblock_hash IN ${sql(microblockHashes)}
+      ),
+      event_changes AS (
+        SELECT
+          address,
+          SUM(balance_change) AS balance_change
+        FROM (
+            SELECT
+              sender AS address,
+              SUM(CASE WHEN canonical AND microblock_canonical THEN -amount ELSE amount END) AS balance_change
+            FROM updated_events
+            WHERE asset_event_type_id IN (1, 3) -- Transfers and Burns affect the sender's balance
+            GROUP BY sender
+          UNION ALL
+            SELECT
+              recipient AS address,
+              SUM(CASE WHEN canonical AND microblock_canonical THEN amount ELSE -amount END) AS balance_change
+            FROM updated_events
+            WHERE asset_event_type_id IN (1, 2) -- Transfers and Mints affect the recipient's balance
+            GROUP BY recipient
+        ) AS subquery
+        GROUP BY address
+      )
+      INSERT INTO ft_balances (address, token, balance)
+      SELECT ec.address, 'stx', ec.balance_change
+      FROM event_changes ec
+      ON CONFLICT (address, token)
+      DO UPDATE
+      SET balance = ft_balances.balance + EXCLUDED.balance
+      RETURNING ft_balances.address
+    `;
+  }
+
   async updateStxEvents(sql: PgSqlClient, entries: { tx: DbTx; stxEvents: DbStxEvent[] }[]) {
     const values: StxEventInsertValues[] = [];
     for (const { tx, stxEvents } of entries) {
@@ -3289,7 +3407,13 @@ export class PgWriteStore extends PgStore {
       microblocksAccepted.add(mb);
     });
     updatedEntities.markedCanonical.microblocks += microblocksAccepted.size;
+    updatedEntities.markedCanonical.microblockHashes.push(
+      ...microCanonicalUpdateResult.acceptedMicroblocks
+    );
     updatedEntities.markedNonCanonical.microblocks += microblocksOrphaned.size;
+    updatedEntities.markedNonCanonical.microblockHashes.push(
+      ...microCanonicalUpdateResult.orphanedMicroblocks
+    );
 
     const markCanonicalResult = await this.markEntitiesCanonical(
       sql,