Merge branch 'master' into XT/sink-batchable

Author: hunterlxt

src/env.rs

@@ -38,6 +38,8 @@ fn poll_queue(tx: mpsc::Sender<CompletionQueue>) {
 pub struct EnvBuilder {
     cq_count: usize,
     name_prefix: Option<String>,
+    after_start: Option<Arc<dyn Fn() + Send + Sync>>,
+    before_stop: Option<Arc<dyn Fn() + Send + Sync>>,
 }
 
 impl EnvBuilder {
@@ -46,6 +48,8 @@ impl EnvBuilder {
         EnvBuilder {
             cq_count: unsafe { grpc_sys::gpr_cpu_num_cores() as usize },
             name_prefix: None,
+            after_start: None,
+            before_stop: None,
         }
     }
 
@@ -67,6 +71,18 @@ impl EnvBuilder {
         self
     }
 
+    /// Execute function `f` after each thread is started but before it starts doing work.
+    pub fn after_start<F: Fn() + Send + Sync + 'static>(mut self, f: F) -> EnvBuilder {
+        self.after_start = Some(Arc::new(f));
+        self
+    }
+
+    /// Execute function `f` before each thread stops.
+    pub fn before_stop<F: Fn() + Send + Sync + 'static>(mut self, f: F) -> EnvBuilder {
+        self.before_stop = Some(Arc::new(f));
+        self
+    }
+
     /// Finalize the [`EnvBuilder`], build the [`Environment`] and initialize the gRPC library.
     pub fn build(self) -> Environment {
         unsafe {
@@ -81,7 +97,19 @@ impl EnvBuilder {
             if let Some(ref prefix) = self.name_prefix {
                 builder = builder.name(format!("{}-{}", prefix, i));
             }
-            let handle = builder.spawn(move || poll_queue(tx_i)).unwrap();
+            let after_start = self.after_start.clone();
+            let before_stop = self.before_stop.clone();
+            let handle = builder
+                .spawn(move || {
+                    if let Some(f) = after_start {
+                        f();
+                    }
+                    poll_queue(tx_i);
+                    if let Some(f) = before_stop {
+                        f();
+                    }
+                })
+                .unwrap();
             handles.push(handle);
         }
         for _ in 0..self.cq_count {

tests-and-examples/tests/cases/misc.rs

@@ -27,7 +27,16 @@ impl Greeter for PeerService {
 
 #[test]
 fn test_peer() {
-    let env = Arc::new(EnvBuilder::new().build());
+    let counter_add = Arc::new(AtomicI32::new(0));
+    let counter_collect = counter_add.clone();
+    let env = Arc::new(
+        EnvBuilder::new()
+            .cq_count(2)
+            .after_start(move || {
+                counter_add.fetch_add(1, Ordering::Relaxed);
+            })
+            .build(),
+    );
     let service = create_greeter(PeerService);
     let mut server = ServerBuilder::new(env.clone())
         .register_service(service)
@@ -43,6 +52,7 @@ fn test_peer() {
     let resp = client.say_hello(&req).unwrap();
 
     assert!(resp.get_message().contains("127.0.0.1"), "{:?}", resp);
+    assert_eq!(counter_collect.load(Ordering::Relaxed), 2);
 }
 
 #[derive(Clone)]