feat: add cache handler for principal activity including mempool transactions (#2100)

rafaelcr · web-flow · commit 2370c211e957 · 2024-10-01T10:42:36.000-06:00
* feat: mempool principal cache

* fix: unify principal activity query

* fix: narrow mempool search

* test: explicit receipt times
diff --git a/src/api/controllers/cache-controller.ts b/src/api/controllers/cache-controller.ts
@@ -23,6 +23,8 @@ enum ETagType {
   transaction = 'transaction',
   /** Etag based on the confirmed balance of a single principal (STX address or contract id) */
   principal = 'principal',
+  /** Etag based on `principal` but also including its mempool transactions */
+  principalMempool = 'principal_mempool',
 }
 
 /** Value that means the ETag did get calculated but it is empty. */
@@ -78,23 +80,18 @@ async function calculateETag(
   etagType: ETagType,
   req: FastifyRequest
 ): Promise<ETag | undefined> {
-  switch (etagType) {
-    case ETagType.chainTip:
-      try {
+  try {
+    switch (etagType) {
+      case ETagType.chainTip:
         const chainTip = await db.getChainTip(db.sql);
         if (chainTip.block_height === 0) {
           // This should never happen unless the API is serving requests before it has synced any
           // blocks.
           return;
         }
         return chainTip.microblock_hash ?? chainTip.index_block_hash;
-      } catch (error) {
-        logger.error(error, 'Unable to calculate chain_tip ETag');
-        return;
-      }
 
-    case ETagType.mempool:
-      try {
+      case ETagType.mempool:
         const digest = await db.getMempoolTxDigest();
         if (!digest.found) {
           // This should never happen unless the API is serving requests before it has synced any
@@ -106,13 +103,8 @@ async function calculateETag(
           return ETAG_EMPTY;
         }
         return digest.result.digest;
-      } catch (error) {
-        logger.error(error, 'Unable to calculate mempool etag');
-        return;
-      }
 
-    case ETagType.transaction:
-      try {
+      case ETagType.transaction:
         const tx_id = (req.params as { tx_id: string }).tx_id;
         const normalizedTxId = normalizeHashString(tx_id);
         if (normalizedTxId === false) {
@@ -129,23 +121,21 @@ async function calculateETag(
           status.result.status.toString(),
         ];
         return sha256(elements.join(':'));
-      } catch (error) {
-        logger.error(error, 'Unable to calculate transaction etag');
-        return;
-      }
 
-    case ETagType.principal:
-      try {
+      case ETagType.principal:
+      case ETagType.principalMempool:
         const params = req.params as { address?: string; principal?: string };
         const principal = params.address ?? params.principal;
         if (!principal) return ETAG_EMPTY;
-        const activity = await db.getPrincipalLastActivityTxIds(principal);
-        const text = `${activity.stx_tx_id}:${activity.ft_tx_id}:${activity.nft_tx_id}`;
-        return sha256(text);
-      } catch (error) {
-        logger.error(error, 'Unable to calculate principal etag');
-        return;
-      }
+        const activity = await db.getPrincipalLastActivityTxIds(
+          principal,
+          etagType == ETagType.principalMempool
+        );
+        if (!activity.length) return ETAG_EMPTY;
+        return sha256(activity.join(':'));
+    }
+  } catch (error) {
+    logger.error(error, `Unable to calculate ${etagType} etag`);
   }
 }
 
@@ -193,3 +183,7 @@ export async function handleTransactionCache(request: FastifyRequest, reply: Fas
 export async function handlePrincipalCache(request: FastifyRequest, reply: FastifyReply) {
   return handleCache(ETagType.principal, request, reply);
 }
+
+export async function handlePrincipalMempoolCache(request: FastifyRequest, reply: FastifyReply) {
+  return handleCache(ETagType.principalMempool, request, reply);
+}
diff --git a/src/api/routes/address.ts b/src/api/routes/address.ts
@@ -16,9 +16,8 @@ import {
 import { InvalidRequestError, InvalidRequestErrorType, NotFoundError } from '../../errors';
 import { decodeClarityValueToRepr } from 'stacks-encoding-native-js';
 import {
-  handleChainTipCache,
-  handleMempoolCache,
   handlePrincipalCache,
+  handlePrincipalMempoolCache,
   handleTransactionCache,
 } from '../controllers/cache-controller';
 import { PgStore } from '../../datastore/pg-store';
@@ -45,7 +44,6 @@ import {
   AddressTransactionWithTransfers,
   AddressTransactionWithTransfersSchema,
   InboundStxTransfer,
-  InboundStxTransferSchema,
 } from '../schemas/entities/addresses';
 import { PaginatedResponse } from '../schemas/util';
 import { MempoolTransaction, MempoolTransactionSchema } from '../schemas/entities/transactions';
@@ -151,7 +149,7 @@ export const AddressRoutes: FastifyPluginAsync<
       schema: {
         operationId: 'get_account_balance',
         summary: 'Get account balances',
-        description: `Retrieves total account balance information for a given Address or Contract Identifier. This includes the balances of  STX Tokens, Fungible Tokens and Non-Fungible Tokens for the account.`,
+        description: `Retrieves total account balance information for a given Address or Contract Identifier. This includes the balances of STX Tokens, Fungible Tokens and Non-Fungible Tokens for the account.`,
         tags: ['Accounts'],
         params: Type.Object({
           principal: PrincipalSchema,
@@ -629,7 +627,7 @@ export const AddressRoutes: FastifyPluginAsync<
   fastify.get(
     '/:principal/mempool',
     {
-      preHandler: handleMempoolCache,
+      preHandler: handlePrincipalMempoolCache,
       schema: {
         operationId: 'get_address_mempool_transactions',
         summary: 'Transactions for address',
@@ -676,7 +674,7 @@ export const AddressRoutes: FastifyPluginAsync<
   fastify.get(
     '/:principal/nonces',
     {
-      preHandler: handleMempoolCache,
+      preHandler: handlePrincipalMempoolCache,
       schema: {
         operationId: 'get_account_nonces',
         summary: 'Get the latest nonce used by an account',
diff --git a/src/datastore/pg-store.ts b/src/datastore/pg-store.ts
@@ -4407,41 +4407,56 @@ export class PgStore extends BasePgStore {
 
   /** Retrieves the last transaction IDs with STX, FT and NFT activity for a principal */
   async getPrincipalLastActivityTxIds(
-    principal: string
-  ): Promise<{ stx_tx_id: string | null; ft_tx_id: string | null; nft_tx_id: string | null }> {
-    const result = await this.sql<
-      { stx_tx_id: string | null; ft_tx_id: string | null; nft_tx_id: string | null }[]
-    >`
-      WITH last_stx AS (
-        SELECT tx_id
-        FROM principal_stx_txs
-        WHERE principal = ${principal} AND canonical = true AND microblock_canonical = true
-        ORDER BY block_height DESC, microblock_sequence DESC, tx_index DESC
-        LIMIT 1
-      ),
-      last_ft AS (
-        SELECT tx_id
-        FROM ft_events
-        WHERE (sender = ${principal} OR recipient = ${principal})
-          AND canonical = true
-          AND microblock_canonical = true
-        ORDER BY block_height DESC, microblock_sequence DESC, tx_index DESC, event_index DESC
-        LIMIT 1
-      ),
-      last_nft AS (
-        SELECT tx_id
-        FROM nft_events
-        WHERE (sender = ${principal} OR recipient = ${principal})
-          AND canonical = true
-          AND microblock_canonical = true
-        ORDER BY block_height DESC, microblock_sequence DESC, tx_index DESC, event_index DESC
-        LIMIT 1
+    principal: string,
+    includeMempool: boolean = false
+  ): Promise<string[]> {
+    const result = await this.sql<{ tx_id: string }[]>`
+      WITH activity AS (
+        (
+          SELECT tx_id
+          FROM principal_stx_txs
+          WHERE principal = ${principal} AND canonical = true AND microblock_canonical = true
+          ORDER BY block_height DESC, microblock_sequence DESC, tx_index DESC
+          LIMIT 1
+        )
+        UNION
+        (
+          SELECT tx_id
+          FROM ft_events
+          WHERE (sender = ${principal} OR recipient = ${principal})
+            AND canonical = true
+            AND microblock_canonical = true
+          ORDER BY block_height DESC, microblock_sequence DESC, tx_index DESC, event_index DESC
+          LIMIT 1
+        )
+        UNION
+        (
+          SELECT tx_id
+          FROM nft_events
+          WHERE (sender = ${principal} OR recipient = ${principal})
+            AND canonical = true
+            AND microblock_canonical = true
+          ORDER BY block_height DESC, microblock_sequence DESC, tx_index DESC, event_index DESC
+          LIMIT 1
+        )
+        ${
+          includeMempool
+            ? this.sql`UNION
+            (
+              SELECT tx_id
+              FROM mempool_txs
+              WHERE pruned = false AND
+                (sender_address = ${principal}
+                OR sponsor_address = ${principal}
+                OR token_transfer_recipient_address = ${principal})
+              ORDER BY receipt_time DESC, sender_address DESC, nonce DESC
+              LIMIT 1
+            )`
+            : this.sql``
+        }
       )
-      SELECT
-        (SELECT tx_id FROM last_stx) AS stx_tx_id,
-        (SELECT tx_id FROM last_ft) AS ft_tx_id,
-        (SELECT tx_id FROM last_nft) AS nft_tx_id
+      SELECT DISTINCT tx_id FROM activity WHERE tx_id IS NOT NULL
     `;
-    return result[0];
+    return result.map(r => r.tx_id);
   }
 }
diff --git a/tests/api/cache-control.test.ts b/tests/api/cache-control.test.ts
@@ -730,4 +730,92 @@ describe('cache-control tests', () => {
     expect(request10.status).toBe(304);
     expect(request10.text).toBe('');
   });
+
+  test('principal mempool cache control', async () => {
+    const sender_address = 'SP3FXEKSA6D4BW3TFP2BWTSREV6FY863Y90YY7D8G';
+    const url = `/extended/v1/address/${sender_address}/mempool`;
+    await db.update(
+      new TestBlockBuilder({
+        block_height: 1,
+        index_block_hash: '0x01',
+        parent_index_block_hash: '0x00',
+      }).build()
+    );
+
+    // ETag zero.
+    const request1 = await supertest(api.server).get(url);
+    expect(request1.status).toBe(200);
+    expect(request1.type).toBe('application/json');
+    const etag0 = request1.headers['etag'];
+
+    // Add STX tx.
+    await db.updateMempoolTxs({
+      mempoolTxs: [testMempoolTx({ tx_id: '0x0001', receipt_time: 1000, sender_address })],
+    });
+
+    // Valid ETag.
+    const request2 = await supertest(api.server).get(url);
+    expect(request2.status).toBe(200);
+    expect(request2.type).toBe('application/json');
+    expect(request2.headers['etag']).toBeTruthy();
+    const etag1 = request2.headers['etag'];
+    expect(etag1).not.toEqual(etag0);
+
+    // Cache works with valid ETag.
+    const request3 = await supertest(api.server).get(url).set('If-None-Match', etag1);
+    expect(request3.status).toBe(304);
+    expect(request3.text).toBe('');
+
+    // Add sponsor tx.
+    await db.updateMempoolTxs({
+      mempoolTxs: [
+        testMempoolTx({ tx_id: '0x0002', receipt_time: 2000, sponsor_address: sender_address }),
+      ],
+    });
+
+    // Cache is now a miss.
+    const request4 = await supertest(api.server).get(url).set('If-None-Match', etag1);
+    expect(request4.status).toBe(200);
+    expect(request4.type).toBe('application/json');
+    expect(request4.headers['etag']).not.toEqual(etag1);
+    const etag2 = request4.headers['etag'];
+
+    // Cache works with new ETag.
+    const request5 = await supertest(api.server).get(url).set('If-None-Match', etag2);
+    expect(request5.status).toBe(304);
+    expect(request5.text).toBe('');
+
+    // Add token recipient tx.
+    await db.updateMempoolTxs({
+      mempoolTxs: [
+        testMempoolTx({
+          tx_id: '0x0003',
+          receipt_time: 3000,
+          token_transfer_recipient_address: sender_address,
+        }),
+      ],
+    });
+
+    // Cache is now a miss.
+    const request6 = await supertest(api.server).get(url).set('If-None-Match', etag2);
+    expect(request6.status).toBe(200);
+    expect(request6.type).toBe('application/json');
+    expect(request6.headers['etag']).not.toEqual(etag2);
+    const etag3 = request6.headers['etag'];
+
+    // Cache works with new ETag.
+    const request7 = await supertest(api.server).get(url).set('If-None-Match', etag3);
+    expect(request7.status).toBe(304);
+    expect(request7.text).toBe('');
+
+    // Change mempool with no changes to this address.
+    await db.updateMempoolTxs({
+      mempoolTxs: [testMempoolTx({ tx_id: '0x0004', receipt_time: 4000 })],
+    });
+
+    // Cache still works.
+    const request8 = await supertest(api.server).get(url).set('If-None-Match', etag3);
+    expect(request8.status).toBe(304);
+    expect(request8.text).toBe('');
+  });
 });
diff --git a/tests/utils/test-builders.ts b/tests/utils/test-builders.ts
@@ -301,6 +301,8 @@ interface TestMempoolTxArgs {
   nonce?: number;
   fee_rate?: bigint;
   raw_tx?: string;
+  sponsor_address?: string;
+  receipt_time?: number;
 }
 
 /**
@@ -316,12 +318,12 @@ export function testMempoolTx(args?: TestMempoolTxArgs): DbMempoolTxRaw {
     nonce: args?.nonce ?? 0,
     raw_tx: args?.raw_tx ?? '0x01234567',
     type_id: args?.type_id ?? DbTxTypeId.TokenTransfer,
-    receipt_time: (new Date().getTime() / 1000) | 0,
+    receipt_time: args?.receipt_time ?? (new Date().getTime() / 1000) | 0,
     status: args?.status ?? DbTxStatus.Pending,
     post_conditions: '0x01f5',
     fee_rate: args?.fee_rate ?? 1234n,
     sponsored: false,
-    sponsor_address: undefined,
+    sponsor_address: args?.sponsor_address,
     origin_hash_mode: 1,
     sender_address: args?.sender_address ?? SENDER_ADDRESS,
     token_transfer_amount: args?.token_transfer_amount ?? 1234n,
