From 0368724cf78f6f5f3fdacd9892ff33a081319b28 Mon Sep 17 00:00:00 2001
From: sinhaashish <ashi.sinha.87@gmail.com>
Date: Thu, 13 Jun 2024 12:08:59 +0000
Subject: [PATCH] feat(topology/poolaffinitykey):  add topology/poolaffinitykey
 code to csi-driver

Signed-off-by: sinhaashish <ashi.sinha.87@gmail.com>
---
 .../src/bin/controller/controller.rs          | 11 +++++
 control-plane/csi-driver/src/context.rs       | 40 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/control-plane/csi-driver/src/bin/controller/controller.rs b/control-plane/csi-driver/src/bin/controller/controller.rs
index b17e83e78..3587392f6 100644
--- a/control-plane/csi-driver/src/bin/controller/controller.rs
+++ b/control-plane/csi-driver/src/bin/controller/controller.rs
@@ -1075,6 +1075,7 @@ fn context_into_topology(context: &CreateParams) -> CreateVolumeTopology {
     let mut pool_inclusive_label_topology: HashMap<String, String> = HashMap::new();
     let mut node_inclusive_label_topology: HashMap<String, String> = HashMap::new();
     let mut node_exclusive_label_topology: HashMap<String, String> = HashMap::new();
+    let mut pool_affinity_label_topology: Vec<String> = Vec::<String>::new();
     pool_inclusive_label_topology.insert(dsp_created_by_key(), String::from(DSP_OPERATOR));
     pool_inclusive_label_topology.extend(
         context
@@ -1090,6 +1091,13 @@ fn context_into_topology(context: &CreateParams) -> CreateVolumeTopology {
             .clone()
             .unwrap_or_default(),
     );
+    pool_affinity_label_topology.extend(
+        context
+            .publish_params()
+            .pool_affinity_topology_key()
+            .clone()
+            .unwrap_or_default(),
+    );
     node_inclusive_label_topology.extend(
         context
             .publish_params()
@@ -1111,14 +1119,17 @@ fn context_into_topology(context: &CreateParams) -> CreateVolumeTopology {
             .clone()
             .unwrap_or_default(),
     );
+
     CreateVolumeTopology::new(
         Some(models::NodeTopology::labelled(LabelledTopology {
             exclusion: node_exclusive_label_topology,
             inclusion: node_inclusive_label_topology,
+            affinitykey: Default::default(),
         })),
         Some(PoolTopology::labelled(LabelledTopology {
             exclusion: Default::default(),
             inclusion: pool_inclusive_label_topology,
+            affinitykey: pool_affinity_label_topology,
         })),
     )
 }
diff --git a/control-plane/csi-driver/src/context.rs b/control-plane/csi-driver/src/context.rs
index c023bc04b..e17edc3ee 100644
--- a/control-plane/csi-driver/src/context.rs
+++ b/control-plane/csi-driver/src/context.rs
@@ -58,6 +58,8 @@ pub enum Parameters {
     QuiesceFs,
     #[strum(serialize = "poolAffinityTopologyLabel")]
     PoolAffinityTopologyLabel,
+    #[strum(serialize = "poolAffinityTopologyKey")]
+    PoolAffinityTopologyKey,
     #[strum(serialize = "poolHasTopologyKey")]
     PoolHasTopologyKey,
     #[strum(serialize = "nodeAffinityTopologyLabel")]
@@ -116,6 +118,27 @@ impl Parameters {
         })
     }
 
+    fn parse_topology_param_vec(
+        value: Option<&String>,
+    ) -> Result<Option<Vec<String>>, tonic::Status> {
+        Ok(match value {
+            Some(labels) => {
+                let mut result_vec = Vec::new();
+                for label in labels.split('\n') {
+                    if !label.is_empty() {
+                        result_vec.push(label.to_string())
+                    } else {
+                        return Err(tonic::Status::invalid_argument(format!(
+                            "Invalid label : {value:?}"
+                        )));
+                    }
+                }
+                Some(result_vec)
+            }
+            None => None,
+        })
+    }
+
     fn parse_u32(value: Option<&String>) -> Result<Option<u32>, ParseIntError> {
         Ok(match value {
             Some(value) => value.parse::<u32>().map(Some)?,
@@ -176,6 +199,12 @@ impl Parameters {
     ) -> Result<Option<HashMap<String, String>>, tonic::Status> {
         Self::parse_topology_param(value)
     }
+    /// Parse the value for `Self::PoolAffinityTopologyKey`.
+    pub fn pool_affinity_topology_key(
+        value: Option<&String>,
+    ) -> Result<Option<Vec<String>>, tonic::Status> {
+        Self::parse_topology_param_vec(value)
+    }
     /// Parse the value for `Self::PoolHasTopologyKey`.
     pub fn pool_has_topology_key(
         value: Option<&String>,
@@ -217,6 +246,7 @@ pub struct PublishParams {
     fs_type: Option<FileSystem>,
     fs_id: Option<Uuid>,
     pool_affinity_topology_label: Option<HashMap<String, String>>,
+    pool_affinity_topology_key: Option<Vec<String>>,
     pool_has_topology_key: Option<HashMap<String, String>>,
     node_affinity_topology_label: Option<HashMap<String, String>>,
     node_has_topology_key: Option<HashMap<String, String>>,
@@ -247,6 +277,10 @@ impl PublishParams {
     pub fn pool_affinity_topology_label(&self) -> &Option<HashMap<String, String>> {
         &self.pool_affinity_topology_label
     }
+    /// Get the `Parameters::PoolAffinityTopologyKey` value.
+    pub fn pool_affinity_topology_key(&self) -> &Option<Vec<String>> {
+        &self.pool_affinity_topology_key
+    }
     /// Get the `Parameters::PoolHasTopologyKey` value.
     pub fn pool_has_topology_key(&self) -> &Option<HashMap<String, String>> {
         &self.pool_has_topology_key
@@ -319,6 +353,11 @@ impl TryFrom<&HashMap<String, String>> for PublishParams {
         )
         .map_err(|_| tonic::Status::invalid_argument("Invalid pool_affinity_topology_label"))?;
 
+        let pool_affinity_topology_key = Parameters::pool_affinity_topology_key(
+            args.get(Parameters::PoolAffinityTopologyKey.as_ref()),
+        )
+        .map_err(|_| tonic::Status::invalid_argument("Invalid pool_affinity_topology_key"))?;
+
         let pool_has_topology_key =
             Parameters::pool_has_topology_key(args.get(Parameters::PoolHasTopologyKey.as_ref()))
                 .map_err(|_| tonic::Status::invalid_argument("Invalid pool_has_topology_key"))?;
@@ -356,6 +395,7 @@ impl TryFrom<&HashMap<String, String>> for PublishParams {
             fs_type,
             fs_id,
             pool_affinity_topology_label,
+            pool_affinity_topology_key,
             pool_has_topology_key,
             node_affinity_topology_label,
             node_has_topology_key,