Bigtable: update google proto files and allow configuration of max_me…

…ssage_size (#34740) * Update proto files with tonic-build v0.9.2 * Manually ignore invalid doc-tests * Add new ReadRowsRequest fields * Add LedgerStorageConfig::max_message_size and default value * Add BigtableConnection::max_message_size and use on client creation * Add max_message_size to RpcBigtableConfig and make const pub * Add solana-validator cli arg (cherry picked from commit 166be29)
solana-labs · Jan 11, 2024 · 0fb7b43 · 0fb7b43
1 parent 6af8e3b
commit 0fb7b43
Show file tree

Hide file tree

Showing 11 changed files with 1,777 additions and 116 deletions.
diff --git a/ledger-tool/src/bigtable.rs b/ledger-tool/src/bigtable.rs
@@ -585,6 +585,7 @@ async fn get_bigtable(
                 credential_type: CredentialType::Filepath(Some(args.crediential_path.unwrap())),
                 instance_name: args.instance_name,
                 app_profile_id: args.app_profile_id,
+                max_message_size: solana_storage_bigtable::DEFAULT_MAX_MESSAGE_SIZE,
             },
         )
         .await

diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs
@@ -169,6 +169,7 @@ pub struct RpcBigtableConfig {
     pub bigtable_instance_name: String,
     pub bigtable_app_profile_id: String,
     pub timeout: Option<Duration>,
+    pub max_message_size: usize,
 }
 
 impl Default for RpcBigtableConfig {
@@ -180,6 +181,7 @@ impl Default for RpcBigtableConfig {
             bigtable_instance_name,
             bigtable_app_profile_id,
             timeout: None,
+            max_message_size: solana_storage_bigtable::DEFAULT_MAX_MESSAGE_SIZE,
         }
     }
 }

diff --git a/rpc/src/rpc_service.rs b/rpc/src/rpc_service.rs
@@ -406,6 +406,7 @@ impl JsonRpcService {
                 ref bigtable_instance_name,
                 ref bigtable_app_profile_id,
                 timeout,
+                max_message_size,
             }) = config.rpc_bigtable_config
             {
                 let bigtable_config = solana_storage_bigtable::LedgerStorageConfig {
@@ -414,6 +415,7 @@ impl JsonRpcService {
                     credential_type: CredentialType::Filepath(None),
                     instance_name: bigtable_instance_name.clone(),
                     app_profile_id: bigtable_app_profile_id.clone(),
+                    max_message_size,
                 };
                 runtime
                     .block_on(solana_storage_bigtable::LedgerStorage::new_with_config(

diff --git a/storage-bigtable/proto/google.api.rs b/storage-bigtable/proto/google.api.rs
diff --git a/storage-bigtable/proto/google.bigtable.v2.rs b/storage-bigtable/proto/google.bigtable.v2.rs
diff --git a/storage-bigtable/proto/google.rpc.rs b/storage-bigtable/proto/google.rpc.rs
@@ -8,12 +8,14 @@
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Status {
-    /// The status code, which should be an enum value of \[google.rpc.Code][google.rpc.Code\].
+    /// The status code, which should be an enum value of
+    /// \[google.rpc.Code][google.rpc.Code\].
     #[prost(int32, tag = "1")]
     pub code: i32,
     /// A developer-facing error message, which should be in English. Any
     /// user-facing error message should be localized and sent in the
-    /// \[google.rpc.Status.details][google.rpc.Status.details\] field, or localized by the client.
+    /// \[google.rpc.Status.details][google.rpc.Status.details\] field, or localized
+    /// by the client.
     #[prost(string, tag = "2")]
     pub message: ::prost::alloc::string::String,
     /// A list of messages that carry the error details.  There is a common set of

diff --git a/storage-bigtable/src/bigtable.rs b/storage-bigtable/src/bigtable.rs
@@ -112,6 +112,7 @@ pub struct BigTableConnection {
     table_prefix: String,
     app_profile_id: String,
     timeout: Option<Duration>,
+    max_message_size: usize,
 }
 
 impl BigTableConnection {
@@ -132,11 +133,18 @@ impl BigTableConnection {
         read_only: bool,
         timeout: Option<Duration>,
         credential_type: CredentialType,
+        max_message_size: usize,
     ) -> Result<Self> {
         match std::env::var("BIGTABLE_EMULATOR_HOST") {
             Ok(endpoint) => {
                 info!("Connecting to bigtable emulator at {}", endpoint);
-                Self::new_for_emulator(instance_name, app_profile_id, &endpoint, timeout)
+                Self::new_for_emulator(
+                    instance_name,
+                    app_profile_id,
+                    &endpoint,
+                    timeout,
+                    max_message_size,
+                )
             }
 
             Err(_) => {
@@ -201,6 +209,7 @@ impl BigTableConnection {
                     table_prefix,
                     app_profile_id: app_profile_id.to_string(),
                     timeout,
+                    max_message_size,
                 })
             }
         }
@@ -211,6 +220,7 @@ impl BigTableConnection {
         app_profile_id: &str,
         endpoint: &str,
         timeout: Option<Duration>,
+        max_message_size: usize,
     ) -> Result<Self> {
         Ok(Self {
             access_token: None,
@@ -220,6 +230,7 @@ impl BigTableConnection {
             table_prefix: format!("projects/emulator/instances/{instance_name}/tables/"),
             app_profile_id: app_profile_id.to_string(),
             timeout,
+            max_message_size,
         })
     }
 
@@ -245,7 +256,9 @@ impl BigTableConnection {
                 }
                 Ok(req)
             },
-        );
+        )
+        .max_decoding_message_size(self.max_message_size)
+        .max_encoding_message_size(self.max_message_size);
         BigTable {
             access_token: self.access_token.clone(),
             client,
@@ -458,6 +471,8 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
                         ],
                     })),
                 }),
