add ability to close/stop scheduler and executor

python/examples/executor.py

@@ -27,3 +27,5 @@
 executor.wait_for_termination()
 # %%
 # %%
+executor.close()
+# %%

python/examples/scheduler.py

@@ -26,3 +26,4 @@
 # %%
 scheduler.wait_for_termination()
 # %%
+scheduler.close()

python/src/cluster.rs

@@ -96,6 +96,19 @@ impl PyScheduler {
             None => Ok(()),
         }
     }
+
+    pub fn close(&mut self) -> PyResult<()> {
+        let mut handle = None;
+        std::mem::swap(&mut self.handle, &mut handle);
+
+        match handle {
+            Some(handle) => handle.abort(),
+            None => {}
+        };
+
+        Ok(())
+    }
+
     #[classattr]
     pub fn version() -> &'static str {
         ballista_core::BALLISTA_VERSION
@@ -199,6 +212,18 @@ impl PyExecutor {
         }
     }
 
+    pub fn close(&mut self) -> PyResult<()> {
+        let mut handle = None;
+        std::mem::swap(&mut self.handle, &mut handle);
+
+        match handle {
+            Some(handle) => handle.abort(),
+            None => {}
+        };
+
+        Ok(())
+    }
+
     #[classattr]
     pub fn version() -> &'static str {
         ballista_core::BALLISTA_VERSION