feat(meta): add affinity hooks (#9596)

* feat(meta): add affinity hooks

Signed-off-by: WenyXu <wenymedia@gmail.com>

* test(meta): cover load-based affinity fallback queries

Signed-off-by: WenyXu <wenymedia@gmail.com>

* refactor(meta): borrow affinity requirements during selection

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(meta): remove unused affinity error variant

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>

src/cmd/src/standalone.rs

@@ -681,6 +681,8 @@ impl StartCommand {
         ));
 
         let ddl_context = DdlContext {
+            affinity_resolver: None,
+            peer_alloc_context_provider: None,
             node_manager: node_manager.clone(),
             cache_invalidator: layered_cache_registry.clone(),
             memory_region_keeper: Arc::new(MemoryRegionKeeper::default()),

src/common/meta/AGENTS.md

@@ -11,6 +11,7 @@ and durable DDL procedures. Metasrv process and service wiring live in
 
 | Area | Path | Entry point |
 | --- | --- | --- |
+| Affinity | `src/common/meta/src/affinity.rs` | Requirements, reference resolution, node matching, and allocation-context provider |
 | Metadata model | `src/common/meta/src/key.rs`, `src/common/meta/src/key/` | Typed keys, values, and `TableMetadataManager` |
 | KV storage | `src/common/meta/src/kv_backend.rs`, `src/common/meta/src/kv_backend/` | `KvBackend`, transactions, memory/etcd/RDS backends |
 | DDL procedures | `src/common/meta/src/ddl.rs`, `src/common/meta/src/ddl/`, `src/common/meta/src/ddl_manager.rs` | Durable DDL state machines and task dispatch |
@@ -31,6 +32,10 @@ and durable DDL procedures. Metasrv process and service wiring live in
   `src/table/src/requests.rs` — classification, validation, logical-table
   eligibility, and the metadata-only alter flow come with it. Do not add
   per-option special cases to the alter conversion or DDL routing.
+- Create-table preparation obtains allocation inputs from the optional query
+  provider in `DdlContext`. Install providers before registering procedure
+  loaders so restored procedures use the same runtime dependencies. Affinity
+  requirements and fallback policy are runtime inputs, not persisted fields.
 - `KvBackend` behavior changes should extend the shared tests in
   `src/common/meta/src/kv_backend/test.rs`.
 

src/common/meta/src/affinity.rs

@@ -0,0 +1,254 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Resolve reference attributes into requirements, then evaluate candidates.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use crate::DatanodeId;
+use crate::cluster::NodeInfo;
+use crate::error::Result;
+use crate::peer::PeerAllocContext;
+use crate::rpc::ddl::QueryContext;
+
+/// Named attributes used to express affinity requirements.
+pub type AffinityAttributes = HashMap<String, String>;
+
+/// Identifies the node whose attributes define affinity requirements.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum AffinityReference {
+    /// Uses the datanode's last-known attributes, even if its lease has expired.
+    Datanode(DatanodeId),
+    /// Uses the frontend identified by its advertised address.
+    Frontend(String),
+}
+
+/// Behavior when affinity cannot be satisfied.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub enum AffinityFallback {
+    /// Returns an error instead of selecting without affinity.
+    #[default]
+    Error,
+    /// Allows selection among eligible nodes without affinity.
+    Unconstrained,
+}
+
+/// Builds allocation inputs from the query that initiated an operation.
+#[async_trait::async_trait]
+pub trait PeerAllocContextProvider: Send + Sync {
+    /// Resolves query-specific allocation inputs using metadata, without selecting peers.
+    async fn build(&self, query_context: &QueryContext) -> Result<PeerAllocContext>;
+}
+
+/// Shared allocation-context provider.
+pub type PeerAllocContextProviderRef = Arc<dyn PeerAllocContextProvider>;
+
+/// Nonempty attributes a candidate must match.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AffinityRequirements {
+    attributes: AffinityAttributes,
+}
+
+impl AffinityRequirements {
+    /// Returns no requirements when the supplied attributes are empty.
+    pub fn new(attributes: AffinityAttributes) -> Option<Self> {
+        (!attributes.is_empty()).then_some(Self { attributes })
+    }
+
+    /// Returns the attributes that candidates are evaluated against.
+    pub fn attributes(&self) -> &AffinityAttributes {
+        &self.attributes
+    }
+}
+
+/// Resolves references using current metadata, without selecting candidates.
+#[async_trait::async_trait]
+pub trait AffinityResolver: Send + Sync {
+    /// Returns one result per reference, in order.
+    /// None means requirements are unavailable; Err means resolution failed.
+    async fn resolve(
+        &self,
+        references: &[AffinityReference],
+    ) -> Result<Vec<Option<AffinityRequirements>>>;
+}
+
+/// Shared requirement resolver.
+pub type AffinityResolverRef = Arc<dyn AffinityResolver>;
+
+/// Evaluates candidate nodes without I/O or scheduling decisions.
+pub trait Affinity: Send + Sync {
+    /// Evaluates a node snapshot; implementations choose which fields to inspect.
+    fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool;
+}
+
+/// Accepts every candidate without imposing affinity constraints.
+pub struct NoopAffinity;
+
+impl Affinity for NoopAffinity {
+    fn matches(&self, _: &AffinityRequirements, _: &NodeInfo) -> bool {
+        true
+    }
+}
+
+/// Shared affinity evaluator.
+pub type AffinityRef = Arc<dyn Affinity>;
+
+#[cfg(test)]
+mod tests {
+    use std::sync::RwLock;
+
+    use super::*;
+    use crate::cluster::{DatanodeStatus, NodeStatus};
+    use crate::peer::Peer;
+
+    struct NodeAffinityResolver {
+        attributes: RwLock<HashMap<AffinityReference, AffinityAttributes>>,
+    }
+
+    #[async_trait::async_trait]
+    impl AffinityResolver for NodeAffinityResolver {
+        async fn resolve(
+            &self,
+            references: &[AffinityReference],
+        ) -> Result<Vec<Option<AffinityRequirements>>> {
+            let snapshot = self.attributes.read().unwrap();
+            Ok(references
+                .iter()
+                .map(|reference| {
+                    snapshot
+                        .get(reference)
+                        .cloned()
+                        .and_then(AffinityRequirements::new)
+                })
+                .collect())
+        }
+    }
+
+    struct ExactAffinity;
+
+    impl Affinity for ExactAffinity {
+        fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool {
+            requirements
+                .attributes()
+                .iter()
+                .all(|(key, expected)| candidate.env_vars.get(key) == Some(expected))
+        }
+    }
+
+    fn zone(name: &str) -> AffinityAttributes {
+        HashMap::from([("AZ".to_string(), name.to_string())])
+    }
+
+    fn node(id: DatanodeId, env_vars: AffinityAttributes) -> NodeInfo {
+        NodeInfo {
+            peer: Peer::empty(id),
+            last_activity_ts: 0,
+            status: NodeStatus::Datanode(DatanodeStatus {
+                rcus: 0,
+                wcus: 0,
+                leader_regions: 0,
+                follower_regions: 0,
+                workloads: Default::default(),
+            }),
+            version: String::new(),
+            git_commit: String::new(),
+            start_time_ms: 0,
+            total_cpu_millicores: 0,
+            total_memory_bytes: 0,
+            cpu_usage_millicores: 0,
+            memory_usage_bytes: 0,
+            hostname: String::new(),
+            env_vars,
+        }
+    }
+
+    #[tokio::test]
+    async fn test_resolve_reference_and_match_candidates() {
+        let resolver = NodeAffinityResolver {
+            attributes: RwLock::new(HashMap::from([(AffinityReference::Datanode(1), zone("a"))])),
+        };
+        let references = [AffinityReference::Datanode(1)];
+        let resolved = resolver.resolve(&references).await.unwrap();
+        assert_eq!(resolved, vec![AffinityRequirements::new(zone("a"))]);
+        let requirements = resolved[0].as_ref().unwrap();
+        let candidates = [
+            node(2, zone("a")),
+            node(3, zone("b")),
+            node(4, HashMap::new()),
+        ];
+        assert_eq!(
+            candidates
+                .iter()
+                .map(|node| (node.peer.id, ExactAffinity.matches(requirements, node)))
+                .collect::<Vec<_>>(),
+            vec![(2, true), (3, false), (4, false),]
+        );
+
+        assert!(
+            candidates
+                .iter()
+                .all(|node| NoopAffinity.matches(requirements, node))
+        );
+
+        // A new allocation attempt resolves current attributes rather than a saved requirement.
+        resolver
+            .attributes
+            .write()
+            .unwrap()
+            .insert(AffinityReference::Datanode(1), zone("b"));
+        let resolved = resolver.resolve(&references).await.unwrap();
+        let requirements = resolved[0].as_ref().unwrap();
+        let matched: Vec<_> = candidates
+            .iter()
+            .filter(|node| ExactAffinity.matches(requirements, node))
+            .map(|node| node.peer.id)
+            .collect();
+        assert_eq!(matched, vec![3]);
+    }
+
+    #[tokio::test]
+    async fn test_resolution_preserves_order_and_unavailable_requirements() {
+        let resolver = NodeAffinityResolver {
+            attributes: RwLock::new(HashMap::from([
+                (AffinityReference::Datanode(1), zone("a")),
+                (AffinityReference::Datanode(2), HashMap::new()),
+                (
+                    AffinityReference::Frontend("frontend:4001".into()),
+                    zone("b"),
+                ),
+            ])),
+        };
+        let references = [
+            AffinityReference::Frontend("missing:4001".into()),
+            AffinityReference::Datanode(1),
+            AffinityReference::Frontend("frontend:4001".into()),
+            AffinityReference::Datanode(2),
+            AffinityReference::Datanode(3),
+            AffinityReference::Frontend("frontend:4001".into()),
+        ];
+        assert_eq!(
+            resolver.resolve(&references).await.unwrap(),
+            vec![
+                None,
+                AffinityRequirements::new(zone("a")),
+                AffinityRequirements::new(zone("b")),
+                None,
+                None,
+                AffinityRequirements::new(zone("b")),
+            ]
+        );
+        assert!(AffinityRequirements::new(HashMap::new()).is_none());
+    }
+}