+                request_stats_view: 0,
+                reversed: false,
             })
             .await?
             .into_inner();
@@ -483,6 +498,8 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
                 filter: Some(RowFilter {
                     filter: Some(row_filter::Filter::StripValueTransformer(true)),
                 }),
+                request_stats_view: 0,
+                reversed: false,
             })
             .await?
             .into_inner();
@@ -534,6 +551,8 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
                     // Only return the latest version of each cell
                     filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
                 }),
+                request_stats_view: 0,
+                reversed: false,
             })
             .await?
             .into_inner();
@@ -566,6 +585,8 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
                     // Only return the latest version of each cell
                     filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
                 }),
+                request_stats_view: 0,
+                reversed: false,
             })
             .await?
             .into_inner();
@@ -599,6 +620,8 @@ impl<F: FnMut(Request<()>) -> InterceptedRequestResult> BigTable<F> {
                     // Only return the latest version of each cell
                     filter: Some(row_filter::Filter::CellsPerColumnLimitFilter(1)),
                 }),
+                request_stats_view: 0,
+                reversed: false,
             })
             .await?
             .into_inner();

diff --git a/storage-bigtable/src/lib.rs b/storage-bigtable/src/lib.rs
@@ -376,6 +376,7 @@ impl From<LegacyTransactionByAddrInfo> for TransactionByAddrInfo {
 
 pub const DEFAULT_INSTANCE_NAME: &str = "solana-ledger";
 pub const DEFAULT_APP_PROFILE_ID: &str = "default";
+pub const DEFAULT_MAX_MESSAGE_SIZE: usize = 64 * 1024 * 1024; // 64MB
 
 #[derive(Debug)]
 pub enum CredentialType {
@@ -390,6 +391,7 @@ pub struct LedgerStorageConfig {
     pub credential_type: CredentialType,
     pub instance_name: String,
     pub app_profile_id: String,
+    pub max_message_size: usize,
 }
 
 impl Default for LedgerStorageConfig {
@@ -400,6 +402,7 @@ impl Default for LedgerStorageConfig {
             credential_type: CredentialType::Filepath(None),
             instance_name: DEFAULT_INSTANCE_NAME.to_string(),
             app_profile_id: DEFAULT_APP_PROFILE_ID.to_string(),
+            max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
         }
     }
 }
@@ -466,6 +469,7 @@ impl LedgerStorage {
                 app_profile_id,
                 endpoint,
                 timeout,
+                LedgerStorageConfig::default().max_message_size,
             )?,
             stats,
         })
