add teardown to Configurator trait

Author: KSDaemon

rust/cubestore/cubestore/src/cluster/mod.rs

@@ -251,6 +251,13 @@ impl Configurator for WorkerConfigurator {
         };
         Ok(config)
     }
+
+    fn teardown() {
+        let teardown = SELECT_WORKER_SHUTDOWN.read().unwrap();
+        if teardown.is_some() {
+            teardown.as_ref().unwrap()();
+        }
+    }
 }
 
 #[cfg(not(target_os = "windows"))]
@@ -369,12 +376,23 @@ lazy_static! {
         std::sync::RwLock::new(None);
 }
 
+lazy_static! {
+    static ref SELECT_WORKER_SHUTDOWN: std::sync::RwLock<Option<Box<dyn Fn() + Send + Sync>>> =
+        std::sync::RwLock::new(None);
+}
+
 pub fn register_select_worker_setup(f: fn(&Runtime)) {
     let mut setup = SELECT_WORKER_SETUP.write().unwrap();
     assert!(setup.is_none(), "select worker setup already registered");
     *setup = Some(Box::new(f));
 }
 
+pub fn register_select_worker_teardown(f: fn()) {
+    let mut setup = SELECT_WORKER_SHUTDOWN.write().unwrap();
+    assert!(setup.is_none(), "select worker teardown already registered");
+    *setup = Some(Box::new(f));
+}
+
 pub fn register_select_worker_configure_fn(f: fn() -> BoxFuture<'static, Config>) {
     let mut func = SELECT_WORKER_CONFIGURE_FN.write().unwrap();
     assert!(

rust/cubestore/cubestore/src/cluster/worker_pool.rs

@@ -376,6 +376,14 @@ pub struct WorkerProcessArgs<C: Configurator, P: WorkerProcessing, S: ServicesTr
     timeout: Duration,
 }
 
+struct TeardownGuard<C: Configurator>(PhantomData<C>);
+
+impl<C: Configurator> Drop for TeardownGuard<C> {
+    fn drop(&mut self) {
+        C::teardown();
+    }
+}
+
 pub fn worker_main<C, P, S>(a: WorkerProcessArgs<C, P, S>) -> i32
 where
     C: Configurator,
@@ -399,6 +407,7 @@ where
     tokio_builder.thread_stack_size(stack_size);
     let runtime = tokio_builder.build().unwrap();
     C::setup(&runtime);
+    let _teardown_guard = TeardownGuard::<C>(PhantomData);
     runtime.block_on(async move {
         let services_client = S::connect(services_sender, services_reciever, timeout);
         let config = match C::configure(services_client).await {
@@ -503,6 +512,8 @@ mod tests {
             config.configure_injector().await;
             Ok(config)
         }
+
+        fn teardown() {}
     }
 
     pub struct Processor;
@@ -711,6 +722,8 @@ mod tests {
             let config = TestConfig { services_client };
             Ok(config)
         }
+
+        fn teardown() {}
     }
 
     pub struct ServProcessor;

rust/cubestore/cubestore/src/cluster/worker_services.rs

@@ -34,6 +34,8 @@ pub trait Configurator: Send + Sync + 'static {
             dyn Callable<Request = Self::ServicesRequest, Response = Self::ServicesResponse>,
         >,
     ) -> Result<Self::Config, CubeError>;
+
+    fn teardown();
 }
 
 #[async_trait]