src/common/meta/src/cluster.rs

@@ -106,7 +106,7 @@ fn calculate_node_id(addr: &str) -> u64 {
 }
 
 /// The information of a node in the cluster.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct NodeInfo {
     /// The peer information. [node_id, address]
     pub peer: Peer,
@@ -148,7 +148,7 @@ pub enum Role {
     Metasrv,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum NodeStatus {
     Datanode(DatanodeStatus),
     Frontend(FrontendStatus),
@@ -171,7 +171,7 @@ impl NodeStatus {
 }
 
 /// The status of a datanode.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct DatanodeStatus {
     /// The read capacity units during this period.
     pub rcus: i64,
@@ -186,7 +186,7 @@ pub struct DatanodeStatus {
 }
 
 /// The status of a frontend.
-#[derive(Debug, Default, Serialize, Deserialize)]
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct FrontendStatus {
     /// The workloads of the frontend.
     #[serde(default)]
@@ -194,7 +194,7 @@ pub struct FrontendStatus {
 }
 
 /// The status of a flownode.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct FlownodeStatus {
     /// The workloads of the flownode.
     #[serde(default)]
@@ -202,7 +202,7 @@ pub struct FlownodeStatus {
 }
 
 /// The status of a metasrv.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct MetasrvStatus {
     pub is_leader: bool,
 }

src/common/meta/src/ddl.rs

@@ -18,6 +18,7 @@ use std::time::Duration;
 use store_api::storage::{RegionId, TableId};
 
 use crate::DatanodeId;
+use crate::affinity::{AffinityResolverRef, PeerAllocContextProviderRef};
 use crate::cache_invalidator::CacheInvalidatorRef;
 use crate::ddl::flow_meta::FlowMetadataAllocatorRef;
 use crate::ddl::table_meta::TableMetadataAllocatorRef;
@@ -105,6 +106,10 @@ impl RegionFailureDetectorController for NoopRegionFailureDetectorControl {
 /// The context of ddl.
 #[derive(Clone)]
 pub struct DdlContext {
+    /// Optional resolver for automatic allocation affinity.
+    pub affinity_resolver: Option<AffinityResolverRef>,
+    /// Builds allocation inputs from the initiating query.
+    pub peer_alloc_context_provider: Option<PeerAllocContextProviderRef>,
     /// Sends querying and requests to nodes.
     pub node_manager: NodeManagerRef,
     /// Cache invalidation.

src/common/meta/src/ddl/create_table.rs

@@ -165,7 +165,10 @@ impl CreateTableProcedure {
             return Ok(Status::done_with_output(table_id));
         }
 
-        self.data.state = CreateTableState::DatanodeCreateRegions;
+        let alloc_context = match &self.context.peer_alloc_context_provider {
+            Some(provider) => provider.build(&self.data.query_context).await?,
+            None => PeerAllocContext::default(),
+        };
         let TableMetadata {
             table_id,
             table_route,
@@ -173,14 +176,10 @@ impl CreateTableProcedure {
         } = self
             .context
             .table_metadata_allocator
-            .create_with_context(
-                &self.data.task,
-                &PeerAllocContext {
-                    extensions: self.data.query_context.extensions.clone(),
-                },
-            )
+            .create_with_context(&self.data.task, &alloc_context)
             .await?;
         self.set_allocated_metadata(table_id, table_route, region_wal_options);
+        self.data.state = CreateTableState::DatanodeCreateRegions;
 
         Ok(Status::executing(true))
     }

src/common/meta/src/ddl/tests/create_table.rs

@@ -398,3 +398,143 @@ async fn test_memory_region_keeper_guard_dropped_on_procedure_done() {
             .contains(datanode_id, region_id)
     );
 }
+
+#[tokio::test]
+async fn test_on_prepare_uses_peer_alloc_context_provider() {
+    use std::sync::Mutex;
+
+    use crate::affinity::{AffinityRequirements, PeerAllocContextProvider};
+    use crate::ddl::table_meta::TableMetadataAllocator;
+    use crate::error;
+    use crate::peer::{PeerAllocContext, PeerAllocator};
+    use crate::rpc::ddl::QueryContext;
+
+    #[derive(Clone, Copy, PartialEq, Eq)]
+    enum ProviderMode {
+        Absent,
+        Available,
+        Error,
+    }
+
+    struct Provider {
+        mode: ProviderMode,
+        calls: Mutex<Vec<QueryContext>>,
+    }
+
+    #[async_trait::async_trait]
+    impl PeerAllocContextProvider for Provider {
+        async fn build(&self, query_context: &QueryContext) -> Result<PeerAllocContext> {
+            self.calls.lock().unwrap().push(query_context.clone());
+            if self.mode == ProviderMode::Error {
+                return error::UnexpectedSnafu {
+                    err_msg: "allocation context lookup failed",
+                }
+                .fail();
+            }
+            Ok(PeerAllocContext {
+                affinity: AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])),
+            })
+        }
+    }
+
+    #[derive(Default)]
+    struct Allocator(Mutex<Vec<PeerAllocContext>>);
+
+    #[async_trait::async_trait]
+    impl PeerAllocator for Allocator {
+        async fn alloc(&self, num: usize, ctx: &PeerAllocContext) -> Result<Vec<Peer>> {
+            self.0.lock().unwrap().push(ctx.clone());
+            Ok(vec![Peer::empty(2); num])
+        }
+    }
+
+    for (mode, table_exists) in [
+        (ProviderMode::Available, false),
+        (ProviderMode::Absent, false),
+        (ProviderMode::Error, false),
+        (ProviderMode::Error, true),
+    ] {
+        let provider = Arc::new(Provider {
+            mode,
+            calls: Mutex::new(Vec::new()),
+        });
+        let allocator = Arc::new(Allocator::default());
+        let mut context = new_ddl_context(Arc::new(MockDatanodeManager::new(())));
+        if mode != ProviderMode::Absent {
+            context.peer_alloc_context_provider = Some(provider.clone());
+        }
+        context.table_metadata_allocator = Arc::new(TableMetadataAllocator::with_peer_allocator(
+            context.table_metadata_allocator.table_id_allocator(),
+            context.table_metadata_allocator.wal_options_allocator(),
+            allocator.clone(),
+        ));
+        let mut task = test_create_table_task("provider_table");
+        task.create_table.create_if_not_exists = true;
+        if table_exists {
+            context
+                .table_metadata_manager
+                .create_table_metadata(
+                    task.table_info.clone(),
+                    TableRouteValue::physical(vec![]),
+                    HashMap::new(),
+                )
+                .await
+                .unwrap();
+        }
+        let query_context = QueryContext {
+            current_catalog: "test_catalog".into(),
+            current_schema: "test_schema".into(),
+            timezone: "UTC".into(),
+            extensions: HashMap::from([("custom-key".into(), "custom-value".into())]),
+            channel: 1,
+            snapshot_seqs: HashMap::from([(1, 10)]),
+            sst_min_sequences: HashMap::from([(1, 2)]),
+        };
+        let mut procedure =
+            CreateTableProcedure::new_with_query_context(task, query_context.clone(), context)
+                .unwrap();
+        let result = procedure.on_prepare().await;
+        if table_exists {
+            assert_matches!(result.unwrap(), Status::Done { .. });
+            assert!(provider.calls.lock().unwrap().is_empty());
+            assert!(allocator.0.lock().unwrap().is_empty());
+            continue;
+        }
+        let expected_calls = if mode == ProviderMode::Absent {
+            vec![]
+        } else {
+            vec![query_context]
+        };
+        assert_eq!(*provider.calls.lock().unwrap(), expected_calls);
+        if mode == ProviderMode::Error {
+            let error = result.unwrap_err();
+            assert_matches!(error, Error::Unexpected { .. });
+            assert!(
+                error
+                    .to_string()
+                    .contains("allocation context lookup failed")
+            );
+            assert_matches!(procedure.data.state, CreateTableState::Prepare);
+            assert!(allocator.0.lock().unwrap().is_empty());
+        } else {
+            result.unwrap();
+            assert_matches!(
+                procedure.data.state,
+                CreateTableState::DatanodeCreateRegions
+            );
+            let calls = allocator.0.lock().unwrap();
+            assert_eq!(calls.len(), 1);
+            let expected = if mode == ProviderMode::Available {
+                PeerAllocContext {
+                    affinity: AffinityRequirements::new(HashMap::from([(
+                        "pool".into(),
+                        "a".into(),
+                    )])),
+                }
+            } else {
+                PeerAllocContext::default()
+            };
+            assert_eq!(calls[0].affinity, expected.affinity);
+        }
+    }
+}

src/common/meta/src/ddl_manager.rs

