From 6b4fdb33593a113be48a78e6766758b48683a15f Mon Sep 17 00:00:00 2001
From: Joseph Koshakow <koshy44@gmail.com>
Date: Thu, 8 Feb 2024 09:53:40 -0500
Subject: [PATCH] stash: Speed up large batch test (#25055)

test_stash_batch_large_number_updates was added in
dad4ed8c7124f5a5dd2e47f3597c7f98004407a0 to test the max number of args
in a CRDB update statement. Previously, we were testing a number of
updates that was order of magnitudes larger than the max number of args
in a CRDB update statement. That caused the test to be slow and
sometimes hang forever. This commit updates the test to lower the number
of updates while still ensuring that we end up with more than the max
number of args.

Works towards resolving #25003
---
 src/stash/src/tests.rs       | 8 +++++++-
 src/stash/src/transaction.rs | 2 +-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/stash/src/tests.rs b/src/stash/src/tests.rs
index f23d1892f081c..50879a0d0bca3 100644
--- a/src/stash/src/tests.rs
+++ b/src/stash/src/tests.rs
@@ -11,6 +11,7 @@ use std::collections::{BTreeMap, BTreeSet};
 use std::convert::Infallible;
 use std::time::Duration;
 
+use crate::transaction::MAX_INSERT_ARGUMENTS;
 use crate::{
     AppendBatch, Data, DebugStashFactory, Stash, StashCollection, StashError, StashFactory,
     TableTransaction, Timestamp, TypedCollection, INSERT_BATCH_SPLIT_SIZE,
@@ -570,7 +571,12 @@ async fn test_stash_batch_large_number_updates() {
     Stash::with_debug_stash(|mut stash| async move {
         let col = collection::<i64, i64>(&mut stash, "c1").await.unwrap();
         let mut batch = make_batch(&col, &mut stash).await.unwrap();
-        for i in 0..500_000 {
+        // Back of the envelope math would produce 12 batches of updates.
+        //
+        // Each update statement takes 4 arguments, so we have a total of
+        // MAX_INSERT_ARGUMENTS * 4 * 3 arguments, leading to 12 batches of updates.
+        for i in 0..(MAX_INSERT_ARGUMENTS * 3) {
+            let i = i.into();
             col.append_to_batch(&mut batch, &i, &(i + 1), 1);
         }
         append(&mut stash, vec![batch]).await.unwrap();
diff --git a/src/stash/src/transaction.rs b/src/stash/src/transaction.rs
index cd2f42dc661de..f35d8827a2155 100644
--- a/src/stash/src/transaction.rs
+++ b/src/stash/src/transaction.rs
@@ -43,7 +43,7 @@ pub const INSERT_BATCH_SPLIT_SIZE: usize = 2 * 1024 * 1024;
 /// [`tokio_postgres`] has a maximum number of arguments it supports when executing a query. This
 /// is the limit at which to split a batch to make sure we don't try to include too many elements
 /// in any one update.
-const MAX_INSERT_ARGUMENTS: u16 = u16::MAX / 4;
+pub(crate) const MAX_INSERT_ARGUMENTS: u16 = u16::MAX / 4;
 
 impl Stash {
     /// Transactionally executes closure `f`.