Explorar el Código

Chunk id-batch posting primitives to a safe statement size

The id-batch primitives (get_postings, get_posting_states, reserve,
release, deactivate) matched every id in one statement via an OR of
equality pairs. That OR chain grows with the batch, and SQLite rejects it
once the expression tree passes its depth limit (default 1000), well
before the bind-parameter limits are reached; PostgreSQL has its own
ceilings too. A caller passing a large id set would hit a hard database
error.

Split every id batch into fixed-size chunks, keeping the two write
primitives' chunks inside their existing single transaction so the claim
or release stays atomic, and summing the affected-row counts. Reads
accumulate across chunks. The batch size the primitives accept now has no
practical ceiling.

Also encode hex without a per-byte allocation, and add a test that drives
a batch larger than one chunk through reserve, read, and deactivate.
Cesar Rodas hace 1 semana
padre
commit
1d2cd7fc82
Se han modificado 2 ficheros con 277 adiciones y 176 borrados
  1. 223 176
      crates/kuatia-storage-sql/src/lib.rs
  2. 54 0
      crates/kuatia-storage-sql/tests/sqlite.rs

+ 223 - 176
crates/kuatia-storage-sql/src/lib.rs

@@ -189,9 +189,11 @@ fn deserialize_json<T: serde::de::DeserializeOwned>(s: &str) -> Result<T, StoreE
 /// opaque saga bytes are stored as hex `TEXT` so a row is legible in any SQL
 /// client and matches the hex form used in logs and `Debug` output.
 fn to_hex(bytes: &[u8]) -> String {
+    const HEX: &[u8; 16] = b"0123456789abcdef";
     let mut s = String::with_capacity(bytes.len() * 2);
-    for b in bytes {
-        s.push_str(&format!("{b:02x}"));
+    for &b in bytes {
+        s.push(HEX[(b >> 4) as usize] as char);
+        s.push(HEX[(b & 0x0f) as usize] as char);
     }
     s
 }
@@ -304,12 +306,23 @@ fn filter_source(filter: PostingFilter) -> &'static str {
     }
 }
 
+/// Maximum posting ids matched by a single statement. `id_predicate` expands to
+/// an `OR` of `n` equality pairs, so the binding constraint is SQLite's
+/// expression-tree depth limit (`SQLITE_MAX_EXPR_DEPTH`, default 1000), which a
+/// chain of `n` `OR`s reaches at roughly `n` deep. It caps well before the
+/// bind-parameter limits (SQLite 32766, PostgreSQL 65535) that `2 * n (+1)`
+/// parameters would hit. `500` stays comfortably under the expression-depth
+/// limit; callers that pass more ids are chunked, so the id-batch primitives
+/// have no practical ceiling on batch size.
+const MAX_IDS_PER_QUERY: usize = 500;
+
 /// Build a portable predicate matching a set of posting ids:
 /// `(transfer_id = $s AND idx = $s+1) OR (transfer_id = $s+2 AND idx = $s+3) ...`
 /// starting at placeholder `$start`. Row-value `IN ((a, b), ...)` is not
 /// portable across SQLite and PostgreSQL; an `OR` of equality pairs is. The
 /// caller binds each id as `(hex(transfer), idx as i16)` in order, matching the
