From 6ee41d52e70afd330604f94e8de0fd26b2682250 Mon Sep 17 00:00:00 2001
From: andrussal <salumets.andrus@gmail.com>
Date: Tue, 4 Aug 2026 04:26:56 +0200
Subject: [PATCH] fix(tf): clean up local cluster resources

---
 .../deployers/local/src/cluster.rs            |  2 ++
 .../deployers/local/src/env/mod.rs            |  6 +++++
 .../deployers/local/src/node_control/mod.rs   |  4 ++++
 .../deployers/local/src/provisioner.rs        | 23 ++++++++++++++++++-
 4 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/testing-framework/deployers/local/src/cluster.rs b/testing-framework/deployers/local/src/cluster.rs
index adaa031..ab7b24c 100644
--- a/testing-framework/deployers/local/src/cluster.rs
+++ b/testing-framework/deployers/local/src/cluster.rs
@@ -29,6 +29,7 @@ impl<E: LocalDeployerEnv> LocalClusterOwner<E> {
     fn close(&self) {
         if !self.closed.swap(true, Ordering::AcqRel) {
             self.nodes.stop_all();
+            E::cleanup_local_cluster(self.nodes.deployment());
         }
     }
 
@@ -68,6 +69,7 @@ impl<E: LocalDeployerEnv> Clone for LocalCluster<E> {
 
 impl<E: LocalDeployerEnv> LocalCluster<E> {
     pub(crate) fn empty(deployment: E::Deployment, keep_tempdir: bool) -> Self {
+        E::prepare_local_cluster(&deployment);
         let nodes = NodeManager::new_with_seed(
             deployment.clone(),
             NodeClients::default(),
diff --git a/testing-framework/deployers/local/src/env/mod.rs b/testing-framework/deployers/local/src/env/mod.rs
index 6cbdcb0..3df5f50 100644
--- a/testing-framework/deployers/local/src/env/mod.rs
+++ b/testing-framework/deployers/local/src/env/mod.rs
@@ -75,6 +75,12 @@ pub trait LocalDeployerEnv: Application + Sized
 where
     <Self as Application>::NodeConfig: Clone + Send + Sync + 'static,
 {
+    /// Prepares application-specific resources associated with a local cluster.
+    fn prepare_local_cluster(_deployment: &Self::Deployment) {}
+
+    /// Releases application-specific resources associated with a local cluster.
+    fn cleanup_local_cluster(_deployment: &Self::Deployment) {}
+
     /// Returns named ports that should be reserved in addition to the main
     /// network port for each node.
     fn local_port_names() -> &'static [&'static str] {
diff --git a/testing-framework/deployers/local/src/node_control/mod.rs b/testing-framework/deployers/local/src/node_control/mod.rs
index 0a6f1af..9e4c2bc 100644
--- a/testing-framework/deployers/local/src/node_control/mod.rs
+++ b/testing-framework/deployers/local/src/node_control/mod.rs
@@ -85,6 +85,10 @@ pub struct NodeManagerSeed {
 }
 
 impl<E: LocalDeployerEnv> NodeManager<E> {
+    pub(crate) const fn deployment(&self) -> &E::Deployment {
+        &self.descriptors
+    }
+
     pub async fn spawn_initial_nodes(
         descriptors: &E::Deployment,
         keep_tempdir: bool,
diff --git a/testing-framework/deployers/local/src/provisioner.rs b/testing-framework/deployers/local/src/provisioner.rs
index 9e97d05..3cd6c91 100644
--- a/testing-framework/deployers/local/src/provisioner.rs
+++ b/testing-framework/deployers/local/src/provisioner.rs
@@ -268,7 +268,11 @@ async fn run_retry_attempt<E: LocalDeployerEnv>(
 
 #[cfg(test)]
 mod tests {
-    use std::{collections::HashMap, path::Path};
+    use std::{
+        collections::HashMap,
+        path::Path,
+        sync::atomic::{AtomicUsize, Ordering},
+    };
 
     use testing_framework_core::{
         scenario::{
@@ -298,6 +302,9 @@ mod tests {
 
     struct TestEnv;
 
+    static PREPARE_CALLS: AtomicUsize = AtomicUsize::new(0);
+    static CLEANUP_CALLS: AtomicUsize = AtomicUsize::new(0);
+
     #[async_trait::async_trait]
     impl Application for TestEnv {
         type Deployment = EmptyDeployment;
@@ -311,6 +318,14 @@ mod tests {
 
     #[async_trait::async_trait]
     impl LocalDeployerEnv for TestEnv {
+        fn prepare_local_cluster(_deployment: &Self::Deployment) {
+            PREPARE_CALLS.fetch_add(1, Ordering::Relaxed);
+        }
+
+        fn cleanup_local_cluster(_deployment: &Self::Deployment) {
+            CLEANUP_CALLS.fetch_add(1, Ordering::Relaxed);
+        }
+
         fn build_node_config(
             _topology: &Self::Deployment,
             _index: usize,
@@ -346,6 +361,8 @@ mod tests {
 
     #[tokio::test]
     async fn on_demand_managed_unit_exposes_shared_control_and_cleanup() {
+        PREPARE_CALLS.store(0, Ordering::Relaxed);
+        CLEANUP_CALLS.store(0, Ordering::Relaxed);
         let request = ClusterRequest::managed(EmptyDeployment)
             .with_start_mode(ClusterStartMode::OnDemand)
             .with_control(ClusterControlRequest::Full);
@@ -356,6 +373,7 @@ mod tests {
         let (cluster, mut unit) = provisioned.into_parts();
         let cluster = cluster.expect("managed unit should expose its concrete local handle");
 
+        assert_eq!(PREPARE_CALLS.load(Ordering::Relaxed), 1);
         assert_eq!(
             unit.control_profile(),
             ClusterControlProfile::ManualControlled
@@ -366,7 +384,10 @@ mod tests {
         unit.take_cleanup()
             .expect("managed unit should own cleanup")
             .cleanup();
+        assert_eq!(CLEANUP_CALLS.load(Ordering::Relaxed), 1);
         assert!(cluster.stop_all().is_err(), "cleanup must lock out clones");
+        drop(cluster);
+        assert_eq!(CLEANUP_CALLS.load(Ordering::Relaxed), 1);
     }
 
     #[tokio::test]