@@ -479,13 +483,15 @@ impl LedgerStorage {
             instance_name,
             app_profile_id,
             credential_type,
+            max_message_size,
         } = config;
         let connection = bigtable::BigTableConnection::new(
             instance_name.as_str(),
             app_profile_id.as_str(),
             read_only,
             timeout,
             credential_type,
+            max_message_size,
         )
         .await?;
         Ok(Self { stats, connection })

diff --git a/validator/src/bin/solana-test-validator.rs b/validator/src/bin/solana-test-validator.rs
@@ -430,6 +430,7 @@ fn main() {
                 String
             ),
             timeout: None,
+            ..RpcBigtableConfig::default()
         })
     } else {
         None

diff --git a/validator/src/cli.rs b/validator/src/cli.rs
@@ -878,6 +878,15 @@ pub fn app<'a>(version: &'a str, default_args: &'a DefaultArgs) -> App<'a, 'a> {
                 .default_value(&default_args.rpc_bigtable_app_profile_id)
                 .help("Bigtable application profile id to use in requests")
         )
+        .arg(
+            Arg::with_name("rpc_bigtable_max_message_size")
+                .long("rpc-bigtable-max-message-size")
+                .value_name("BYTES")
+                .validator(is_parsable::<usize>)
+                .takes_value(true)
+                .default_value(&default_args.rpc_bigtable_max_message_size)
+                .help("Max encoding and decoding message size used in Bigtable Grpc client"),
+        )
         .arg(
             Arg::with_name("rpc_pubsub_worker_threads")
                 .long("rpc-pubsub-worker-threads")
@@ -1908,6 +1917,7 @@ pub struct DefaultArgs {
     pub rpc_bigtable_timeout: String,
     pub rpc_bigtable_instance_name: String,
     pub rpc_bigtable_app_profile_id: String,
+    pub rpc_bigtable_max_message_size: String,
     pub rpc_max_request_body_size: String,
     pub rpc_pubsub_worker_threads: String,
 
@@ -1993,6 +2003,8 @@ impl DefaultArgs {
             rpc_bigtable_instance_name: solana_storage_bigtable::DEFAULT_INSTANCE_NAME.to_string(),
             rpc_bigtable_app_profile_id: solana_storage_bigtable::DEFAULT_APP_PROFILE_ID
                 .to_string(),
+            rpc_bigtable_max_message_size: solana_storage_bigtable::DEFAULT_MAX_MESSAGE_SIZE
+                .to_string(),
             rpc_pubsub_worker_threads: "4".to_string(),
             accountsdb_repl_threads: num_cpus::get().to_string(),
             accounts_filler_count: "0".to_string(),

diff --git a/validator/src/main.rs b/validator/src/main.rs
@@ -1237,6 +1237,7 @@ pub fn main() {
             timeout: value_t!(matches, "rpc_bigtable_timeout", u64)
                 .ok()
                 .map(Duration::from_secs),
+            max_message_size: value_t_or_exit!(matches, "rpc_bigtable_max_message_size", usize),
         })
     } else {
         None
-Original file line number
+Diff line change
@@ Expand Up / @@ -430,6 +430,7 @@ fn main() { @@
                     String
                 ),
                 timeout: None,
+                ..RpcBigtableConfig::default()
             })
         } else {
             None
@@ Expand Down @@