-/// placeholder sequence. `ids` must be non-empty.
+/// placeholder sequence. `ids` must be non-empty and no longer than
+/// [`MAX_IDS_PER_QUERY`]; larger sets are split into chunks by the caller.
 fn id_predicate(count: usize, start: u32) -> String {
     (0..count)
         .map(|i| {
@@ -540,33 +553,36 @@ impl PostingStore for SqlStore {
             return Ok(Vec::new());
         }
 
-        // One set-based query for the whole batch instead of one probe per id,
-        // reusing the portable `id_predicate` and binding each id in order as
-        // `(hex(transfer), idx as i16)`.
-        let sql = format!(
-            "SELECT * FROM postings WHERE {}",
-            id_predicate(ids.len(), 1)
-        );
-        let mut q = sqlx::query(&sql);
-        for id in ids {
-            q = q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
-        }
-        let rows = q
-            .fetch_all(&self.pool)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
-
-        // Index the fetched postings by the same `(hex, idx)` key that was bound.
-        let mut found: HashMap<(String, i16), Posting> = HashMap::with_capacity(rows.len());
-        for row in &rows {
-            let posting = row_to_posting(row)?;
-            let key = (
-                envelope_id_to_hex(&posting.id.transfer),
-                posting.id.index as i16,
+        // Set-based query per chunk instead of one probe per id, reusing the
+        // portable `id_predicate` and binding each id in order as
+        // `(hex(transfer), idx as i16)`. Chunked so a large batch never exceeds
+        // the backend's bind-parameter limit (see `MAX_IDS_PER_QUERY`).
+        let mut found: HashMap<(String, i16), Posting> = HashMap::with_capacity(ids.len());
+        for chunk in ids.chunks(MAX_IDS_PER_QUERY) {
+            let sql = format!(
+                "SELECT * FROM postings WHERE {}",
+                id_predicate(chunk.len(), 1)
             );
-            found.insert(key, posting);
+            let mut q = sqlx::query(&sql);
+            for id in chunk {
+                q = q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            let rows = q
+                .fetch_all(&self.pool)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
+
+            // Index the fetched postings by the same `(hex, idx)` key that was bound.
+            for row in &rows {
+                let posting = row_to_posting(row)?;
+                let key = (
+                    envelope_id_to_hex(&posting.id.transfer),
+                    posting.id.index as i16,
+                );
+                found.insert(key, posting);
+            }
         }
 
         // Return in input order, erroring on the first id absent from the batch
@@ -628,48 +644,10 @@ impl PostingStore for SqlStore {
         }
 
         // One set-based query per state table instead of up to three probes per
-        // id. Each reuses the portable `id_predicate` (an OR of equality pairs;
-        // row-value `IN` is not portable across SQLite and PostgreSQL) and binds
-        // every id in order as `(hex(transfer), idx as i16)`.
-        let predicate = id_predicate(ids.len(), 1);
-
-        let active_sql = format!("SELECT transfer_id, idx FROM active_postings WHERE {predicate}");
-        let mut active_q = sqlx::query(&active_sql);
-        for id in ids {
-            active_q = active_q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
-        }
-        let active_rows = active_q
-            .fetch_all(&self.pool)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
-
-        let reserved_sql = format!(
-            "SELECT transfer_id, idx, reservation FROM reserved_postings WHERE {predicate}"
-        );
-        let mut reserved_q = sqlx::query(&reserved_sql);
-        for id in ids {
-            reserved_q = reserved_q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
-        }
-        let reserved_rows = reserved_q
-            .fetch_all(&self.pool)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
-
-        let spent_sql = format!("SELECT transfer_id, idx FROM postings WHERE {predicate}");
-        let mut spent_q = sqlx::query(&spent_sql);
-        for id in ids {
-            spent_q = spent_q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
-        }
-        let spent_rows = spent_q
-            .fetch_all(&self.pool)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
+        // id, reusing the portable `id_predicate` (an OR of equality pairs;
+        // row-value `IN` is not portable across SQLite and PostgreSQL) and
+        // binding every id in order as `(hex(transfer), idx as i16)`. Chunked so
+        // a large batch never exceeds the bind-parameter limit.
 
         // Key membership by the same `(hex, idx)` values that were bound, so the
         // per-id lookup below matches without decoding transfer ids back.
@@ -683,20 +661,63 @@ impl PostingStore for SqlStore {
             Ok((transfer_id, idx))
         };
 
-        let mut active: HashSet<(String, i16)> = HashSet::with_capacity(active_rows.len());
-        for row in &active_rows {
-            active.insert(row_key(row)?);
-        }
-        let mut reserved: HashMap<(String, i16), i64> = HashMap::with_capacity(reserved_rows.len());
-        for row in &reserved_rows {
-            let rid: i64 = row
-                .try_get("reservation")
+        let mut active: HashSet<(String, i16)> = HashSet::new();
+        let mut reserved: HashMap<(String, i16), i64> = HashMap::new();
+        let mut spent: HashSet<(String, i16)> = HashSet::new();
+
+        for chunk in ids.chunks(MAX_IDS_PER_QUERY) {
+            let predicate = id_predicate(chunk.len(), 1);
+
+            let active_sql =
+                format!("SELECT transfer_id, idx FROM active_postings WHERE {predicate}");
+            let mut active_q = sqlx::query(&active_sql);
+            for id in chunk {
+                active_q = active_q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            let active_rows = active_q
+                .fetch_all(&self.pool)
+                .await
                 .map_err(|e| StoreError::Internal(e.to_string()))?;
-            reserved.insert(row_key(row)?, rid);
-        }
-        let mut spent: HashSet<(String, i16)> = HashSet::with_capacity(spent_rows.len());
-        for row in &spent_rows {
-            spent.insert(row_key(row)?);
+            for row in &active_rows {
+                active.insert(row_key(row)?);
+            }
+
+            let reserved_sql = format!(
+                "SELECT transfer_id, idx, reservation FROM reserved_postings WHERE {predicate}"
+            );
+            let mut reserved_q = sqlx::query(&reserved_sql);
+            for id in chunk {
+                reserved_q = reserved_q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            let reserved_rows = reserved_q
+                .fetch_all(&self.pool)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
+            for row in &reserved_rows {
+                let rid: i64 = row
+                    .try_get("reservation")
+                    .map_err(|e| StoreError::Internal(e.to_string()))?;
+                reserved.insert(row_key(row)?, rid);
+            }
+
+            let spent_sql = format!("SELECT transfer_id, idx FROM postings WHERE {predicate}");
+            let mut spent_q = sqlx::query(&spent_sql);
+            for id in chunk {
+                spent_q = spent_q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            let spent_rows = spent_q
+                .fetch_all(&self.pool)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
+            for row in &spent_rows {
+                spent.insert(row_key(row)?);
+            }
         }
 
         // Reconstruct each id's state in input order, preserving the active >
@@ -799,43 +820,49 @@ impl PostingStore for SqlStore {
             .await
             .map_err(|e| StoreError::Internal(e.to_string()))?;
 
-        // Reservation is $1; each id pair follows starting at $2.
-        let insert_sql = format!(
-            "INSERT INTO reserved_postings (transfer_id, idx, owner, subaccount, asset, value, reservation) \
-             SELECT transfer_id, idx, owner, subaccount, asset, value, $1 FROM active_postings WHERE {} \
-             ON CONFLICT (transfer_id, idx) DO NOTHING",
-            id_predicate(ids.len(), 2)
-        );
-        let mut insert_q = sqlx::query(&insert_sql).bind(reservation.0);
-        for id in ids {
-            insert_q = insert_q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
-        }
-        insert_q
-            .execute(&mut *tx)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
+        // Chunked so a large id set stays under the bind-parameter limit; all
+        // chunks share one transaction so the whole claim is atomic.
+        let mut claimed: u64 = 0;
+        for chunk in ids.chunks(MAX_IDS_PER_QUERY) {
+            // Reservation is $1; each id pair follows starting at $2.
+            let insert_sql = format!(
+                "INSERT INTO reserved_postings (transfer_id, idx, owner, subaccount, asset, value, reservation) \
+                 SELECT transfer_id, idx, owner, subaccount, asset, value, $1 FROM active_postings WHERE {} \
+                 ON CONFLICT (transfer_id, idx) DO NOTHING",
+                id_predicate(chunk.len(), 2)
+            );
+            let mut insert_q = sqlx::query(&insert_sql).bind(reservation.0);
+            for id in chunk {
+                insert_q = insert_q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            insert_q
+                .execute(&mut *tx)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
 
-        let delete_sql = format!(
-            "DELETE FROM active_postings WHERE {}",
-            id_predicate(ids.len(), 1)
-        );
-        let mut delete_q = sqlx::query(&delete_sql);
-        for id in ids {
-            delete_q = delete_q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
+            let delete_sql = format!(
+                "DELETE FROM active_postings WHERE {}",
+                id_predicate(chunk.len(), 1)
+            );
+            let mut delete_q = sqlx::query(&delete_sql);
+            for id in chunk {
+                delete_q = delete_q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            let del = delete_q
+                .execute(&mut *tx)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
+            claimed += del.rows_affected();
         }
-        let del = delete_q
-            .execute(&mut *tx)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
 
         tx.commit()
             .await
             .map_err(|e| StoreError::Internal(e.to_string()))?;
-        Ok(del.rows_affected())
+        Ok(claimed)
     }
 
     async fn release_postings(
@@ -856,43 +883,49 @@ impl PostingStore for SqlStore {
             .await
             .map_err(|e| StoreError::Internal(e.to_string()))?;
 
-        // Reservation is $1; each id pair follows starting at $2.
-        let insert_sql = format!(
-            "INSERT INTO active_postings (transfer_id, idx, owner, subaccount, asset, value) \
-             SELECT transfer_id, idx, owner, subaccount, asset, value FROM reserved_postings \
-             WHERE ({}) AND reservation = $1 ON CONFLICT (transfer_id, idx) DO NOTHING",
-            id_predicate(ids.len(), 2)
-        );
-        let mut insert_q = sqlx::query(&insert_sql).bind(reservation.0);
-        for id in ids {
-            insert_q = insert_q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
-        }
-        insert_q
-            .execute(&mut *tx)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
+        // Chunked so a large id set stays under the bind-parameter limit; all
+        // chunks share one transaction.
+        let mut released: u64 = 0;
+        for chunk in ids.chunks(MAX_IDS_PER_QUERY) {
+            // Reservation is $1; each id pair follows starting at $2.
+            let insert_sql = format!(
+                "INSERT INTO active_postings (transfer_id, idx, owner, subaccount, asset, value) \
+                 SELECT transfer_id, idx, owner, subaccount, asset, value FROM reserved_postings \
+                 WHERE ({}) AND reservation = $1 ON CONFLICT (transfer_id, idx) DO NOTHING",
+                id_predicate(chunk.len(), 2)
+            );
+            let mut insert_q = sqlx::query(&insert_sql).bind(reservation.0);
+            for id in chunk {
+                insert_q = insert_q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            insert_q
+                .execute(&mut *tx)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
 
-        let delete_sql = format!(
-            "DELETE FROM reserved_postings WHERE ({}) AND reservation = $1",
-            id_predicate(ids.len(), 2)
-        );
-        let mut delete_q = sqlx::query(&delete_sql).bind(reservation.0);
-        for id in ids {
-            delete_q = delete_q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
+            let delete_sql = format!(
+                "DELETE FROM reserved_postings WHERE ({}) AND reservation = $1",
+                id_predicate(chunk.len(), 2)
+            );
+            let mut delete_q = sqlx::query(&delete_sql).bind(reservation.0);
+            for id in chunk {
+                delete_q = delete_q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            let del = delete_q
+                .execute(&mut *tx)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
+            released += del.rows_affected();
         }
-        let del = delete_q
-            .execute(&mut *tx)
-            .await
-            .map_err(|e| StoreError::Internal(e.to_string()))?;
 
         tx.commit()
             .await
             .map_err(|e| StoreError::Internal(e.to_string()))?;
-        Ok(del.rows_affected())
+        Ok(released)
     }
 
     async fn deactivate_postings(
@@ -900,44 +933,58 @@ impl PostingStore for SqlStore {
         ids: &[PostingId],
         reservation: Option<ReservationId>,
     ) -> Result<u64, StoreError> {
-        // Dumb instruction over the whole id set: a single DELETE removes the
-        // ids from an index so they become spent (present only in the immutable
-        // table). `rows_affected` is the count; the caller interprets a shortfall.
+        // Dumb instruction over the whole id set: a DELETE removes the ids from
+        // an index so they become spent (present only in the immutable table).
+        // `rows_affected` is the count; the caller interprets a shortfall.
+        // Chunked under one transaction so a large id set stays within the
+        // bind-parameter limit while the removal stays atomic.
         if ids.is_empty() {
             return Ok(0);
         }
-        let (sql, rid) = match reservation {
-            // Raw path: remove from the active index.
-            None => (
-                format!(
-                    "DELETE FROM active_postings WHERE {}",
-                    id_predicate(ids.len(), 1)
+        let mut tx = self
+            .pool
+            .begin()
+            .await
+            .map_err(|e| StoreError::Internal(e.to_string()))?;
+        let mut removed: u64 = 0;
+        for chunk in ids.chunks(MAX_IDS_PER_QUERY) {
+            let (sql, rid) = match reservation {
+                // Raw path: remove from the active index.
+                None => (
+                    format!(
+                        "DELETE FROM active_postings WHERE {}",
+                        id_predicate(chunk.len(), 1)
+                    ),
+                    None,
                 ),
-                None,
-            ),
-            // Saga path: remove only the rows reserved by `rid`.
-            Some(rid) => (
-                format!(
-                    "DELETE FROM reserved_postings WHERE ({}) AND reservation = $1",
-                    id_predicate(ids.len(), 2)
+                // Saga path: remove only the rows reserved by `rid`.
+                Some(rid) => (
+                    format!(
+                        "DELETE FROM reserved_postings WHERE ({}) AND reservation = $1",
+                        id_predicate(chunk.len(), 2)
+                    ),
+                    Some(rid),
                 ),
-                Some(rid),
-            ),
-        };
-        let mut q = sqlx::query(&sql);
-        if let Some(rid) = rid {
-            q = q.bind(rid.0);
-        }
-        for id in ids {
-            q = q
-                .bind(envelope_id_to_hex(&id.transfer))
-                .bind(id.index as i16);
+            };
+            let mut q = sqlx::query(&sql);
+            if let Some(rid) = rid {
+                q = q.bind(rid.0);
+            }
+            for id in chunk {
+                q = q
+                    .bind(envelope_id_to_hex(&id.transfer))
+                    .bind(id.index as i16);
+            }
+            let res = q
+                .execute(&mut *tx)
+                .await
+                .map_err(|e| StoreError::Internal(e.to_string()))?;
+            removed += res.rows_affected();
         }
-        let res = q
-            .execute(&self.pool)
+        tx.commit()
             .await
             .map_err(|e| StoreError::Internal(e.to_string()))?;
-        Ok(res.rows_affected())
+        Ok(removed)
     }
 
     async fn insert_postings(&self, postings: &[Posting]) -> Result<u64, StoreError> {

+ 54 - 0
crates/kuatia-storage-sql/tests/sqlite.rs

@@ -350,3 +350,57 @@ async fn migrate_is_idempotent() {
     store.migrate().await.unwrap();
     store.migrate().await.unwrap();
 }
+
+/// The id-batch primitives chunk a large id set so it never exceeds the
+/// backend's bind-parameter limit (SQLite caps a statement at 32766 variables).
+/// A batch larger than one chunk must reserve, read, and deactivate every id
+/// exactly as a small one does.
+#[tokio::test]
+async fn id_batch_primitives_chunk_large_batches() {
+    // Comfortably larger than the internal chunk size so the batch spans
+    // multiple statements.
+    const N: u16 = 9000;
+    let store = new_store().await;
+
+    let tid = EnvelopeId([0xcd; 32]);
+    let postings: Vec<Posting> = (0..N)
+        .map(|i| {
+            Posting::new(
+                PostingId {
+                    transfer: tid,
+                    index: i,
+                },
+                AccountId::new(1),
+                AssetId::new(1),
+                Cent::from(1),
+            )
+        })
+        .collect();
+    let ids: Vec<PostingId> = postings.iter().map(|p| p.id).collect();
+
+    assert_eq!(
+        store.insert_postings(&postings).await.unwrap(),
+        u64::from(N)
+    );
+
+    // Reserve the whole batch in one call: every id is claimed across chunks.
+    let rid = ReservationId::new(1);
+    assert_eq!(
+        store.reserve_postings(&ids, rid).await.unwrap(),
+        u64::from(N)
+    );
+
+    // Reads span chunks and return every id.
+    let states = store.get_posting_states(&ids).await.unwrap();
+    assert_eq!(states.len(), N as usize);
+    assert!(states.iter().all(|s| *s == PostingState::Reserved(rid)));
+    assert_eq!(store.get_postings(&ids).await.unwrap().len(), N as usize);
+
+    // Deactivating the whole batch spends every id.
+    assert_eq!(
+        store.deactivate_postings(&ids, Some(rid)).await.unwrap(),
+        u64::from(N)
+    );
+    let after = store.get_posting_states(&ids).await.unwrap();
+    assert!(after.iter().all(|s| *s == PostingState::Spent));
+}