Pick Add blocking callback to EnvBuilder (#474) (#475)

Signed-off-by: Xintao <[email protected]>

Author: hunterlxt

grpc-sys/build.rs

@@ -381,9 +381,9 @@ fn main() {
     };
 
     if get_env("CARGO_CFG_TARGET_OS").map_or(false, |s| s == "windows") {
-        // At lease win7
-        cc.define("_WIN32_WINNT", Some("0x0700"));
-        bind_config = bind_config.clang_arg("-D _WIN32_WINNT=0x0700");
+        // At lease vista
+        cc.define("_WIN32_WINNT", Some("0x600"));
+        bind_config = bind_config.clang_arg("-D _WIN32_WINNT=0x600");
     }
 
     if get_env("GRPCIO_SYS_USE_PKG_CONFIG").map_or(false, |s| s == "1") {

proto/Cargo.toml

@@ -27,6 +27,6 @@ protobuf = "2"
 lazy_static = { version = "1.3", optional = true }
 
 [build-dependencies]
-protobuf-build = { version = "0.11", default-features = false }
+protobuf-build = { version = "=0.11.3", default-features = false }
 grpcio-compiler = { path = "../compiler", version = "0.5.0", default-features = false }
 walkdir = "2.2"

src/call/server.rs

@@ -317,7 +317,7 @@ macro_rules! impl_unary_sink {
                 $t {
                     call: Some(call),
                     write_flags: 0,
-                    ser: ser,
+                    ser,
                 }
             }
 
@@ -349,8 +349,8 @@ macro_rules! impl_unary_sink {
 
                 $rt {
                     call: self.call.take().unwrap(),
-                    cq_f: cq_f,
-                    err: err,
+                    cq_f,
+                    err,
                 }
             }
         }
@@ -415,7 +415,7 @@ macro_rules! impl_stream_sink {
                     status: RpcStatus::ok(),
                     flushed: false,
                     closed: false,
-                    ser: ser,
+                    ser,
                 }
             }
 
@@ -439,8 +439,8 @@ macro_rules! impl_stream_sink {
 
                 $ft {
                     call: self.call.take().unwrap(),
-                    fail_f: fail_f,
-                    err: err,
+                    fail_f,
+                    err,
                 }
             }
         }

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

@@ -25,7 +25,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)
@@ -41,6 +50,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)]