make fallback page size a knob + allow falling back to page size 1 in the worst case (#43374)

GitOrigin-RevId: fdd1f12adc5ee8462053813a101b11f7618a3158

Author: gautamg795

crates/common/src/knobs.rs

@@ -868,6 +868,15 @@ pub static MYSQL_MAX_DYNAMIC_SMART_CHUNK_SIZE: LazyLock<usize> =
 pub static MYSQL_MAX_CHUNK_BYTES: LazyLock<usize> =
     LazyLock::new(|| env_config("MYSQL_MAX_CHUNK_BYTES", 10 << 20));
 
+/// Page size to fall back to when Vitess rejects large result sets.
+/// Vitess limits query results to 64MiB. As documents can be up to 1MiB (plus
+/// some overhead) and system documents can be larger still, we may need to fall
+/// back to a much smaller page size if we hit the limit while loading
+/// documents. If the fallback page size still fails, we'll try page size 1
+/// before giving up.
+pub static MYSQL_FALLBACK_PAGE_SIZE: LazyLock<u32> =
+    LazyLock::new(|| env_config("MYSQL_FALLBACK_PAGE_SIZE", 5));
+
 /// Timeout for all operations on MySQL connections, Vitess timeout is 20s so
 /// set lower than that
 pub static MYSQL_TIMEOUT: LazyLock<u64> = LazyLock::new(|| env_config("MYSQL_TIMEOUT_SECONDS", 19));

crates/mysql/src/lib.rs

@@ -55,6 +55,7 @@ use common::{
     },
     interval::Interval,
     knobs::{
+        MYSQL_FALLBACK_PAGE_SIZE,
         MYSQL_MAX_QUERY_BATCH_SIZE,
         MYSQL_MAX_QUERY_DYNAMIC_BATCH_SIZE,
         MYSQL_MIN_QUERY_BATCH_SIZE,
@@ -128,12 +129,6 @@ use crate::{
     },
 };
 
-// Vitess limits query results to 64MiB.
-// As documents can be up to 1MiB (plus some overhead) and system documents can
-// be larger still, we may need to fall back to a much smaller page size if we
-// hit the limit while loading documents.
-const FALLBACK_PAGE_SIZE: u32 = 5;
-
 #[derive(Clone, Debug)]
 pub struct MySqlInstanceName {
     raw: String,
@@ -764,18 +759,26 @@ impl<RT: Runtime> MySqlReader<RT> {
                             .message
                             .contains("trying to send message larger than max") =>
                 {
-                    if page_size == FALLBACK_PAGE_SIZE {
+                    if page_size == 1 {
                         anyhow::bail!(
-                            "Failed to load documents with fallback page size \
-                             `{FALLBACK_PAGE_SIZE}`: {}",
+                            "Failed to load documents with minimum page size `1`: {}",
+                            db_err.message,
+                        );
+                    }
+                    if page_size == *MYSQL_FALLBACK_PAGE_SIZE {
+                        tracing::warn!(
+                            "Falling back to page size `1` due to repeated server error: {}",
                             db_err.message
                         );
+                        page_size = 1;
+                    } else {
+                        tracing::warn!(
+                            "Falling back to page size `{}` due to server error: {}",
+                            *MYSQL_FALLBACK_PAGE_SIZE,
+                            db_err.message
+                        );
+                        page_size = *MYSQL_FALLBACK_PAGE_SIZE;
                     }
-                    tracing::warn!(
-                        "Falling back to page size `{FALLBACK_PAGE_SIZE}` due to server error: {}",
-                        db_err.message
-                    );
-                    page_size = FALLBACK_PAGE_SIZE;
                     continue;
                 },
                 Err(e) => Err(e),
@@ -931,19 +934,26 @@ impl<RT: Runtime> MySqlReader<RT> {
                                 .message
                                 .contains("trying to send message larger than max") =>
                     {
-                        if batch_size == FALLBACK_PAGE_SIZE {
+                        if batch_size == 1 {
                             anyhow::bail!(
-                                "Failed to load documents with fallback page size \
-                                 `{FALLBACK_PAGE_SIZE}`: {}",
+                                "Failed to load index rows with minimum page size `1`: {}",
                                 db_err.message
                             );
                         }
-                        tracing::warn!(
-                            "Falling back to page size `{FALLBACK_PAGE_SIZE}` due to server \
-                             error: {}",
-                            db_err.message
-                        );
-                        batch_size = FALLBACK_PAGE_SIZE;
+                        if batch_size == *MYSQL_FALLBACK_PAGE_SIZE {
+                            tracing::warn!(
+                                "Falling back to page size `1` due to repeated server error: {}",
+                                db_err.message
+                            );
+                            batch_size = 1;
+                        } else {
+                            tracing::warn!(
+                                "Falling back to page size `{}` due to server error: {}",
+                                *MYSQL_FALLBACK_PAGE_SIZE,
+                                db_err.message
+                            );
+                            batch_size = *MYSQL_FALLBACK_PAGE_SIZE;
+                        }
                         fallback = true;
                         continue;
                     },