Removed properties should be known to the property manager.

cas/scheduler/BUILD

@@ -156,6 +156,7 @@ rust_library(
         ":scheduler",
         "//config",
         "//util:error",
+        "@crate_index//:parking_lot",
         "@crate_index//:tokio",
     ],
 )

cas/scheduler/property_modifier_scheduler.rs

@@ -12,35 +12,68 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use parking_lot::Mutex;
 use tokio::sync::watch;
 
 use action_messages::{ActionInfo, ActionInfoHashKey, ActionState};
-use config::schedulers::PropertyModification;
+use config::schedulers::{PropertyModification, PropertyType};
 use error::{Error, ResultExt};
 use platform_property_manager::PlatformPropertyManager;
 use scheduler::ActionScheduler;
 
 pub struct PropertyModifierScheduler {
     modifications: Vec<config::schedulers::PropertyModification>,
     scheduler: Arc<dyn ActionScheduler>,
+    property_managers: Mutex<HashMap<String, Arc<PlatformPropertyManager>>>,
 }
 
 impl PropertyModifierScheduler {
     pub fn new(config: &config::schedulers::PropertyModifierScheduler, scheduler: Arc<dyn ActionScheduler>) -> Self {
         Self {
             modifications: config.modifications.clone(),
             scheduler,
+            property_managers: Mutex::new(HashMap::new()),
         }
     }
 }
 
 #[async_trait]
 impl ActionScheduler for PropertyModifierScheduler {
     async fn get_platform_property_manager(&self, instance_name: &str) -> Result<Arc<PlatformPropertyManager>, Error> {
-        self.scheduler.get_platform_property_manager(instance_name).await
+        {
+            let property_managers = self.property_managers.lock();
+            if let Some(property_manager) = property_managers.get(instance_name) {
+                return Ok(property_manager.clone());
+            }
+        }
+        let property_manager = self.scheduler.get_platform_property_manager(instance_name).await?;
+        let mut known_properties = property_manager.get_known_properties().clone();
+        for modification in &self.modifications {
+            match modification {
+                PropertyModification::Remove(name) => {
+                    known_properties.entry(name.into()).or_insert(PropertyType::Priority);
+                }
+                PropertyModification::Add(_) => (),
+            }
+        }
+        let property_manager = {
+            let mut property_managers = self.property_managers.lock();
+            match property_managers.entry(instance_name.into()) {
+                Entry::Vacant(new_entry) => {
+                    let property_manager = Arc::new(PlatformPropertyManager::new(known_properties));
+                    new_entry.insert(property_manager.clone());
+                    property_manager
+                }
+                // We lost the race, use the other manager.
+                Entry::Occupied(old_entry) => old_entry.get().clone(),
+            }
+        };
+        Ok(property_manager)
     }
 
     async fn add_action(&self, mut action_info: ActionInfo) -> Result<watch::Receiver<Arc<ActionState>>, Error> {

cas/scheduler/tests/property_modifier_scheduler_test.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::HashMap;
+use std::iter::FromIterator;
 use std::sync::Arc;
 use std::time::UNIX_EPOCH;
 
@@ -53,22 +54,6 @@ mod property_modifier_scheduler_tests {
     use super::*;
     use pretty_assertions::assert_eq; // Must be declared in every module.
 
-    #[tokio::test]
-    async fn platform_property_manager_call_passed() -> Result<(), Error> {
-        let context = make_modifier_scheduler(vec![]);
-        let platform_property_manager = Arc::new(PlatformPropertyManager::new(HashMap::new()));
-        let instance_name = INSTANCE_NAME.to_string();
-        let (actual_manager, actual_instance_name) = join!(
-            context.modifier_scheduler.get_platform_property_manager(&instance_name),
-            context
-                .mock_scheduler
-                .expect_get_platform_property_manager(Ok(platform_property_manager.clone())),
-        );
-        assert_eq!(Arc::as_ptr(&platform_property_manager), Arc::as_ptr(&actual_manager?));
-        assert_eq!(instance_name, actual_instance_name);
-        Ok(())
-    }
-
     #[tokio::test]
     async fn add_action_adds_property() -> Result<(), Error> {
         let name = "name".to_string();
@@ -243,4 +228,42 @@ mod property_modifier_scheduler_tests {
         assert_eq!(action_name, actual_action_name);
         Ok(())
     }
+
+    #[tokio::test]
+    async fn remove_adds_to_underlying_manager() -> Result<(), Error> {
+        let name = "name".to_string();
+        let context = make_modifier_scheduler(vec![PropertyModification::Remove(name.clone())]);
+        let scheduler_property_manager = Arc::new(PlatformPropertyManager::new(HashMap::new()));
+        let get_property_manager_fut = context
+            .mock_scheduler
+            .expect_get_platform_property_manager(Ok(scheduler_property_manager));
+        let property_manager_fut = context.modifier_scheduler.get_platform_property_manager(INSTANCE_NAME);
+        let (actual_instance_name, property_manager) = join!(get_property_manager_fut, property_manager_fut);
+        assert_eq!(
+            HashMap::<_, _>::from_iter([(name, PropertyType::Priority)]),
+            *property_manager?.get_known_properties()
+        );
+        assert_eq!(actual_instance_name, INSTANCE_NAME);
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn remove_retains_type_in_underlying_manager() -> Result<(), Error> {
+        let name = "name".to_string();
+        let context = make_modifier_scheduler(vec![PropertyModification::Remove(name.clone())]);
+        let scheduler_property_manager = Arc::new(PlatformPropertyManager::new(HashMap::<_, _>::from_iter([(
+            name.clone(),
+            PropertyType::Exact,
+        )])));
+        let get_property_manager_fut = context
+            .mock_scheduler
+            .expect_get_platform_property_manager(Ok(scheduler_property_manager));
+        let property_manager_fut = context.modifier_scheduler.get_platform_property_manager(INSTANCE_NAME);
+        let (_, property_manager) = join!(get_property_manager_fut, property_manager_fut);
+        assert_eq!(
+            HashMap::<_, _>::from_iter([(name, PropertyType::Exact)]),
+            *property_manager?.get_known_properties()
+        );
+        Ok(())
+    }
 }

gencargo/property_modifier_scheduler/Cargo.toml

@@ -20,6 +20,7 @@ doctest = false
 
 [dependencies]
 async-trait = { workspace = true }
+parking_lot = { workspace = true }
 tokio = { workspace = true }
 
 # Local libraries.