feat(common): batch insert receipts via mpsc (#285)

* feat(common): use channel to batch insert

Signed-off-by: Gustavo Inacio <[email protected]>

* fix(common): only cancel when queue empty

Signed-off-by: Gustavo Inacio <[email protected]>

---------

Signed-off-by: Gustavo Inacio <[email protected]>

Author: gusinacio

.sqlx/query-85ff59ba99e8c59b2273bde484cc555a3268c17806410275b9ef943a293d7867.json

@@ -9,21 +9,25 @@ use crate::tap::checks::timestamp_check::TimestampCheck;
 use crate::{escrow_accounts::EscrowAccounts, prelude::Allocation};
 use alloy::dyn_abi::Eip712Domain;
 use eventuals::Eventual;
+use receipt_store::{DatabaseReceipt, InnerContext};
 use sqlx::PgPool;
 use std::fmt::Debug;
 use std::time::Duration;
 use std::{collections::HashMap, sync::Arc};
 use tap_core::receipt::checks::ReceiptCheck;
 use thegraph_core::Address;
+use tokio::sync::mpsc::{self, Sender};
+use tokio_util::sync::CancellationToken;
 use tracing::error;
 
 mod checks;
 mod receipt_store;
 
-#[derive(Clone)]
 pub struct IndexerTapContext {
-    pgpool: PgPool,
     domain_separator: Arc<Eip712Domain>,
+
+    receipt_producer: Sender<DatabaseReceipt>,
+    cancelation_token: CancellationToken,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -54,9 +58,22 @@ impl IndexerTapContext {
     }
 
     pub async fn new(pgpool: PgPool, domain_separator: Eip712Domain) -> Self {
+        const MAX_RECEIPT_QUEUE_SIZE: usize = 1000;
+        let (tx, rx) = mpsc::channel(MAX_RECEIPT_QUEUE_SIZE);
+        let cancelation_token = CancellationToken::new();
+        let inner = InnerContext { pgpool };
+        Self::spawn_store_receipt_task(inner, rx, cancelation_token.clone());
+
         Self {
-            pgpool,
+            cancelation_token,
+            receipt_producer: tx,
             domain_separator: Arc::new(domain_separator),
         }
     }
 }
+
+impl Drop for IndexerTapContext {
+    fn drop(&mut self) {
+        self.cancelation_token.cancel();
+    }
+}

.sqlx/query-ed054cb84e667373b57cbb62dbd4a96ec4b8b54f04355188eb7c16703e41e7b3.json

@@ -1,18 +1,101 @@
 // Copyright 2023-, GraphOps and Semiotic Labs.
 // SPDX-License-Identifier: Apache-2.0
 
-use alloy::hex::ToHexExt;
+use alloy::{dyn_abi::Eip712Domain, hex::ToHexExt};
 use anyhow::anyhow;
 use bigdecimal::num_bigint::BigInt;
-use sqlx::types::BigDecimal;
+use sqlx::{types::BigDecimal, PgPool};
 use tap_core::{
     manager::adapters::ReceiptStore,
     receipt::{state::Checking, ReceiptWithState},
 };
+use tokio::{select, sync::mpsc::Receiver, task::JoinHandle};
+use tokio_util::sync::CancellationToken;
 use tracing::error;
 
 use super::{AdapterError, IndexerTapContext};
 
+#[derive(Clone)]
+pub struct InnerContext {
+    pub pgpool: PgPool,
+}
+
+impl InnerContext {
+    async fn store_receipts(&self, receipts: Vec<DatabaseReceipt>) -> Result<(), AdapterError> {
+        let receipts_len = receipts.len();
+        let mut signers = Vec::with_capacity(receipts_len);
+        let mut signatures = Vec::with_capacity(receipts_len);
+        let mut allocation_ids = Vec::with_capacity(receipts_len);
+        let mut timestamps = Vec::with_capacity(receipts_len);
+        let mut nonces = Vec::with_capacity(receipts_len);
+        let mut values = Vec::with_capacity(receipts_len);
+
+        for receipt in receipts {
+            signers.push(receipt.signer_address);
+            signatures.push(receipt.signature);
+            allocation_ids.push(receipt.allocation_id);
+            timestamps.push(receipt.timestamp_ns);
+            nonces.push(receipt.nonce);
+            values.push(receipt.value);
+        }
+        sqlx::query!(
+            r#"INSERT INTO scalar_tap_receipts (
+                signer_address,
+                signature,
+                allocation_id,
+                timestamp_ns,
+                nonce,
+                value
+            ) SELECT * FROM UNNEST(
+                $1::CHAR(40)[],
+                $2::BYTEA[],
+                $3::CHAR(40)[],
+                $4::NUMERIC(20)[],
+                $5::NUMERIC(20)[],
+                $6::NUMERIC(40)[]
+            )"#,
+            &signers,
+            &signatures,
+            &allocation_ids,
+            &timestamps,
+            &nonces,
+            &values,
+        )
+        .execute(&self.pgpool)
+        .await
+        .map_err(|e| {
+            error!("Failed to store receipt: {}", e);
+            anyhow!(e)
+        })?;
+
+        Ok(())
+    }
+}
+
+impl IndexerTapContext {
+    pub fn spawn_store_receipt_task(
+        inner_context: InnerContext,
+        mut receiver: Receiver<DatabaseReceipt>,
+        cancelation_token: CancellationToken,
+    ) -> JoinHandle<()> {
+        const BUFFER_SIZE: usize = 100;
+        tokio::spawn(async move {
+            loop {
+                let mut buffer = Vec::with_capacity(BUFFER_SIZE);
+                select! {
+                    biased;
+                    _ = receiver.recv_many(&mut buffer, BUFFER_SIZE) => {
+                        if let Err(e) = inner_context.store_receipts(buffer).await {
+                            error!("Failed to store receipts: {}", e);
+                        }
+                    }
+                    _ = cancelation_token.cancelled() => { break },
+                }
+            }
+        })
+    }
+}
+
 #[async_trait::async_trait]
 impl ReceiptStore for IndexerTapContext {
     type AdapterError = AdapterError;
@@ -21,38 +104,53 @@ impl ReceiptStore for IndexerTapContext {
         &self,
         receipt: ReceiptWithState<Checking>,
     ) -> Result<u64, Self::AdapterError> {
+        let db_receipt = DatabaseReceipt::from_receipt(receipt, &self.domain_separator)?;
+        self.receipt_producer.send(db_receipt).await.map_err(|e| {
+            error!("Failed to queue receipt for storage: {}", e);
+            anyhow!(e)
+        })?;
+
+        // We don't need receipt_ids
+        Ok(0)
+    }
+}
+
+pub struct DatabaseReceipt {
+    signer_address: String,
+    signature: Vec<u8>,
+    allocation_id: String,
+    timestamp_ns: BigDecimal,
+    nonce: BigDecimal,
+    value: BigDecimal,
+}
+
+impl DatabaseReceipt {
+    fn from_receipt(
+        receipt: ReceiptWithState<Checking>,
+        separator: &Eip712Domain,
+    ) -> anyhow::Result<Self> {
         let receipt = receipt.signed_receipt();
-        let allocation_id = receipt.message.allocation_id;
-        let encoded_signature = receipt.signature.as_bytes().to_vec();
+        let allocation_id = receipt.message.allocation_id.encode_hex();
+        let signature = receipt.signature.as_bytes().to_vec();
 
-        let receipt_signer = receipt
-            .recover_signer(self.domain_separator.as_ref())
+        let signer_address = receipt
+            .recover_signer(separator)
             .map_err(|e| {
                 error!("Failed to recover receipt signer: {}", e);
                 anyhow!(e)
-            })?;
-
-        // TODO: consider doing this in another async task to avoid slowing down the paid query flow.
-        sqlx::query!(
-            r#"
-                INSERT INTO scalar_tap_receipts (signer_address, signature, allocation_id, timestamp_ns, nonce, value)
-                VALUES ($1, $2, $3, $4, $5, $6)
-            "#,
-            receipt_signer.encode_hex(),
-            encoded_signature,
-            allocation_id.encode_hex(),
-            BigDecimal::from(receipt.message.timestamp_ns),
-            BigDecimal::from(receipt.message.nonce),
-            BigDecimal::from(BigInt::from(receipt.message.value)),
-        )
-        .execute(&self.pgpool)
-        .await
-        .map_err(|e| {
-            error!("Failed to store receipt: {}", e);
-            anyhow!(e)
-        })?;
+            })?
+            .encode_hex();
 
-        // We don't need receipt_ids
-        Ok(0)
+        let timestamp_ns = BigDecimal::from(receipt.message.timestamp_ns);
+        let nonce = BigDecimal::from(receipt.message.nonce);
+        let value = BigDecimal::from(BigInt::from(receipt.message.value));
+        Ok(Self {
+            allocation_id,
+            nonce,
+            signature,
+            signer_address,
+            timestamp_ns,
+            value,
+        })
     }
 }