feat(python): add graceful shutdown method

.changeset/kind-flowers-relate.md

@@ -0,0 +1,5 @@
+---
+"python-sdk": minor
+---
+
+Add shutdown() method for graceful client shutdown.

.changeset/loud-lions-confess.md

@@ -0,0 +1,5 @@
+---
+"eppo_core": minor
+---
+
+Add by-ref graceful shutdown for BackgroundThread.

eppo_core/src/background/thread.rs

@@ -1,10 +1,12 @@
+use std::sync::Mutex;
+
 use super::runtime::BackgroundRuntime;
 
 /// An owning handle to a background thread running tokio runtime.
 ///
 /// When the handle is dropped, the tokio runtime is commanded to exit and the thread shuts down.
 pub struct BackgroundThread {
-    join_handle: std::thread::JoinHandle<()>,
+    join_handle: Mutex<Option<std::thread::JoinHandle<()>>>,
     runtime: BackgroundRuntime<tokio::runtime::Handle>,
 }
 
@@ -23,11 +25,13 @@ impl BackgroundThread {
         let join_handle = std::thread::Builder::new()
             .name("eppo-background".to_owned())
             .spawn(move || {
+                log::info!(target: "eppo", "BackgroundThread: started");
                 runtime.block_on(wait);
+                log::info!(target: "eppo", "BackgroundThread: exiting");
             })?;
 
         Ok(BackgroundThread {
-            join_handle,
+            join_handle: Mutex::new(Some(join_handle)),
             runtime: background_runtime,
         })
     }
@@ -38,16 +42,32 @@ impl BackgroundThread {
 
     /// Command the associated background thread to exit (without waiting for it to complete).
     ///
-    /// Prefer `graceful_shutdown()` if you have the time to wait.
+    /// Prefer `shutdown()` if you have the time to wait.
     pub fn kill(&self) {
         self.runtime.stop();
     }
 
     /// Command background activities to stop and wait for thread to terminate.
-    pub fn graceful_shutdown(self) {
+    pub fn shutdown(&self) {
         self.runtime.stop();
 
-        let _ = self.join_handle.join();
+        let join_handle = {
+            // scope to keep mutex lock short
+            let Ok(mut join_handle) = self.join_handle.lock() else {
+                return;
+            };
+            join_handle.take()
+        };
+
+        if let Some(join_handle) = join_handle {
+            let _ = join_handle.join();
+        }
+    }
+
+    /// Command background activities to stop and wait for thread to terminate.
+    #[deprecated]
+    pub fn graceful_shutdown(self) {
+        self.shutdown();
     }
 }
 
@@ -69,6 +89,6 @@ mod tests {
 
         assert_eq!(received, true);
 
-        background_thread.graceful_shutdown();
+        background_thread.shutdown();
     }
 }

eppo_core/src/configuration_fetcher.rs

@@ -31,8 +31,10 @@ impl ConfigurationFetcher {
     pub fn new(config: ConfigurationFetcherConfig) -> ConfigurationFetcher {
         let builder = reqwest::Client::builder();
         let client = match builder.build() {
-            Err(e) => { panic!("Reqwest client build failed {:?}", e); }
-            Ok(client) => client
+            Err(e) => {
+                panic!("Reqwest client build failed {:?}", e);
+            }
+            Ok(client) => client,
         };
 
         ConfigurationFetcher {

eppo_core/src/event_ingestion/event.rs

@@ -15,4 +15,4 @@ pub(super) struct Event {
     pub event_type: String,
 
     pub payload: serde_json::Value,
-}
+}

python-sdk/python/eppo_client/_eppo_client.pyi

@@ -143,6 +143,7 @@ class EppoClient:
     def set_is_graceful_mode(self, is_graceful_mode: bool): ...
     def is_initialized(self) -> bool: ...
     def wait_for_initialization(self) -> None: ...
+    def shutdown(self) -> None: ...
 
 class ContextAttributes:
     def __new__(
@@ -154,7 +155,7 @@ class ContextAttributes:
     def empty() -> ContextAttributes: ...
     @staticmethod
     def from_dict(
-        attributes: Dict[str, Union[str, int, float, bool, None]]
+        attributes: Dict[str, Union[str, int, float, bool, None]],
     ) -> ContextAttributes: ...
     @property
     def numeric_attributes(self) -> Dict[str, float]: ...

python-sdk/src/client.rs

@@ -469,6 +469,23 @@ impl EppoClient {
         }
     }
 
+    /// Shutdown the client and wait for the background thread to exit.
+    ///
+    /// It is recommended to call this method before exiting the program so the background thread
+    /// has a chance to cleanup its resources and avoid getting killed, producing scary (but
+    /// harmless) panic backtrace.
+    pub fn shutdown(&self, py: Python) {
+        log::info!(target: "eppo", "shutting down client");
+        // Release python GIL before we try blocking and waiting for the background thread to
+        // join. This is required because pyo3_log in background thread might attempt acquiring GIL,
+        // which would result in deadlock as GIL is held by the current thread which is waiting for
+        // background thread to exit.
+        py.allow_threads(|| {
+            if let Some(thread) = &self.background_thread {
+                thread.shutdown();
+            }
+        })
+    }
     /// Returns a set of all flag keys that have been initialized.
     /// This can be useful to debug the initialization process.
     ///
@@ -699,18 +716,12 @@ impl EppoClient {
             .call_method1(py, intern!(py, "log_bandit_action"), (event,))?;
         Ok(())
     }
-
-    pub fn shutdown(&self) {
-        if let Some(thread) = &self.background_thread {
-            // Using `.kill()` instead of `.shutdown()` here because we don't need to wait for the
-            // poller thread to exit.
-            thread.kill();
-        }
-    }
 }
 
 impl Drop for EppoClient {
     fn drop(&mut self) {
-        self.shutdown();
+        if let Some(thread) = &self.background_thread {
+            thread.kill();
+        }
     }
 }

python-sdk/src/init.rs

@@ -35,7 +35,7 @@ pub fn init(config: Bound<ClientConfig>) -> PyResult<Py<EppoClient>> {
         std::mem::replace(&mut *instance, Some(client))
     };
     if let Some(existing) = existing {
-        existing.get().shutdown();
+        existing.get().shutdown(py);
         existing.drop_ref(py);
     }
 

python-sdk/tests/test_client.py

@@ -24,6 +24,15 @@ def test_wait_for_initialization():
     assert client.is_initialized() == True
 
 
+@pytest.mark.rust_only
+def test_shutdown():
+    client = init("ufc", wait_for_init=False)
+
+    client.shutdown()
+    # safe to call multiple times:
+    client.shutdown()
+
+
 def test_get_flag_keys_none():
     client = init("ufc", wait_for_init=False)
     assert client.get_flag_keys() == set()