@@ -33,6 +33,7 @@ use snafu::{OptionExt, ResultExt, ensure};
 use store_api::storage::{RegionId, TableId};
 use table::table_name::TableName;
 
+use crate::affinity::PeerAllocContextProviderRef;
 use crate::ddl::alter_database::AlterDatabaseProcedure;
 use crate::ddl::alter_logical_tables::AlterLogicalTablesProcedure;
 use crate::ddl::alter_table::{AlterTableProcedure, RegionRouteChanged, only_sets_skip_wal};
@@ -300,6 +301,15 @@ impl DdlManager {
         self
     }
 
+    /// Sets the provider used to build allocation inputs from queries.
+    pub fn with_peer_alloc_context_provider(
+        mut self,
+        provider: PeerAllocContextProviderRef,
+    ) -> Self {
+        self.ddl_context.peer_alloc_context_provider = Some(provider);
+        self
+    }
+
     pub fn with_create_database_metadata_committer(
         mut self,
         committer: CreateDatabaseMetadataCommitterRef,
@@ -1875,6 +1885,8 @@ mod tests {
 
         let ddl_manager = DdlManager::new(
             DdlContext {
+                affinity_resolver: None,
+                peer_alloc_context_provider: None,
                 node_manager: Arc::new(DummyDatanodeManager),
                 cache_invalidator: Arc::new(DummyCacheInvalidator),
                 table_metadata_manager,
@@ -1941,6 +1953,8 @@ mod tests {
 
         DdlManager::new(
             DdlContext {
+                affinity_resolver: None,
+                peer_alloc_context_provider: None,
                 node_manager: Arc::new(DummyDatanodeManager),
                 cache_invalidator: Arc::new(DummyCacheInvalidator),
                 table_metadata_manager,

src/common/meta/src/lib.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod affinity;
 pub mod cache;
 pub mod cache_invalidator;
 pub mod cluster;

src/common/meta/src/peer.rs

@@ -12,12 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::HashMap;
 use std::sync::Arc;
 
 pub use api::v1::meta::Peer;
 use api::v1::meta::heartbeat_request::NodeWorkloads;
 
+use crate::affinity::AffinityRequirements;
 use crate::cluster::NodeInfo;
 use crate::error::Error;
 use crate::{DatanodeId, FlownodeId};
@@ -80,7 +80,8 @@ pub type PeerDiscoveryRef = Arc<dyn PeerDiscovery>;
 
 #[derive(Debug, Clone, Default)]
 pub struct PeerAllocContext {
-    pub extensions: HashMap<String, String>,
+    /// Resolved requirements passed to node selection.
+    pub affinity: Option<AffinityRequirements>,
 }
 
 /// [`PeerAllocator`] allocates [`Peer`]s for creating region or flow.

src/common/meta/src/test_util.rs

@@ -198,6 +198,8 @@ pub fn new_ddl_context_with_kv_backend(
                 .build(),
         )));
     DdlContext {
+        affinity_resolver: None,
+        peer_alloc_context_provider: None,
         node_manager,
         cache_invalidator: Arc::new(DummyCacheInvalidator),
         memory_region_keeper: Arc::new(MemoryRegionKeeper::new()),

src/meta-srv/AGENTS.md

@@ -64,6 +64,12 @@ API under `src/meta-srv/src/service/admin/`.
   must be idempotent for crash recovery.
 - **Cache invalidation**: new metadata types need matching invalidation in
   `cache_invalidator.rs` and the heartbeat publish handler.
+- **Affinity**: `common-meta/affinity.rs` defines requirements, matching, batch
+  resolution, and the query allocation-context provider. The selector factory
+  optionally supplies a resolver for allocation inputs. DDL configurators can
+  install a query provider before procedure loaders are registered. Each selector
+  allows one unconstrained retry on candidate shortage only when explicitly
+  configured. Procedure state is unchanged.
 - **Selector stats fields**: keep `selector/` implementations in sync with the
   stats they consume.
 

src/meta-srv/src/metasrv.rs

@@ -26,6 +26,7 @@ use common_base::readable_size::ReadableSize;
 use common_config::{Configurable, DEFAULT_DATA_HOME};
 use common_event_recorder::EventRecorderOptions;
 use common_greptimedb_telemetry::GreptimeDBTelemetryTask;
+use common_meta::affinity::{AffinityFallback, AffinityRef, AffinityResolverRef};
 use common_meta::cache_invalidator::CacheInvalidatorRef;
 use common_meta::ddl::allocator::resource_id::ResourceIdAllocatorRef;
 use common_meta::ddl_manager::DdlManagerRef;
@@ -486,6 +487,10 @@ impl Display for SelectTarget {
 #[derive(Clone)]
 pub struct SelectorContext {
     pub peer_discovery: PeerDiscoveryRef,
+    /// Evaluator shared by datanode and flownode selection.
+    pub affinity: AffinityRef,
+    /// Whether selection may retry without affinity when candidates are insufficient.
+    pub affinity_fallback: AffinityFallback,
 }
 
 pub type SelectorRef = Arc<dyn Selector<Context = SelectorContext, Output = Vec<Peer>>>;
@@ -507,6 +512,14 @@ pub struct SelectorFactoryContext {
 /// Builds the final datanode selector metasrv should use.
 pub trait SelectorFactory: Send + Sync {
     fn build(&self, ctx: SelectorFactoryContext) -> SelectorRef;
+
+    /// Builds the resolver shared by allocation and failover callers.
+    fn build_affinity_resolver(
+        &self,
+        _meta_peer_client: MetaPeerClientRef,
+    ) -> Option<AffinityResolverRef> {
+        None
+    }
 }
 
 /// Shared selector factory plugin registered through [`common_base::Plugins`].

src/meta-srv/src/metasrv/builder.rs

@@ -22,6 +22,7 @@ use client::inserter::InsertOptions;
 use common_base::Plugins;
 use common_catalog::consts::{MIN_USER_FLOW_ID, MIN_USER_TABLE_ID};
 use common_event_recorder::{DEFAULT_COMPACTION_TIME_WINDOW, EventRecorderImpl, EventRecorderRef};
+use common_meta::affinity::{AffinityFallback, AffinityRef, NoopAffinity};
 use common_meta::ddl::flow_meta::FlowMetadataAllocator;
 use common_meta::ddl::table_meta::{TableMetadataAllocator, TableMetadataAllocatorRef};
 use common_meta::ddl::{
@@ -66,7 +67,8 @@ use crate::handler::region_lease_handler::{CustomizedRegionLeaseRenewerRef, Regi
 use crate::handler::{HeartbeatHandlerGroupBuilder, HeartbeatMailbox, Pushers};
 use crate::metasrv::{
     ElectionRef, FLOW_ID_SEQ, METASRV_DATA_DIR, Metasrv, MetasrvInfo, MetasrvOptions,
-    RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorRef, TABLE_ID_SEQ,
+    RegionStatAwareSelectorRef, SelectTarget, SelectorContext, SelectorFactoryRef, SelectorRef,
+    TABLE_ID_SEQ,
 };
 use crate::peer::MetasrvPeerAllocator;
 use crate::procedure::region_migration::DefaultContextFactory;
@@ -274,6 +276,14 @@ impl MetasrvBuilder {
         ));
 
         let selector_ctx = SelectorContext {
+            affinity_fallback: plugins
+                .as_ref()
+                .and_then(|plugins| plugins.get::<AffinityFallback>())
+                .unwrap_or_default(),
+            affinity: plugins
+                .as_ref()
+                .and_then(|plugins| plugins.get::<AffinityRef>())
+                .unwrap_or_else(|| Arc::new(NoopAffinity)),
             peer_discovery: meta_peer_client.clone(),
         };
 
@@ -385,6 +395,10 @@ impl MetasrvBuilder {
             ),
         ));
         region_migration_manager.try_start()?;
+        let affinity_resolver = plugins
+            .as_ref()
+            .and_then(|plugins| plugins.get::<SelectorFactoryRef>())
+            .and_then(|factory| factory.build_affinity_resolver(meta_peer_client.clone()));
         let region_supervisor_selector = plugins
             .as_ref()
             .and_then(|plugins| plugins.get::<RegionStatAwareSelectorRef>());
@@ -422,6 +436,8 @@ impl MetasrvBuilder {
         let topic_stats_registry = Arc::new(TopicStatsRegistry::default());
 
         let ddl_context = DdlContext {
+            peer_alloc_context_provider: None,
+            affinity_resolver,
             node_manager: node_manager.clone(),
             cache_invalidator: cache_invalidator.clone(),
             memory_region_keeper: memory_region_keeper.clone(),
@@ -795,10 +811,99 @@ mod tests {
         options.gc.enable = false;
         options.gc.max_concurrent_tables = 0;
 
-        MetasrvBuilder::new()
+        let metasrv = MetasrvBuilder::new()
             .options(options)
             .build()
             .await
             .unwrap();
+        let context = metasrv.ddl_manager().create_context();
+        assert!(context.affinity_resolver.is_none());
+        assert!(context.peer_alloc_context_provider.is_none());
+    }
+
+    #[tokio::test]
+    async fn test_builder_injects_affinity_components() {
+        use std::sync::atomic::{AtomicUsize, Ordering};
+
+        use common_meta::affinity::{
+            Affinity, AffinityRef, AffinityReference, AffinityRequirements, AffinityResolver,
+            AffinityResolverRef,
+        };
+        use common_meta::cluster::NodeInfo;
+        use common_meta::error::Result as MetaResult;
+
+        use crate::metasrv::{SelectorFactory, SelectorFactoryContext};
+
+        struct TestAffinity;
+
+        impl Affinity for TestAffinity {
+            fn matches(&self, _: &AffinityRequirements, _: &NodeInfo) -> bool {
+                true
+            }
+        }
+
+        struct TestResolver;
+
+        #[async_trait::async_trait]
+        impl AffinityResolver for TestResolver {
+            async fn resolve(
+                &self,
+                references: &[AffinityReference],
+            ) -> MetaResult<Vec<Option<AffinityRequirements>>> {
+                Ok(vec![None; references.len()])
+            }
+        }
+
+        struct RecordingFactory {
+            calls: Arc<AtomicUsize>,
+            resolver: AffinityResolverRef,
+        }
+
+        impl SelectorFactory for RecordingFactory {
+            fn build_affinity_resolver(&self, _: MetaPeerClientRef) -> Option<AffinityResolverRef> {
+                self.calls.fetch_add(1, Ordering::Relaxed);
+                Some(self.resolver.clone())
+            }
+
+            fn build(&self, ctx: SelectorFactoryContext) -> SelectorRef {
+                ctx.base_selector
+            }
+        }
+
+        let calls = Arc::new(AtomicUsize::new(0));
+        let resolver: AffinityResolverRef = Arc::new(TestResolver);
+        let affinity: AffinityRef = Arc::new(TestAffinity);
+        let plugins = Plugins::new();
+        plugins.insert(affinity.clone());
+        plugins.insert(AffinityFallback::Unconstrained);
+        plugins.insert(Arc::new(RecordingFactory {
+            calls: calls.clone(),
+            resolver: resolver.clone(),
+        }) as SelectorFactoryRef);
+        let metasrv = MetasrvBuilder::new()
+            .options(MetasrvOptions {
+                enable_telemetry: false,
+                ..Default::default()
+            })
+            .plugins(plugins)
+            .build()
+            .await
+            .unwrap();
+
+        assert_eq!(
+            metasrv.selector_ctx().affinity_fallback,
+            AffinityFallback::Unconstrained
+        );
+        assert_eq!(calls.load(Ordering::Relaxed), 1);
+        assert!(Arc::ptr_eq(&metasrv.selector_ctx().affinity, &affinity));
+        assert!(Arc::ptr_eq(
+            metasrv
+                .ddl_manager()
+                .create_context()
+                .affinity_resolver
+                .as_ref()
+                .unwrap(),
+            &resolver
+        ));
     }
 }

src/meta-srv/src/peer.rs

@@ -21,8 +21,8 @@ use common_meta::peer::{Peer, PeerAllocContext, PeerAllocator};
 use snafu::{ResultExt, ensure};
 
 use crate::discovery::utils::accept_ingest_workload;
-use crate::error::{Result, TooManyPartitionsSnafu};
-use crate::metasrv::{SelectorContext, SelectorRef};
+use crate::error::{self, Result, TooManyPartitionsSnafu};
+use crate::metasrv::{SelectTarget, SelectorContext, SelectorRef};
 use crate::selector::SelectorOptions;
 
 pub struct MetasrvPeerAllocator {
@@ -67,18 +67,23 @@ impl MetasrvPeerAllocator {
             );
         }
 
-        self.selector
-            .select(
-                &self.ctx,
-                SelectorOptions {
-                    min_required_items,
-                    allow_duplication: true,
-                    exclude_peer_ids: HashSet::new(),
-                    workload_filter: Some(accept_ingest_workload),
-                    extensions: alloc_context.extensions.clone(),
-                },
-            )
-            .await
+        let options = SelectorOptions {
+            affinity: alloc_context.affinity.as_ref(),
+            min_required_items,
+            allow_duplication: true,
+            exclude_peer_ids: HashSet::new(),
+            workload_filter: Some(accept_ingest_workload),
+        };
+        let peers = self.selector.select(&self.ctx, options).await?;
+        ensure!(
+            peers.len() >= min_required_items,
+            error::NoEnoughAvailableNodeSnafu {
+                required: min_required_items,
+                available: peers.len(),
+                select_target: SelectTarget::Datanode,
+            }
+        );
+        Ok(peers)
     }
 }
 
@@ -91,3 +96,187 @@ impl PeerAllocator for MetasrvPeerAllocator {
             .context(ExternalSnafu)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+    use std::ptr;
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+
+    use common_meta::affinity::{
+        Affinity, AffinityFallback, AffinityRef, AffinityRequirements, NoopAffinity,
+    };
+    use common_meta::cluster::NodeInfo;
+
+    use super::*;
+    use crate::selector::Selector;
+    use crate::selector::round_robin::RoundRobinSelector;
+    use crate::test_util::{create_meta_peer_client, put_datanodes};
+
+    struct OnlySecondNode;
+
+    impl Affinity for OnlySecondNode {
+        fn matches(&self, _: &AffinityRequirements, node: &NodeInfo) -> bool {
+            node.peer.id == 2
+        }
+    }
+
+    #[derive(Clone, Copy)]
+    enum Outcome {
+        Base,
+        BackendError,
+        Short,
+    }
+
+    struct CountingSelector {
+        base: RoundRobinSelector,
+        calls: AtomicUsize,
+        requirements_address: AtomicUsize,
+        outcome: Outcome,
+    }
+
+    #[async_trait]
+    impl Selector for CountingSelector {
+        type Context = SelectorContext;
+        type Output = Vec<Peer>;
+
+        async fn select(
+            &self,
+            ctx: &SelectorContext,
+            opts: SelectorOptions<'_>,
+        ) -> Result<Self::Output> {
+            self.calls.fetch_add(1, Ordering::Relaxed);
+            self.requirements_address.store(
+                opts.affinity
+                    .map_or(0, |requirements| ptr::from_ref(requirements) as usize),
+                Ordering::Relaxed,
+            );
+            assert!(opts.allow_duplication);
+            assert!(opts.exclude_peer_ids.is_empty());
+            assert!(opts.workload_filter.is_some());
+            assert_eq!(opts.min_required_items, 2);
+            match self.outcome {
+                Outcome::Base => self.base.select(ctx, opts).await,
+                Outcome::BackendError => error::UnexpectedSnafu {
+                    violated: "selection backend failed",
+                }
+                .fail(),
+                Outcome::Short => Ok(vec![Peer::empty(1)]),
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_allocator_delegates_fallback_and_validates_results() {
+        let discovery = create_meta_peer_client();
+        put_datanodes(&discovery, vec![Peer::empty(1)]).await;
+        for (has_requirements, fallback, noop, outcome, expected_calls, succeeds) in [
+            (
+                true,
+                AffinityFallback::Error,
+                false,
+                Outcome::Base,
+                1,
+                false,
+            ),
+            (
+                true,
+                AffinityFallback::Unconstrained,
+                false,
+                Outcome::Base,
+                1,
+                true,
+            ),
+            (true, AffinityFallback::Error, true, Outcome::Base, 1, true),
+            (
+                false,
+                AffinityFallback::Unconstrained,
+                false,
+                Outcome::Base,
+                1,
+                true,
+            ),
+            (
+                true,
+                AffinityFallback::Unconstrained,
+                false,
+                Outcome::BackendError,
+                1,
+                false,
+            ),
+            (
+                true,
+                AffinityFallback::Error,
+                false,
+                Outcome::Short,
+                1,
+                false,
+            ),
+            (
+                true,
+                AffinityFallback::Unconstrained,
+                false,
+                Outcome::Short,
+                1,
+                false,
+            ),
+            (
+                false,
+                AffinityFallback::Unconstrained,
+                false,
+                Outcome::Short,
+                1,
+                false,
+            ),
+        ] {
+            let selector = Arc::new(CountingSelector {
+                base: RoundRobinSelector::default(),
+                calls: AtomicUsize::new(0),
+                requirements_address: AtomicUsize::new(0),
+                outcome,
+            });
+            let affinity: AffinityRef = if noop {
+                Arc::new(NoopAffinity)
+            } else {
+                Arc::new(OnlySecondNode)
+            };
+            let allocator = MetasrvPeerAllocator::new(
+                SelectorContext {
+                    affinity_fallback: fallback,
+                    peer_discovery: discovery.clone(),
+                    affinity,
+                },
+                selector.clone(),
+            );
+            let alloc_context = PeerAllocContext {
+                affinity: has_requirements.then(|| {
+                    AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])).unwrap()
+                }),
+            };
+            let result = allocator.alloc(2, &alloc_context).await;
+            assert_eq!(
+                selector.requirements_address.load(Ordering::Relaxed),
+                alloc_context
+                    .affinity
+                    .as_ref()
+                    .map_or(0, |requirements| { ptr::from_ref(requirements) as usize }),
+            );
+            assert_eq!(selector.calls.load(Ordering::Relaxed), expected_calls);
+            if succeeds {
+                assert_eq!(result.unwrap(), vec![Peer::empty(1); 2]);
+            } else {
+                match outcome {
+                    Outcome::BackendError => assert!(matches!(
+                        result.unwrap_err(),
+                        error::Error::Unexpected { .. }
+                    )),
+                    _ => assert!(matches!(
+                        result.unwrap_err(),
+                        error::Error::NoEnoughAvailableNode { .. }
+                    )),
+                }
+            }
+        }
+    }
+}

src/meta-srv/src/procedure/utils.rs

@@ -535,6 +535,8 @@ pub mod test_data {
             Arc::new(SequenceBuilder::new("test", kv_backend).build()),
         ));
         DdlContext {
+            peer_alloc_context_provider: None,
+            affinity_resolver: None,
             node_manager,
             cache_invalidator: Arc::new(MetasrvCacheInvalidator::new(
                 mailbox,

src/meta-srv/src/region/supervisor.rs

@@ -667,7 +667,7 @@ impl RegionSupervisor {
                     allow_duplication: true,
                     exclude_peer_ids,
                     workload_filter: Some(accept_ingest_workload),
-                    extensions: Default::default(),
+                    affinity: None,
                 };
                 let peers = selector.select(&self.selector_context, opt).await?;
                 ensure!(

src/meta-srv/src/selector.rs

@@ -20,9 +20,10 @@ pub mod round_robin;
 pub(crate) mod test_utils;
 pub mod weight_compute;
 pub mod weighted_choose;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 
 use api::v1::meta::heartbeat_request::NodeWorkloads;
+use common_meta::affinity::AffinityRequirements;
 use serde::{Deserialize, Serialize};
 use store_api::storage::RegionId;
 use strum::AsRefStr;
@@ -35,7 +36,7 @@ pub trait Selector: Send + Sync {
     type Context;
     type Output;
 
-    async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result<Self::Output>;
+    async fn select(&self, ctx: &Self::Context, opts: SelectorOptions<'_>) -> Result<Self::Output>;
 }
 
 /// A selector that aware of region statistics
@@ -56,8 +57,10 @@ pub trait RegionStatAwareSelector: Send + Sync {
     ) -> Result<Self::Output>;
 }
 
-#[derive(Debug)]
-pub struct SelectorOptions {
+#[derive(Debug, Clone)]
+pub struct SelectorOptions<'a> {
+    /// Resolved requirements used to filter candidate nodes.
+    pub affinity: Option<&'a AffinityRequirements>,
     /// Minimum number of selected results.
     pub min_required_items: usize,
     /// Whether duplicates are allowed in the selected result, default false.
@@ -66,18 +69,16 @@ pub struct SelectorOptions {
     pub exclude_peer_ids: HashSet<u64>,
     /// The filter to select the peers based on their workloads.
     pub workload_filter: Option<fn(&NodeWorkloads) -> bool>,
-    /// Additional placement context, such as the frontend identity that originated the request.
-    pub extensions: HashMap<String, String>,
 }
 
-impl Default for SelectorOptions {
+impl Default for SelectorOptions<'_> {
     fn default() -> Self {
         Self {
+            affinity: None,
             min_required_items: 1,
             allow_duplication: false,
             exclude_peer_ids: HashSet::new(),
             workload_filter: None,
-            extensions: HashMap::new(),
         }
     }
 }

src/meta-srv/src/selector/common.rs

@@ -32,7 +32,7 @@ pub fn filter_out_excluded_peers(
 }
 
 /// According to the `opts`, choose peers from the `weight_array` through `weighted_choose`.
-pub fn choose_items<W>(opts: &SelectorOptions, weighted_choose: &mut W) -> Result<Vec<Peer>>
+pub fn choose_items<W>(opts: &SelectorOptions<'_>, weighted_choose: &mut W) -> Result<Vec<Peer>>
 where
     W: WeightedChoose<Peer>,
 {
@@ -86,13 +86,36 @@ where
 
 #[cfg(test)]
 mod tests {
+    use std::any::Any;
     use std::collections::{HashMap, HashSet};
+    use std::sync::Arc;
+    use std::sync::atomic::{AtomicUsize, Ordering};
 
-    use common_meta::peer::Peer;
+    use api::v1::meta::DatanodeWorkloads;
+    use api::v1::meta::heartbeat_request::NodeWorkloads;
+    use common_meta::affinity::{Affinity, AffinityFallback, AffinityRequirements, NoopAffinity};
+    use common_meta::cluster::{DatanodeStatus, NodeInfo, NodeStatus};
+    use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue, Stat};
+    use common_meta::error::{Error as MetaError, Result as MetaResult};
+    use common_meta::kv_backend::memory::MemoryKvBackend;
+    use common_meta::kv_backend::{KvBackend, KvBackendRef, ResettableKvBackend, TxnService};
+    use common_meta::peer::{Peer, PeerDiscovery};
+    use common_meta::rpc::store::{
+        BatchDeleteRequest, BatchDeleteResponse, BatchGetRequest, BatchGetResponse,
+        BatchPutRequest, BatchPutResponse, DeleteRangeRequest, DeleteRangeResponse, PutRequest,
+        PutResponse, RangeRequest, RangeResponse,
+    };
 
-    use crate::selector::SelectorOptions;
+    use crate::cluster::MetaPeerClientBuilder;
+    use crate::discovery::utils::accept_ingest_workload;
+    use crate::metasrv::{SelectTarget, SelectorContext, SelectorRef};
     use crate::selector::common::{choose_items, filter_out_excluded_peers};
+    use crate::selector::lease_based::LeaseBasedSelector;
+    use crate::selector::load_based::LoadBasedSelector;
+    use crate::selector::round_robin::RoundRobinSelector;
+    use crate::selector::weight_compute::RegionNumsBasedWeightCompute;
     use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem};
+    use crate::selector::{Selector, SelectorOptions};
 
     #[test]
     fn test_choose_peers() {
@@ -136,11 +159,11 @@ mod tests {
 
         for i in 1..=5 {
             let opts = SelectorOptions {
+                affinity: None,
                 min_required_items: i,
                 allow_duplication: false,
                 exclude_peer_ids: HashSet::new(),
                 workload_filter: None,
-                extensions: HashMap::new(),
             };
 
             let selected_peers: HashSet<_> =
@@ -153,11 +176,11 @@ mod tests {
         }
 
         let opts = SelectorOptions {
+            affinity: None,
             min_required_items: 6,
             allow_duplication: false,
             exclude_peer_ids: HashSet::new(),
             workload_filter: None,
-            extensions: HashMap::new(),
         };
 
         let selected_result =
@@ -166,11 +189,11 @@ mod tests {
 
         for i in 1..=50 {
             let opts = SelectorOptions {
+                affinity: None,
                 min_required_items: i,
                 allow_duplication: true,
                 exclude_peer_ids: HashSet::new(),
                 workload_filter: None,
-                extensions: HashMap::new(),
             };
 
             let selected_peers =
@@ -205,4 +228,508 @@ mod tests {
         assert_eq!(weight_array.len(), 1);
         assert_eq!(weight_array[0].item.id, 2);
     }
+
+    struct ExactAffinity;
+
+    impl Affinity for ExactAffinity {
+        fn matches(&self, requirements: &AffinityRequirements, candidate: &NodeInfo) -> bool {
+            requirements
+                .attributes()
+                .iter()
+                .all(|(key, value)| candidate.env_vars.get(key) == Some(value))
+        }
+    }
+
+    fn requirements() -> AffinityRequirements {
+        AffinityRequirements::new(HashMap::from([("pool".into(), "a".into())])).unwrap()
+    }
+
+    fn node(id: u64, pool: Option<&str>, accepts_ingest: bool) -> NodeInfo {
+        NodeInfo {
+            peer: Peer::empty(id),
+            last_activity_ts: 0,
+            status: NodeStatus::Datanode(DatanodeStatus {
+                rcus: 0,
+                wcus: 0,
+                leader_regions: 0,
+                follower_regions: 0,
+                workloads: DatanodeWorkloads {
+                    types: if accepts_ingest { vec![0] } else { vec![] },
+                },
+            }),
+            version: String::new(),
+            git_commit: String::new(),
+            start_time_ms: 0,
+            total_cpu_millicores: 0,
+            total_memory_bytes: 0,
+            cpu_usage_millicores: 0,
+            memory_usage_bytes: 0,
+            hostname: String::new(),
+            env_vars: pool
+                .map(|pool| HashMap::from([("pool".into(), pool.into())]))
+                .unwrap_or_default(),
+        }
+    }
+
+    struct Discovery {
+        snapshots: Vec<Vec<NodeInfo>>,
+        fail: bool,
+        reads: AtomicUsize,
+    }
+
+    #[async_trait::async_trait]
+    impl PeerDiscovery for Discovery {
+        async fn active_frontends(&self) -> MetaResult<Vec<NodeInfo>> {
+            Ok(vec![])
+        }
+
+        async fn active_datanodes(
+            &self,
+            filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
+        ) -> MetaResult<Vec<NodeInfo>> {
+            let index = self.reads.fetch_add(1, Ordering::Relaxed);
+            if self.fail {
+                return common_meta::error::UnexpectedSnafu {
+                    err_msg: "discovery unavailable",
+                }
+                .fail();
+            }
+            let index = index.min(self.snapshots.len() - 1);
+            Ok(self.snapshots[index]
+                .iter()
+                .filter(|node| {
+                    let NodeStatus::Datanode(status) = &node.status else {
+                        unreachable!()
+                    };
+                    filter.is_none_or(|filter| {
+                        filter(&NodeWorkloads::Datanode(status.workloads.clone()))
+                    })
+                })
+                .cloned()
+                .collect())
+        }
+
+        async fn active_flownodes(
+            &self,
+            filter: Option<for<'a> fn(&'a NodeWorkloads) -> bool>,
+        ) -> MetaResult<Vec<NodeInfo>> {
+            self.active_datanodes(filter).await
+        }
+    }
+
+    #[tokio::test]
+    async fn test_affinity_constrains_base_selectors() {
+        let affinity_requirements = requirements();
+        for lease_based in [false, true] {
+            for (use_exact_affinity, has_requirements) in
+                [(true, false), (true, true), (false, true)]
+            {
+                let enabled = use_exact_affinity && has_requirements;
+                let selector: SelectorRef = if lease_based {
+                    Arc::new(LeaseBasedSelector)
+                } else {
+                    Arc::new(RoundRobinSelector::new(SelectTarget::Datanode))
+                };
+                let ctx = SelectorContext {
+                    affinity_fallback: AffinityFallback::Error,
+                    affinity: if use_exact_affinity {
+                        Arc::new(ExactAffinity)
+                    } else {
+                        Arc::new(NoopAffinity)
+                    },
+                    peer_discovery: Arc::new(Discovery {
+                        snapshots: vec![vec![
+                            node(1, Some("b"), true),
+                            node(2, Some("a"), true),
+                            node(3, Some("a"), true),
+                            node(4, None, true),
+                            node(5, Some("a"), false),
+                        ]],
+                        reads: AtomicUsize::new(0),
+                        fail: false,
+                    }),
+                };
+                // Selecting a complete round makes the lease-based result deterministic as a set.
+                let peers = selector
+                    .select(
+                        &ctx,
+                        SelectorOptions {
+                            affinity: has_requirements.then_some(&affinity_requirements),
+                            min_required_items: if enabled { 1 } else { 3 },
+                            exclude_peer_ids: HashSet::from([3]),
+                            workload_filter: Some(accept_ingest_workload),
+                            ..Default::default()
+                        },
+                    )
+                    .await
+                    .unwrap();
+                let ids = peers
+                    .into_iter()
+                    .map(|peer| peer.id)
+                    .collect::<HashSet<_>>();
+                assert_eq!(
+                    ids,
+                    if enabled {
+                        HashSet::from([2])
+                    } else {
+                        HashSet::from([1, 2, 4])
+                    }
+                );
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_round_robin_filters_each_discovery_read_for_both_targets() {
+        let affinity_requirements = requirements();
+        for target in [SelectTarget::Datanode, SelectTarget::Flownode] {
+            let discovery = Arc::new(Discovery {
+                snapshots: vec![
+                    vec![node(1, Some("b"), true), node(2, Some("a"), true)],
+                    vec![
+                        node(1, Some("b"), true),
+                        node(2, Some("a"), true),
+                        node(3, None, true),
+                    ],
+                ],
+                reads: AtomicUsize::new(0),
+                fail: false,
+            });
+            let ctx = SelectorContext {
+                affinity_fallback: AffinityFallback::Error,
+                peer_discovery: discovery.clone(),
+                affinity: Arc::new(ExactAffinity),
+            };
+            let selector = RoundRobinSelector::new(target);
+            let opts = SelectorOptions {
+                affinity: Some(&affinity_requirements),
+                ..Default::default()
+            };
+            for _ in 0..2 {
+                assert_eq!(
+                    selector.select(&ctx, opts.clone()).await.unwrap(),
+                    vec![Peer::empty(2)]
+                );
+            }
+            assert_eq!(discovery.reads.load(Ordering::Relaxed), 2);
+        }
+    }
+
+    #[tokio::test]
+    async fn test_round_robin_excludes_peers_for_both_targets() {
+        let affinity_requirements = requirements();
+        for target in [SelectTarget::Datanode, SelectTarget::Flownode] {
+            let ctx = SelectorContext {
+                affinity_fallback: AffinityFallback::Error,
+                peer_discovery: Arc::new(Discovery {
+                    snapshots: vec![vec![node(1, Some("a"), true), node(2, Some("a"), true)]],
+                    reads: AtomicUsize::new(0),
+                    fail: false,
+                }),
+                affinity: Arc::new(ExactAffinity),
+            };
+            for affinity in [None, Some(&affinity_requirements)] {
+                // Keep the topology fixed and vary only the exclusion set.
+                for (excluded, expected) in [(vec![], vec![1, 2]), (vec![1], vec![2, 2])] {
+                    let selector = RoundRobinSelector::new(target);
+                    let peers = selector
+                        .select(
+                            &ctx,
+                            SelectorOptions {
+                                affinity,
+                                min_required_items: 2,
+                                allow_duplication: true,
+                                exclude_peer_ids: excluded.into_iter().collect(),
+                                ..Default::default()
+                            },
+                        )
+                        .await
+                        .unwrap();
+                    assert_eq!(
+                        peers.into_iter().map(|peer| peer.id).collect::<Vec<_>>(),
+                        expected
+                    );
+                }
+                let selector = RoundRobinSelector::new(target);
+                let error = selector
+                    .select(
+                        &ctx,
+                        SelectorOptions {
+                            affinity,
+                            exclude_peer_ids: HashSet::from([1, 2]),
+                            ..Default::default()
+                        },
+                    )
+                    .await
+                    .unwrap_err();
+                assert!(matches!(
+                    error,
+                    crate::error::Error::NoEnoughAvailableNode { available: 0, .. }
+                ));
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_base_selector_affinity_fallback() {
+        let affinity_requirements = requirements();
+        for lease_based in [false, true] {
+            for fallback in [AffinityFallback::Error, AffinityFallback::Unconstrained] {
+                for has_requirements in [false, true] {
+                    for required in [2, 3] {
+                        let discovery = Arc::new(Discovery {
+                            snapshots: vec![vec![
+                                node(1, Some("a"), true),
+                                node(2, Some("b"), true),
+                                node(3, Some("a"), true),
+                                node(4, Some("a"), false),
+                            ]],
+                            reads: AtomicUsize::new(0),
+                            fail: false,
+                        });
+                        let ctx = SelectorContext {
+                            peer_discovery: discovery.clone(),
+                            affinity: Arc::new(ExactAffinity),
+                            affinity_fallback: fallback,
+                        };
+                        let selector: SelectorRef = if lease_based {
+                            Arc::new(LeaseBasedSelector)
+                        } else {
+                            Arc::new(RoundRobinSelector::new(SelectTarget::Datanode))
+                        };
+                        let result = selector
+                            .select(
+                                &ctx,
+                                SelectorOptions {
+                                    affinity: has_requirements.then_some(&affinity_requirements),
+                                    min_required_items: required,
+                                    allow_duplication: false,
+                                    exclude_peer_ids: HashSet::from([3]),
+                                    workload_filter: Some(accept_ingest_workload),
+                                },
+                            )
+                            .await;
+                        let retries =
+                            has_requirements && fallback == AffinityFallback::Unconstrained;
+                        assert_eq!(
+                            discovery.reads.load(Ordering::Relaxed),
+                            1 + usize::from(retries)
+                        );
+                        let constrained = has_requirements && !retries;
+                        if required == 2 && !constrained {
+                            assert_eq!(
+                                result
+                                    .unwrap()
+                                    .into_iter()
+                                    .map(|peer| peer.id)
+                                    .collect::<HashSet<_>>(),
+                                HashSet::from([1, 2])
+                            );
+                        } else {
+                            let available = if constrained { 1 } else { 2 };
+                            assert!(
+                                matches!(result.unwrap_err(), crate::error::Error::NoEnoughAvailableNode {
+                                required: actual_required, available: actual_available, ..
+                            } if actual_required == required && actual_available == available)
+                            );
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_base_selector_discovery_error_does_not_fallback() {
+        let affinity_requirements = requirements();
+        for lease_based in [false, true] {
+            let discovery = Arc::new(Discovery {
+                snapshots: vec![],
+                reads: AtomicUsize::new(0),
+                fail: true,
+            });
+            let ctx = SelectorContext {
+                peer_discovery: discovery.clone(),
+                affinity: Arc::new(ExactAffinity),
+                affinity_fallback: AffinityFallback::Unconstrained,
+            };
+            let selector: SelectorRef = if lease_based {
+                Arc::new(LeaseBasedSelector)
+            } else {
+                Arc::new(RoundRobinSelector::new(SelectTarget::Datanode))
+            };
+            let error = selector
+                .select(
+                    &ctx,
+                    SelectorOptions {
+                        affinity: Some(&affinity_requirements),
+                        ..Default::default()
+                    },
+                )
+                .await
+                .unwrap_err();
+            assert!(matches!(
+                error,
+                crate::error::Error::ListActiveDatanodes { .. }
+            ));
+            assert_eq!(discovery.reads.load(Ordering::Relaxed), 1);
+        }
+    }
+
+    struct CountingStatBackend {
+        inner: MemoryKvBackend<MetaError>,
+        reads: AtomicUsize,
+        fail: bool,
+    }
+
+    #[async_trait::async_trait]
+    impl TxnService for CountingStatBackend {
+        type Error = MetaError;
+    }
+
+    #[async_trait::async_trait]
+    impl KvBackend for CountingStatBackend {
+        fn name(&self) -> &str {
+            "counting-stat-backend"
+        }
+
+        fn as_any(&self) -> &dyn Any {
+            self
+        }
+
+        async fn batch_get(&self, req: BatchGetRequest) -> MetaResult<BatchGetResponse> {
+            self.reads.fetch_add(1, Ordering::Relaxed);
+            if self.fail {
+                return common_meta::error::UnexpectedSnafu {
+                    err_msg: "stat backend unavailable",
+                }
+                .fail();
+            }
+            self.inner.batch_get(req).await
+        }
+
+        async fn range(&self, req: RangeRequest) -> MetaResult<RangeResponse> {
+            self.inner.range(req).await
+        }
+
+        async fn put(&self, req: PutRequest) -> MetaResult<PutResponse> {
+            self.inner.put(req).await
+        }
+
+        async fn batch_put(&self, req: BatchPutRequest) -> MetaResult<BatchPutResponse> {
+            self.inner.batch_put(req).await
+        }
+
+        async fn delete_range(&self, req: DeleteRangeRequest) -> MetaResult<DeleteRangeResponse> {
+            self.inner.delete_range(req).await
+        }
+
+        async fn batch_delete(&self, req: BatchDeleteRequest) -> MetaResult<BatchDeleteResponse> {
+            self.inner.batch_delete(req).await
+        }
+    }
+
+    impl ResettableKvBackend for CountingStatBackend {
+        fn reset(&self) {
+            self.inner.reset();
+        }
+
+        fn as_kv_backend_ref(self: Arc<Self>) -> KvBackendRef {
+            self
+        }
+    }
+
+    #[tokio::test]
+    async fn test_load_based_affinity_fallback_queries() {
+        let affinity_requirements = requirements();
+        for (has_requirements, fallback, required, fail_stats, expected_reads, available) in [
+            (false, AffinityFallback::Unconstrained, 2, false, 1, 2),
+            (true, AffinityFallback::Error, 2, false, 1, 1),
+            (true, AffinityFallback::Unconstrained, 2, false, 2, 2),
+            (true, AffinityFallback::Unconstrained, 3, false, 2, 2),
+            (true, AffinityFallback::Unconstrained, 2, true, 1, 0),
+        ] {
+            let backend = Arc::new(CountingStatBackend {
+                inner: MemoryKvBackend::new(),
+                reads: AtomicUsize::new(0),
+                fail: fail_stats,
+            });
+            for id in 1..=4 {
+                let value = DatanodeStatValue {
+                    stats: vec![Stat {
+                        id,
+                        addr: format!("node-{id}"),
+                        region_num: id,
+                        ..Default::default()
+                    }],
+                };
+                backend
+                    .put(
+                        PutRequest::new()
+                            .with_key(Vec::<u8>::from(DatanodeStatKey { node_id: id }))
+                            .with_value(Vec::<u8>::try_from(value).unwrap()),
+                    )
+                    .await
+                    .unwrap();
+            }
+            let client = Arc::new(
+                MetaPeerClientBuilder::default()
+                    .election(None)
+                    .in_memory(backend.clone())
+                    .build()
+                    .unwrap(),
+            );
+            let selector = LoadBasedSelector::new(RegionNumsBasedWeightCompute, client);
+            let discovery = Arc::new(Discovery {
+                snapshots: vec![vec![
+                    node(1, Some("a"), true),
+                    node(2, Some("b"), true),
+                    node(3, Some("a"), true),
+                    node(4, Some("a"), false),
+                ]],
+                reads: AtomicUsize::new(0),
+                fail: false,
+            });
+            let ctx = SelectorContext {
+                peer_discovery: discovery.clone(),
+                affinity: Arc::new(ExactAffinity),
+                affinity_fallback: fallback,
+            };
+            let result = selector
+                .select(
+                    &ctx,
+                    SelectorOptions {
+                        affinity: has_requirements.then_some(&affinity_requirements),
+                        min_required_items: required,
+                        allow_duplication: false,
+                        exclude_peer_ids: HashSet::from([3]),
+                        workload_filter: Some(accept_ingest_workload),
+                    },
+                )
+                .await;
+            assert_eq!(discovery.reads.load(Ordering::Relaxed), expected_reads);
+            assert_eq!(backend.reads.load(Ordering::Relaxed), expected_reads);
+            if fail_stats {
+                assert!(matches!(
+                    result.unwrap_err(),
+                    crate::error::Error::KvBackend { .. }
+                ));
+            } else if available < required {
+                assert!(
+                    matches!(result.unwrap_err(), crate::error::Error::NoEnoughAvailableNode {
+                    required: actual_required, available: actual_available, ..
+                } if actual_required == required && actual_available == available)
+                );
+            } else {
+                assert_eq!(
+                    result
+                        .unwrap()
+                        .into_iter()
+                        .map(|peer| peer.id)
+                        .collect::<HashSet<_>>(),
+                    HashSet::from([1, 2])
+                );
+            }
+        }
+    }
 }

src/meta-srv/src/selector/lease_based.rs

@@ -12,10 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use common_meta::affinity::AffinityFallback;
 use common_meta::peer::Peer;
 use snafu::ResultExt;
 
-use crate::error::{ListActiveDatanodesSnafu, Result};
+use crate::error::{self, ListActiveDatanodesSnafu, Result};
 use crate::metasrv::SelectorContext;
 use crate::selector::common::{choose_items, filter_out_excluded_peers};
 use crate::selector::weighted_choose::{RandomWeightedChoose, WeightedItem};
@@ -25,12 +26,12 @@ use crate::selector::{Selector, SelectorOptions};
 #[derive(Default)]
 pub struct LeaseBasedSelector;
 
-#[async_trait::async_trait]
-impl Selector for LeaseBasedSelector {
-    type Context = SelectorContext;
-    type Output = Vec<Peer>;
-
-    async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result<Self::Output> {
+impl LeaseBasedSelector {
+    async fn select_once(
+        &self,
+        ctx: &SelectorContext,
+        opts: &SelectorOptions<'_>,
+    ) -> Result<Vec<Peer>> {
         // 1. get alive datanodes.
         let alive_datanodes = ctx
             .peer_discovery
@@ -38,6 +39,10 @@ impl Selector for LeaseBasedSelector {
             .await
             .context(ListActiveDatanodesSnafu)?
             .into_iter()
+            .filter(|node| {
+                opts.affinity
+                    .is_none_or(|requirements| ctx.affinity.matches(requirements, node))
+            })
             .map(|node| node.peer)
             .collect::<Vec<_>>();
 
@@ -53,8 +58,35 @@ impl Selector for LeaseBasedSelector {
         // 3. choose peers by weight_array.
         filter_out_excluded_peers(&mut weight_array, &opts.exclude_peer_ids);
         let mut weighted_choose = RandomWeightedChoose::new(weight_array);
-        let selected = choose_items(&opts, &mut weighted_choose)?;
+        let selected = choose_items(opts, &mut weighted_choose)?;
 
         Ok(selected)
     }
 }
+
+#[async_trait::async_trait]
+impl Selector for LeaseBasedSelector {
+    type Context = SelectorContext;
+    type Output = Vec<Peer>;
+
+    async fn select(
+        &self,
+        ctx: &Self::Context,
+        mut opts: SelectorOptions<'_>,
+    ) -> Result<Self::Output> {
+        match self.select_once(ctx, &opts).await {
+            Err(error::Error::NoEnoughAvailableNode {
+                required,
+                available,
+                ..
+            }) if opts.affinity.is_some()
+                && ctx.affinity_fallback == AffinityFallback::Unconstrained
+                && available < required =>
+            {
+                opts.affinity = None;
+                self.select_once(ctx, &opts).await
+            }
+            result => result,
+        }
+    }
+}

src/meta-srv/src/selector/load_based.rs

@@ -14,13 +14,14 @@
 
 use std::collections::HashMap;
 
+use common_meta::affinity::AffinityFallback;
 use common_meta::datanode::{DatanodeStatKey, DatanodeStatValue};
 use common_meta::peer::Peer;
 use common_telemetry::debug;
 use snafu::ResultExt;
 
 use crate::cluster::MetaPeerClientRef;
-use crate::error::{ListActiveDatanodesSnafu, Result};
+use crate::error::{self, ListActiveDatanodesSnafu, Result};
 use crate::metasrv::SelectorContext;
 use crate::selector::common::{choose_items, filter_out_excluded_peers};
 use crate::selector::weight_compute::WeightCompute;
@@ -41,15 +42,15 @@ impl<C> LoadBasedSelector<C> {
     }
 }
 
-#[async_trait::async_trait]
-impl<C> Selector for LoadBasedSelector<C>
+impl<C> LoadBasedSelector<C>
 where
     C: WeightCompute<Source = HashMap<DatanodeStatKey, DatanodeStatValue>>,
 {
-    type Context = SelectorContext;
-    type Output = Vec<Peer>;
-
-    async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result<Self::Output> {
+    async fn select_once(
+        &self,
+        ctx: &SelectorContext,
+        opts: &SelectorOptions<'_>,
+    ) -> Result<Vec<Peer>> {
         // 1. get alive datanodes.
         let alive_datanodes = ctx
             .peer_discovery
@@ -57,6 +58,10 @@ where
             .await
             .context(ListActiveDatanodesSnafu)?
             .into_iter()
+            .filter(|node| {
+                opts.affinity
+                    .is_none_or(|requirements| ctx.affinity.matches(requirements, node))
+            })
             .map(|node| node.peer)
             .collect::<Vec<_>>();
 
@@ -77,7 +82,7 @@ where
         filter_out_excluded_peers(&mut weight_array, &opts.exclude_peer_ids);
         // 5. choose peers by weight_array.
         let mut weighted_choose = RandomWeightedChoose::new(weight_array);
-        let selected = choose_items(&opts, &mut weighted_choose)?;
+        let selected = choose_items(opts, &mut weighted_choose)?;
 
         debug!(
             "LoadBasedSelector select peers: {:?}, opts: {:?}.",
@@ -88,6 +93,36 @@ where
     }
 }
 
+#[async_trait::async_trait]
+impl<C> Selector for LoadBasedSelector<C>
+where
+    C: WeightCompute<Source = HashMap<DatanodeStatKey, DatanodeStatValue>>,
+{
+    type Context = SelectorContext;
+    type Output = Vec<Peer>;
+
+    async fn select(
+        &self,
+        ctx: &Self::Context,
+        mut opts: SelectorOptions<'_>,
+    ) -> Result<Self::Output> {
+        match self.select_once(ctx, &opts).await {
+            Err(error::Error::NoEnoughAvailableNode {
+                required,
+                available,
+                ..
+            }) if opts.affinity.is_some()
+                && ctx.affinity_fallback == AffinityFallback::Unconstrained
+                && available < required =>
+            {
+                opts.affinity = None;
+                self.select_once(ctx, &opts).await
+            }
+            result => result,
+        }
+    }
+}
+
 fn filter_out_expired_datanode(
     mut stat_kvs: HashMap<DatanodeStatKey, DatanodeStatValue>,
     datanodes: &[Peer],

src/meta-srv/src/selector/round_robin.rs

@@ -14,11 +14,12 @@
 
 use std::sync::atomic::AtomicUsize;
 
+use common_meta::affinity::AffinityFallback;
 use common_meta::peer::Peer;
 use snafu::{ResultExt, ensure};
 
 use crate::error::{
-    ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result,
+    self, ListActiveDatanodesSnafu, ListActiveFlownodesSnafu, NoEnoughAvailableNodeSnafu, Result,
 };
 use crate::metasrv::{SelectTarget, SelectorContext};
 use crate::selector::{Selector, SelectorOptions};
@@ -52,35 +53,42 @@ impl RoundRobinSelector {
         }
     }
 
-    async fn get_peers(&self, opts: &SelectorOptions, ctx: &SelectorContext) -> Result<Vec<Peer>> {
+    async fn get_peers(
+        &self,
+        opts: &SelectorOptions<'_>,
+        ctx: &SelectorContext,
+    ) -> Result<Vec<Peer>> {
         let mut peers = match self.select_target {
-            SelectTarget::Datanode => {
-                // 1. get alive datanodes.
-                let alive_datanodes = ctx
-                    .peer_discovery
-                    .active_datanodes(opts.workload_filter)
-                    .await
-                    .context(ListActiveDatanodesSnafu)?;
-
-                // 2. filter out excluded datanodes.
-                alive_datanodes
-                    .into_iter()
-                    .map(|node| node.peer)
-                    .filter(|p| !opts.exclude_peer_ids.contains(&p.id))
-                    .collect::<Vec<_>>()
-            }
+            SelectTarget::Datanode => ctx
+                .peer_discovery
+                .active_datanodes(opts.workload_filter)
+                .await
+                .context(ListActiveDatanodesSnafu)?
+                .into_iter()
+                .filter(|node| {
+                    opts.affinity
+                        .is_none_or(|requirements| ctx.affinity.matches(requirements, node))
+                })
+                .map(|node| node.peer)
+                .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id))
+                .collect::<Vec<_>>(),
             SelectTarget::Flownode => ctx
                 .peer_discovery
                 .active_flownodes(opts.workload_filter)
                 .await
                 .context(ListActiveFlownodesSnafu)?
                 .into_iter()
+                .filter(|node| {
+                    opts.affinity
+                        .is_none_or(|requirements| ctx.affinity.matches(requirements, node))
+                })
                 .map(|node| node.peer)
+                .filter(|peer| !opts.exclude_peer_ids.contains(&peer.id))
                 .collect::<Vec<_>>(),
         };
 
         ensure!(
-            !peers.is_empty(),
+            !peers.is_empty() && (opts.allow_duplication || peers.len() >= opts.min_required_items),
             NoEnoughAvailableNodeSnafu {
                 required: opts.min_required_items,
                 available: peers.len(),
@@ -93,15 +101,13 @@ impl RoundRobinSelector {
 
         Ok(peers)
     }
-}
 
-#[async_trait::async_trait]
-impl Selector for RoundRobinSelector {
-    type Context = SelectorContext;
-    type Output = Vec<Peer>;
-
-    async fn select(&self, ctx: &Self::Context, opts: SelectorOptions) -> Result<Vec<Peer>> {
-        let peers = self.get_peers(&opts, ctx).await?;
+    async fn select_once(
+        &self,
+        ctx: &SelectorContext,
+        opts: &SelectorOptions<'_>,
+    ) -> Result<Vec<Peer>> {
+        let peers = self.get_peers(opts, ctx).await?;
         // choose peers
         let mut selected = Vec::with_capacity(opts.min_required_items);
         for _ in 0..opts.min_required_items {
@@ -116,9 +122,39 @@ impl Selector for RoundRobinSelector {
     }
 }
 
+#[async_trait::async_trait]
+impl Selector for RoundRobinSelector {
+    type Context = SelectorContext;
+    type Output = Vec<Peer>;
+
+    async fn select(
+        &self,
+        ctx: &Self::Context,
+        mut opts: SelectorOptions<'_>,
+    ) -> Result<Self::Output> {
+        match self.select_once(ctx, &opts).await {
+            Err(error::Error::NoEnoughAvailableNode {
+                required,
+                available,
+                ..
+            }) if opts.affinity.is_some()
+                && ctx.affinity_fallback == AffinityFallback::Unconstrained
+                && available < required =>
+            {
+                opts.affinity = None;
+                self.select_once(ctx, &opts).await
+            }
+            result => result,
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::collections::HashSet;
+    use std::sync::Arc;
+
+    use common_meta::affinity::NoopAffinity;
 
     use super::*;
     use crate::test_util::{create_meta_peer_client, put_datanodes};
@@ -128,6 +164,8 @@ mod test {
         let selector = RoundRobinSelector::default();
         let meta_peer_client = create_meta_peer_client();
         let ctx = SelectorContext {
+            affinity_fallback: Default::default(),
+            affinity: Arc::new(NoopAffinity),
             peer_discovery: meta_peer_client.clone(),
         };
         // add three nodes
@@ -150,11 +188,11 @@ mod test {
             .select(
                 &ctx,
                 SelectorOptions {
+                    affinity: None,
                     min_required_items: 4,
                     allow_duplication: true,
                     exclude_peer_ids: HashSet::new(),
                     workload_filter: None,
-                    extensions: Default::default(),
                 },
             )
             .await
@@ -169,11 +207,11 @@ mod test {
             .select(
                 &ctx,
                 SelectorOptions {
+                    affinity: None,
                     min_required_items: 2,
                     allow_duplication: true,
                     exclude_peer_ids: HashSet::new(),
                     workload_filter: None,
-                    extensions: Default::default(),
                 },
             )
             .await
@@ -187,6 +225,8 @@ mod test {
         let selector = RoundRobinSelector::new(SelectTarget::Datanode);
         let meta_peer_client = create_meta_peer_client();
         let ctx = SelectorContext {
+            affinity_fallback: Default::default(),
+            affinity: Arc::new(NoopAffinity),
             peer_discovery: meta_peer_client.clone(),
         };
         // add three nodes
@@ -212,11 +252,11 @@ mod test {
             .select(
                 &ctx,
                 SelectorOptions {
+                    affinity: None,
                     min_required_items: 1,
                     allow_duplication: true,
                     exclude_peer_ids: HashSet::from([2, 5]),
                     workload_filter: None,
-                    extensions: Default::default(),
                 },
             )
             .await

src/meta-srv/src/selector/test_utils.rs

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use common_meta::affinity::NoopAffinity;
 use common_meta::kv_backend::memory::MemoryKvBackend;
 use common_meta::peer::Peer;
 use rand::prelude::SliceRandom;
@@ -34,6 +35,8 @@ pub fn new_test_selector_context() -> SelectorContext {
         .unwrap();
 
     SelectorContext {
+        affinity_fallback: Default::default(),
+        affinity: Arc::new(NoopAffinity),
         peer_discovery: meta_peer_client,
     }
 }
@@ -54,7 +57,11 @@ impl Selector for RandomNodeSelector {
     type Context = SelectorContext;
     type Output = Vec<Peer>;
 
-    async fn select(&self, _ctx: &Self::Context, _opts: SelectorOptions) -> Result<Self::Output> {
+    async fn select(
+        &self,
+        _ctx: &Self::Context,
+        _opts: SelectorOptions<'_>,
+    ) -> Result<Self::Output> {
         let mut rng = rand::rng();
         let mut nodes = self.nodes.clone();
         nodes.shuffle(&mut rng);

tests-integration/src/standalone.rs

@@ -281,6 +281,8 @@ impl GreptimeDbStandaloneBuilder {
         let repartition_procedure_factory = Arc::new(StandaloneRepartitionProcedureFactory);
         let ddl_manager = Arc::new(DdlManager::new(
             DdlContext {
+                affinity_resolver: None,
+                peer_alloc_context_provider: None,
                 node_manager: node_manager.clone(),
                 cache_invalidator: cache_registry.clone(),
                 memory_region_keeper: Arc::new(MemoryRegionKeeper::default()),

tests-integration/tests/region_migration.rs

@@ -1224,7 +1224,7 @@ impl Selector for ConstNodeSelector {
     async fn select(
         &self,
         _ctx: &Self::Context,
-        _opts: SelectorOptions,
+        _opts: SelectorOptions<'_>,
     ) -> MetaResult<Self::Output> {
         Ok(self.peers.clone())
     }