add tests for sink dependent batcher configuration

Author: Wumpf

crates/store/re_chunk/src/batcher.rs

@@ -51,6 +51,12 @@ pub struct BatcherHooks {
     #[allow(clippy::type_complexity)]
     pub on_insert: Option<Arc<dyn Fn(&[PendingRow]) + Send + Sync>>,
 
+    /// Called when the batcher's configuration changes.
+    ///
+    /// Called for initial configuration as well as subsequent changes.
+    /// Used for testing.
+    pub on_config_change: Option<Arc<dyn Fn(&ChunkBatcherConfig) + Send + Sync>>,
+
     /// Callback to be run when an Arrow Chunk goes out of scope.
     ///
     /// See [`re_log_types::ArrowRecordBatchReleaseCallback`] for more information.
@@ -62,6 +68,7 @@ pub struct BatcherHooks {
 impl BatcherHooks {
     pub const NONE: Self = Self {
         on_insert: None,
+        on_config_change: None,
         on_release: None,
     };
 }
@@ -70,6 +77,7 @@ impl PartialEq for BatcherHooks {
     fn eq(&self, other: &Self) -> bool {
         let Self {
             on_insert,
+            on_config_change,
             on_release,
         } = self;
 
@@ -79,18 +87,26 @@ impl PartialEq for BatcherHooks {
             _ => false,
         };
 
-        on_insert_eq && on_release == &other.on_release
+        let on_config_change_eq = match (on_config_change, &other.on_config_change) {
+            (Some(a), Some(b)) => Arc::ptr_eq(a, b),
+            (None, None) => true,
+            _ => false,
+        };
+
+        on_insert_eq && on_config_change_eq && on_release == &other.on_release
     }
 }
 
 impl std::fmt::Debug for BatcherHooks {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         let Self {
             on_insert,
+            on_config_change,
             on_release,
         } = self;
         f.debug_struct("BatcherHooks")
             .field("on_insert", &on_insert.as_ref().map(|_| "…"))
+            .field("on_config_change", &on_config_change.as_ref().map(|_| "…"))
             .field("on_release", &on_release)
             .finish()
     }
@@ -617,6 +633,10 @@ fn batching_thread(
         config.flush_num_rows,
         re_format::format_bytes(config.flush_num_bytes as _),
     );
+    // Signal initial config
+    if let Some(on_config_change) = hooks.on_config_change.as_ref() {
+        on_config_change(&config);
+    }
 
     // Set to `true` when a flush is triggered for a reason other than hitting the time threshold,
     // so that the next tick will not unnecessarily fire early.
@@ -683,6 +703,9 @@ fn batching_thread(
                         }
 
                         re_log::trace!("Updated batcher config: {:?}", new_config);
+                        if let Some(on_config_change) = hooks.on_config_change.as_ref() {
+                            on_config_change(&new_config);
+                        }
 
                         config = new_config;
                         rx_tick = crossbeam::channel::tick(config.flush_tick);

crates/top/re_sdk/src/recording_stream.rs

@@ -3069,4 +3069,147 @@ mod tests {
         )
         .unwrap();
     }
+
+    struct BatcherConfigTestSink {
+        config: ChunkBatcherConfig,
+    }
+
+    impl LogSink for BatcherConfigTestSink {
+        fn default_batcher_config(&self) -> ChunkBatcherConfig {
+            self.config.clone()
+        }
+
+        fn send(&self, _msg: LogMsg) {
+            // noop
+        }
+
+        fn flush_blocking(&self) {
+            // noop
+        }
+
+        fn as_any(&self) -> &dyn std::any::Any {
+            self
+        }
+    }
+
+    struct ScopedEnvVarSet {
+        key: &'static str,
+    }
+
+    impl ScopedEnvVarSet {
+        #[allow(unsafe_code)]
+        fn new(key: &'static str, value: &'static str) -> Self {
+            unsafe { std::env::set_var(key, value) };
+            Self { key }
+        }
+    }
+
+    impl Drop for ScopedEnvVarSet {
+        #[allow(unsafe_code)]
+        fn drop(&mut self) {
+            unsafe {
+                std::env::remove_var(self.key);
+            }
+        }
+    }
+
+    const CONFIG_CHANGE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(1);
+
+    #[test]
+
+    fn test_sink_dependent_batcher_config() {
+        let (tx, rx) = std::sync::mpsc::channel();
+
+        let rec = RecordingStreamBuilder::new("rerun_example_test_batcher_config")
+            .batcher_hooks(BatcherHooks {
+                on_config_change: Some(Arc::new(move |config: &ChunkBatcherConfig| {
+                    tx.send(config.clone()).unwrap();
+                })),
+                ..BatcherHooks::NONE
+            })
+            .buffered()
+            .unwrap();
+
+        let new_config = rx
+            .recv_timeout(CONFIG_CHANGE_TIMEOUT)
+            .expect("no config change message received within timeout");
+        assert_eq!(new_config, ChunkBatcherConfig::DEFAULT); // buffered sink uses the default config.
+
+        // Change sink to our custom sink. Will it take over the setting?
+        let injected_config = ChunkBatcherConfig {
+            flush_tick: std::time::Duration::from_secs(123),
+            flush_num_bytes: 123,
+            flush_num_rows: 123,
+            ..ChunkBatcherConfig::DEFAULT
+        };
+        rec.set_sink(Box::new(BatcherConfigTestSink {
+            config: injected_config.clone(),
+        }));
+        let new_config = rx
+            .recv_timeout(CONFIG_CHANGE_TIMEOUT)
+            .expect("no config change message received within timeout");
+
+        assert_eq!(new_config, injected_config);
+
+        // Set flush num bytes through env var and set the sink again.
+        // check that the env var is respected.
+        let _scoped_env_guard = ScopedEnvVarSet::new("RERUN_FLUSH_NUM_BYTES", "456");
+        rec.set_sink(Box::new(BatcherConfigTestSink {
+            config: injected_config.clone(),
+        }));
+        let new_config = rx
+            .recv_timeout(CONFIG_CHANGE_TIMEOUT)
+            .expect("no config change message received within timeout");
+        assert_eq!(
+            new_config,
+            ChunkBatcherConfig {
+                flush_num_bytes: 456,
+                ..injected_config
+            },
+        );
+    }
+
+    #[test]
+    fn test_explicit_batcher_config() {
+        // This environment variable should still override the explicit config.
+        let _scoped_env_guard = ScopedEnvVarSet::new("RERUN_FLUSH_TICK_SECS", "456");
+        let explicit_config = ChunkBatcherConfig {
+            flush_tick: std::time::Duration::from_secs(123),
+            flush_num_bytes: 123,
+            flush_num_rows: 123,
+            ..ChunkBatcherConfig::DEFAULT
+        };
+        let expected_config = ChunkBatcherConfig {
+            flush_tick: std::time::Duration::from_secs(456),
+            ..explicit_config
+        };
+
+        let (tx, rx) = std::sync::mpsc::channel();
+        let rec = RecordingStreamBuilder::new("rerun_example_test_batcher_config")
+            .batcher_config(explicit_config)
+            .batcher_hooks(BatcherHooks {
+                on_config_change: Some(Arc::new(move |config: &ChunkBatcherConfig| {
+                    tx.send(config.clone()).unwrap();
+                })),
+                ..BatcherHooks::NONE
+            })
+            .buffered()
+            .unwrap();
+
+        let new_config = rx
+            .recv_timeout(CONFIG_CHANGE_TIMEOUT)
+            .expect("no config change message received within timeout");
+        assert_eq!(new_config, expected_config);
+
+        // Changing the sink should have no effect since an explicit config is in place.
+        rec.set_sink(Box::new(BatcherConfigTestSink {
+            config: ChunkBatcherConfig::ALWAYS,
+        }));
+        // Don't want to stall the test for CONFIG_CHANGE_TIMEOUT here.
+        let new_config_recv_result = rx.recv_timeout(std::time::Duration::from_millis(100));
+        assert_eq!(
+            new_config_recv_result,
+            Err(std::sync::mpsc::RecvTimeoutError::Timeout)
+        );
+    }
 }