[State Sync] Add backpressure to fast sync receiver.

config/src/config/state_sync_config.rs

@@ -136,7 +136,7 @@ impl Default for StateSyncDriverConfig {
             max_connection_deadline_secs: 10,
             max_consecutive_stream_notifications: 10,
             max_num_stream_timeouts: 12,
-            max_pending_data_chunks: 100,
+            max_pending_data_chunks: 50,
             max_pending_mempool_notifications: 100,
             max_stream_wait_time_ms: 5000,
             num_versions_to_skip_snapshot_sync: 100_000_000, // At 5k TPS, this allows a node to fail for about 6 hours.
@@ -219,9 +219,7 @@ pub struct DataStreamingServiceConfig {
     /// Maximum number of concurrent data client requests (per stream) for state keys/values.
     pub max_concurrent_state_requests: u64,
 
-    /// Maximum channel sizes for each data stream listener. If messages are not
-    /// consumed, they will be dropped (oldest messages first). The remaining
-    /// messages will be retrieved using FIFO ordering.
+    /// Maximum channel sizes for each data stream listener (per stream).
     pub max_data_stream_channel_sizes: u64,
 
     /// Maximum number of notification ID to response context mappings held in
@@ -256,7 +254,7 @@ impl Default for DataStreamingServiceConfig {
             global_summary_refresh_interval_ms: 50,
             max_concurrent_requests: MAX_CONCURRENT_REQUESTS,
             max_concurrent_state_requests: MAX_CONCURRENT_STATE_REQUESTS,
-            max_data_stream_channel_sizes: 300,
+            max_data_stream_channel_sizes: 50,
             max_notification_id_mappings: 300,
             max_num_consecutive_subscriptions: 40, // At ~4 blocks per second, this should last 10 seconds
             max_pending_requests: 50,
@@ -297,7 +295,7 @@ impl Default for DynamicPrefetchingConfig {
         Self {
             enable_dynamic_prefetching: true,
             initial_prefetching_value: 3,
-            max_prefetching_value: 50,
+            max_prefetching_value: 30,
             min_prefetching_value: 3,
             prefetching_value_increase: 1,
             prefetching_value_decrease: 2,

state-sync/data-streaming-service/src/data_stream.rs

@@ -393,8 +393,11 @@ impl<T: AptosDataClientInterface + Send + Clone + 'static> DataStream<T> {
         pending_client_response
     }
 
-    // TODO(joshlind): this function shouldn't be blocking when trying to send! If there are
-    // multiple streams, a single blocked stream could cause them all to block.
+    // TODO(joshlind): this function shouldn't be blocking when trying to send.
+    // If there are multiple streams, a single blocked stream could cause them
+    // all to block. This is acceptable for now (because there is only ever
+    // a single stream in use by the driver) but it should be fixed if we want
+    // to generalize this for multiple streams.
     async fn send_data_notification(
         &mut self,
         data_notification: DataNotification,

state-sync/state-sync-driver/src/bootstrapper.rs

@@ -1030,6 +1030,7 @@ impl<
         if let Err(error) = self
             .storage_synchronizer
             .save_state_values(notification_id, state_value_chunk_with_proof)
+            .await
         {
             self.reset_active_stream(Some(NotificationAndFeedback::new(
                 notification_id,

state-sync/state-sync-driver/src/storage_synchronizer.rs

@@ -96,7 +96,7 @@ pub trait StorageSynchronizerInterface {
     ///
     /// Note: this requires that `initialize_state_synchronizer` has been
     /// called.
-    fn save_state_values(
+    async fn save_state_values(
         &mut self,
         notification_id: NotificationId,
         state_value_chunk_with_proof: StateValueChunkWithProof,
@@ -403,17 +403,20 @@ impl<
         load_pending_data_chunks(self.pending_data_chunks.clone()) > 0
     }
 
-    fn save_state_values(
+    async fn save_state_values(
         &mut self,
         notification_id: NotificationId,
         state_value_chunk_with_proof: StateValueChunkWithProof,
     ) -> Result<(), Error> {
+        // Get the snapshot notifier and create the storage data chunk
         let state_snapshot_notifier = self.state_snapshot_notifier.as_mut().ok_or_else(|| {
             Error::UnexpectedError("The state snapshot receiver has not been initialized!".into())
         })?;
         let storage_data_chunk =
             StorageDataChunk::States(notification_id, state_value_chunk_with_proof);
-        if let Err(error) = state_snapshot_notifier.try_send(storage_data_chunk) {
+
+        // Notify the snapshot receiver of the storage data chunk
+        if let Err(error) = state_snapshot_notifier.send(storage_data_chunk).await {
             Err(Error::UnexpectedError(format!(
                 "Failed to send storage data chunk to state snapshot listener: {:?}",
                 error

state-sync/state-sync-driver/src/tests/mocks.rs

@@ -475,7 +475,7 @@ mock! {
 
         fn pending_storage_data(&self) -> bool;
 
-        fn save_state_values(
+        async fn save_state_values(
             &mut self,
             notification_id: NotificationId,
             state_value_chunk_with_proof: StateValueChunkWithProof,

state-sync/state-sync-driver/src/tests/storage_synchronizer.rs

@@ -750,9 +750,11 @@ async fn test_save_states_completion() {
     // Save multiple state chunks (including the last chunk)
     storage_synchronizer
         .save_state_values(0, create_state_value_chunk_with_proof(false))
+        .await
         .unwrap();
     storage_synchronizer
         .save_state_values(1, create_state_value_chunk_with_proof(true))
+        .await
         .unwrap();
 
     // Verify we get a commit notification
@@ -808,6 +810,7 @@ async fn test_save_states_dropped_error_listener() {
     let notification_id = 0;
     storage_synchronizer
         .save_state_values(notification_id, create_state_value_chunk_with_proof(true))
+        .await
         .unwrap();
 
     // The handler should panic as the commit listener was dropped
@@ -849,13 +852,14 @@ async fn test_save_states_invalid_chunk() {
     let notification_id = 0;
     storage_synchronizer
         .save_state_values(notification_id, create_state_value_chunk_with_proof(false))
+        .await
         .unwrap();
     verify_error_notification(&mut error_listener, notification_id).await;
 }
 
-#[test]
+#[tokio::test]
 #[should_panic]
-fn test_save_states_without_initialize() {
+async fn test_save_states_without_initialize() {
     // Create the storage synchronizer
     let (_, _, _, _, _, mut storage_synchronizer, _) = create_storage_synchronizer(
         create_mock_executor(),
@@ -864,7 +868,10 @@ fn test_save_states_without_initialize() {
 
     // Attempting to save the states should panic as the state
     // synchronizer was not initialized!
-    let _ = storage_synchronizer.save_state_values(0, create_state_value_chunk_with_proof(false));
+    storage_synchronizer
+        .save_state_values(0, create_state_value_chunk_with_proof(false))
+        .await
+        .unwrap();
 }
 
 /// Creates a storage synchronizer for testing