risingwave_meta/controller/

scale.rs

// Copyright 2024 RisingWave Labs
//
// 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.

use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
use std::num::NonZeroUsize;
use std::sync::atomic::{AtomicU32, Ordering};

use anyhow::anyhow;
use itertools::Itertools;
use risingwave_common::bitmap::Bitmap;
use risingwave_common::catalog::{FragmentTypeFlag, FragmentTypeMask, TableId};
use risingwave_common::id::JobId;
use risingwave_common::system_param::AdaptiveParallelismStrategy;
use risingwave_common::util::worker_util::DEFAULT_RESOURCE_GROUP;
use risingwave_connector::source::{SplitId, SplitImpl, SplitMetaData};
use risingwave_meta_model::fragment::DistributionType;
use risingwave_meta_model::prelude::{
    Database, Fragment, FragmentRelation, FragmentSplits, Sink, Source, StreamingJob, Table,
};
use risingwave_meta_model::{
    DatabaseId, DispatcherType, FragmentId, JobStatus, SourceId, StreamingParallelism, WorkerId,
    database, fragment, fragment_relation, fragment_splits, object, sink, source, streaming_job,
    table,
};
use risingwave_meta_model_migration::Condition;
use risingwave_pb::common::WorkerNode;
use risingwave_pb::stream_plan::PbStreamNode;
use sea_orm::{
    ColumnTrait, ConnectionTrait, EntityTrait, JoinType, QueryFilter, QuerySelect, QueryTrait,
    RelationTrait,
};

use crate::MetaResult;
use crate::controller::fragment::{InflightActorInfo, InflightFragmentInfo};
use crate::manager::ActiveStreamingWorkerNodes;
use crate::model::{ActorId, StreamActor, StreamingJobModelContextExt};
use crate::stream::{AssignerBuilder, SplitDiffOptions};

pub(crate) async fn resolve_streaming_job_definition<C>(
    txn: &C,
    job_ids: &HashSet<JobId>,
) -> MetaResult<HashMap<JobId, String>>
where
    C: ConnectionTrait,
{
    let job_ids = job_ids.iter().cloned().collect_vec();

    // including table, materialized view, index
    let common_job_definitions: Vec<(JobId, String)> = Table::find()
        .select_only()
        .columns([
            table::Column::TableId,
            #[cfg(not(debug_assertions))]
            table::Column::Name,
            #[cfg(debug_assertions)]
            table::Column::Definition,
        ])
        .filter(table::Column::TableId.is_in(job_ids.clone()))
        .into_tuple()
        .all(txn)
        .await?;

    let sink_definitions: Vec<(JobId, String)> = Sink::find()
        .select_only()
        .columns([
            sink::Column::SinkId,
            #[cfg(not(debug_assertions))]
            sink::Column::Name,
            #[cfg(debug_assertions)]
            sink::Column::Definition,
        ])
        .filter(sink::Column::SinkId.is_in(job_ids.clone()))
        .into_tuple()
        .all(txn)
        .await?;

    let source_definitions: Vec<(JobId, String)> = Source::find()
        .select_only()
        .columns([
            source::Column::SourceId,
            #[cfg(not(debug_assertions))]
            source::Column::Name,
            #[cfg(debug_assertions)]
            source::Column::Definition,
        ])
        .filter(source::Column::SourceId.is_in(job_ids.clone()))
        .into_tuple()
        .all(txn)
        .await?;

    let definitions: HashMap<JobId, String> = common_job_definitions
        .into_iter()
        .chain(sink_definitions.into_iter())
        .chain(source_definitions.into_iter())
        .collect();

    Ok(definitions)
}

pub async fn load_fragment_info<C>(
    txn: &C,
    actor_id_counter: &AtomicU32,
    database_id: Option<DatabaseId>,
    worker_nodes: &ActiveStreamingWorkerNodes,
    adaptive_parallelism_strategy: AdaptiveParallelismStrategy,
) -> MetaResult<FragmentRenderMap>
where
    C: ConnectionTrait,
{
    let mut query = StreamingJob::find()
        .select_only()
        .column(streaming_job::Column::JobId);

    if let Some(database_id) = database_id {
        query = query
            .join(JoinType::InnerJoin, streaming_job::Relation::Object.def())
            .filter(object::Column::DatabaseId.eq(database_id));
    }

    let jobs: Vec<JobId> = query.into_tuple().all(txn).await?;

    if jobs.is_empty() {
        return Ok(HashMap::new());
    }

    let jobs: HashSet<JobId> = jobs.into_iter().collect();

    let loaded = load_fragment_context_for_jobs(txn, jobs).await?;

    if loaded.is_empty() {
        return Ok(HashMap::new());
    }

    let RenderedGraph { fragments, .. } = render_actor_assignments(
        actor_id_counter,
        worker_nodes.current(),
        adaptive_parallelism_strategy,
        &loaded,
    )?;

    Ok(fragments)
}

#[derive(Debug)]
pub struct TargetResourcePolicy {
    pub resource_group: Option<String>,
    pub parallelism: StreamingParallelism,
}

#[derive(Debug, Clone)]
pub struct WorkerInfo {
    pub parallelism: NonZeroUsize,
    pub resource_group: Option<String>,
}

pub type FragmentRenderMap =
    HashMap<DatabaseId, HashMap<JobId, HashMap<FragmentId, InflightFragmentInfo>>>;

#[derive(Default)]
pub struct RenderedGraph {
    pub fragments: FragmentRenderMap,
    pub ensembles: Vec<NoShuffleEnsemble>,
}

impl RenderedGraph {
    pub fn empty() -> Self {
        Self::default()
    }
}

/// Context loaded asynchronously from database, containing all metadata
/// required to render actor assignments. This separates async I/O from
/// sync rendering logic.
#[derive(Clone, Debug)]
pub struct LoadedFragment {
    pub fragment_id: FragmentId,
    pub job_id: JobId,
    pub fragment_type_mask: FragmentTypeMask,
    pub distribution_type: DistributionType,
    pub vnode_count: usize,
    pub nodes: PbStreamNode,
    pub state_table_ids: HashSet<TableId>,
    pub parallelism: Option<StreamingParallelism>,
}

impl From<fragment::Model> for LoadedFragment {
    fn from(model: fragment::Model) -> Self {
        Self {
            fragment_id: model.fragment_id,
            job_id: model.job_id,
            fragment_type_mask: FragmentTypeMask::from(model.fragment_type_mask),
            distribution_type: model.distribution_type,
            vnode_count: model.vnode_count as usize,
            nodes: model.stream_node.to_protobuf(),
            state_table_ids: model.state_table_ids.into_inner().into_iter().collect(),
            parallelism: model.parallelism,
        }
    }
}

#[derive(Default, Debug, Clone)]
pub struct LoadedFragmentContext {
    pub ensembles: Vec<NoShuffleEnsemble>,
    pub job_fragments: HashMap<JobId, HashMap<FragmentId, LoadedFragment>>,
    pub job_map: HashMap<JobId, streaming_job::Model>,
    pub streaming_job_databases: HashMap<JobId, DatabaseId>,
    pub database_map: HashMap<DatabaseId, database::Model>,
    pub fragment_source_ids: HashMap<FragmentId, SourceId>,
    pub fragment_splits: HashMap<FragmentId, Vec<SplitImpl>>,
}

impl LoadedFragmentContext {
    pub fn is_empty(&self) -> bool {
        if self.ensembles.is_empty() {
            assert!(
                self.job_fragments.is_empty(),
                "non-empty job fragments for empty ensembles: {:?}",
                self.job_fragments
            );
            true
        } else {
            false
        }
    }

    pub fn for_database(&self, database_id: DatabaseId) -> Option<Self> {
        let job_ids: HashSet<JobId> = self
            .streaming_job_databases
            .iter()
            .filter_map(|(job_id, db_id)| (*db_id == database_id).then_some(*job_id))
            .collect();

        if job_ids.is_empty() {
            return None;
        }

        let job_fragments: HashMap<_, _> = job_ids
            .iter()
            .map(|job_id| (*job_id, self.job_fragments[job_id].clone()))
            .collect();

        let fragment_ids: HashSet<_> = job_fragments
            .values()
            .flat_map(|fragments| fragments.keys().copied())
            .collect();

        assert!(
            !fragment_ids.is_empty(),
            "empty fragments for non-empty database {database_id} with jobs {job_ids:?}"
        );

        let ensembles: Vec<NoShuffleEnsemble> = self
            .ensembles
            .iter()
            .filter(|ensemble| {
                if ensemble
                    .components
                    .iter()
                    .any(|fragment_id| fragment_ids.contains(fragment_id))
                {
                    assert!(
                        ensemble
                            .components
                            .iter()
                            .all(|fragment_id| fragment_ids.contains(fragment_id)),
                        "ensemble {ensemble:?} partially exists in database {database_id} with fragments {job_fragments:?}"
                    );
                    true
                } else {
                    false
                }
            })
            .cloned()
            .collect();

        assert!(
            !ensembles.is_empty(),
            "empty ensembles for non-empty database {database_id} with jobs {job_fragments:?}"
        );

        let job_map = job_ids
            .iter()
            .filter_map(|job_id| self.job_map.get(job_id).map(|job| (*job_id, job.clone())))
            .collect();

        let streaming_job_databases = job_ids
            .iter()
            .filter_map(|job_id| {
                self.streaming_job_databases
                    .get(job_id)
                    .map(|db_id| (*job_id, *db_id))
            })
            .collect();

        let database_model = self.database_map[&database_id].clone();
        let database_map = HashMap::from([(database_id, database_model)]);

        let fragment_source_ids = self
            .fragment_source_ids
            .iter()
            .filter(|(fragment_id, _)| fragment_ids.contains(*fragment_id))
            .map(|(fragment_id, source_id)| (*fragment_id, *source_id))
            .collect();

        let fragment_splits = self
            .fragment_splits
            .iter()
            .filter(|(fragment_id, _)| fragment_ids.contains(*fragment_id))
            .map(|(fragment_id, splits)| (*fragment_id, splits.clone()))
            .collect();

        Some(Self {
            ensembles,
            job_fragments,
            job_map,
            streaming_job_databases,
            database_map,
            fragment_source_ids,
            fragment_splits,
        })
    }

    /// Split this loaded context by database in a single ownership pass to avoid cloning large
    /// fragment payloads (for example `stream_node` in `LoadedFragment`).
    pub fn into_database_contexts(self) -> HashMap<DatabaseId, Self> {
        let Self {
            ensembles,
            mut job_fragments,
            mut job_map,
            streaming_job_databases,
            mut database_map,
            mut fragment_source_ids,
            mut fragment_splits,
        } = self;

        let mut contexts = HashMap::<DatabaseId, Self>::new();
        let mut fragment_databases = HashMap::<FragmentId, DatabaseId>::new();
        let mut unresolved_ensembles = 0usize;
        let mut unresolved_ensemble_sample: Option<Vec<FragmentId>> = None;

        for (job_id, database_id) in streaming_job_databases {
            let context = contexts.entry(database_id).or_insert_with(|| {
                let database_model = database_map
                    .remove(&database_id)
                    .expect("database should exist for streaming job");
                Self {
                    ensembles: Vec::new(),
                    job_fragments: HashMap::new(),
                    job_map: HashMap::new(),
                    streaming_job_databases: HashMap::new(),
                    database_map: HashMap::from([(database_id, database_model)]),
                    fragment_source_ids: HashMap::new(),
                    fragment_splits: HashMap::new(),
                }
            });

            let fragments = job_fragments
                .remove(&job_id)
                .expect("job fragments should exist for streaming job");
            for fragment_id in fragments.keys().copied() {
                fragment_databases.insert(fragment_id, database_id);
                if let Some(source_id) = fragment_source_ids.remove(&fragment_id) {
                    context.fragment_source_ids.insert(fragment_id, source_id);
                }
                if let Some(splits) = fragment_splits.remove(&fragment_id) {
                    context.fragment_splits.insert(fragment_id, splits);
                }
            }

            assert!(
                context
                    .job_map
                    .insert(
                        job_id,
                        job_map
                            .remove(&job_id)
                            .expect("streaming job should exist for loaded context"),
                    )
                    .is_none(),
                "duplicated streaming job"
            );
            assert!(
                context.job_fragments.insert(job_id, fragments).is_none(),
                "duplicated job fragments"
            );
            assert!(
                context
                    .streaming_job_databases
                    .insert(job_id, database_id)
                    .is_none(),
                "duplicated job database mapping"
            );
        }

        for ensemble in ensembles {
            let Some(database_id) = ensemble
                .components
                .iter()
                .find_map(|fragment_id| fragment_databases.get(fragment_id).copied())
            else {
                unresolved_ensembles += 1;
                if unresolved_ensemble_sample.is_none() {
                    unresolved_ensemble_sample =
                        Some(ensemble.components.iter().copied().collect());
                }
                continue;
            };

            debug_assert!(
                ensemble
                    .components
                    .iter()
                    .all(|fragment_id| fragment_databases.get(fragment_id) == Some(&database_id)),
                "ensemble {ensemble:?} should belong to a single database"
            );

            contexts
                .get_mut(&database_id)
                .expect("database context should exist for ensemble")
                .ensembles
                .push(ensemble);
        }

        if unresolved_ensembles > 0 {
            tracing::warn!(
                unresolved_ensembles,
                ?unresolved_ensemble_sample,
                known_fragments = fragment_databases.len(),
                "skip ensembles without resolved database while splitting loaded context"
            );
        }
        debug_assert_eq!(
            unresolved_ensembles, 0,
            "all ensembles should be mappable to a database"
        );

        contexts
    }
}

/// Async load stage for fragment-scoped rendering. It resolves all metadata required to later
/// render actor assignments with arbitrary worker sets.
pub async fn load_fragment_context<C>(
    txn: &C,
    ensembles: Vec<NoShuffleEnsemble>,
) -> MetaResult<LoadedFragmentContext>
where
    C: ConnectionTrait,
{
    if ensembles.is_empty() {
        return Ok(LoadedFragmentContext::default());
    }

    let required_fragment_ids: HashSet<_> = ensembles
        .iter()
        .flat_map(|ensemble| ensemble.components.iter().copied())
        .collect();

    let fragment_models = Fragment::find()
        .filter(fragment::Column::FragmentId.is_in(required_fragment_ids.iter().copied()))
        .all(txn)
        .await?;

    let found_fragment_ids: HashSet<_> = fragment_models
        .iter()
        .map(|fragment| fragment.fragment_id)
        .collect();

    if found_fragment_ids.len() != required_fragment_ids.len() {
        let missing = required_fragment_ids
            .difference(&found_fragment_ids)
            .copied()
            .collect_vec();
        return Err(anyhow!("fragments {:?} not found", missing).into());
    }

    let fragment_models: HashMap<_, _> = fragment_models
        .into_iter()
        .map(|fragment| (fragment.fragment_id, fragment))
        .collect();

    let job_ids: HashSet<_> = fragment_models
        .values()
        .map(|fragment| fragment.job_id)
        .collect();

    if job_ids.is_empty() {
        return Ok(LoadedFragmentContext::default());
    }

    let jobs: HashMap<_, _> = StreamingJob::find()
        .filter(streaming_job::Column::JobId.is_in(job_ids.iter().copied().collect_vec()))
        .all(txn)
        .await?
        .into_iter()
        .map(|job| (job.job_id, job))
        .collect();

    let found_job_ids: HashSet<_> = jobs.keys().copied().collect();
    if found_job_ids.len() != job_ids.len() {
        let missing = job_ids.difference(&found_job_ids).copied().collect_vec();
        return Err(anyhow!("streaming jobs {:?} not found", missing).into());
    }

    build_loaded_context(txn, ensembles, fragment_models, jobs).await
}

/// Async load stage for job-scoped rendering. It collects all no-shuffle ensembles and the
/// metadata required to render actor assignments later with a provided worker set.
pub async fn load_fragment_context_for_jobs<C>(
    txn: &C,
    job_ids: HashSet<JobId>,
) -> MetaResult<LoadedFragmentContext>
where
    C: ConnectionTrait,
{
    if job_ids.is_empty() {
        return Ok(LoadedFragmentContext::default());
    }

    let excluded_fragments_query = FragmentRelation::find()
        .select_only()
        .column(fragment_relation::Column::TargetFragmentId)
        .filter(fragment_relation::Column::DispatcherType.eq(DispatcherType::NoShuffle))
        .into_query();

    let condition = Condition::all()
        .add(fragment::Column::JobId.is_in(job_ids.clone()))
        .add(fragment::Column::FragmentId.not_in_subquery(excluded_fragments_query));

    let fragments: Vec<FragmentId> = Fragment::find()
        .select_only()
        .column(fragment::Column::FragmentId)
        .filter(condition)
        .into_tuple()
        .all(txn)
        .await?;

    let ensembles = find_fragment_no_shuffle_dags_detailed(txn, &fragments).await?;

    let fragments = Fragment::find()
        .filter(
            fragment::Column::FragmentId.is_in(
                ensembles
                    .iter()
                    .flat_map(|graph| graph.components.iter())
                    .cloned()
                    .collect_vec(),
            ),
        )
        .all(txn)
        .await?;

    let fragment_map: HashMap<_, _> = fragments
        .into_iter()
        .map(|fragment| (fragment.fragment_id, fragment))
        .collect();

    let job_ids = fragment_map
        .values()
        .map(|fragment| fragment.job_id)
        .collect::<BTreeSet<_>>()
        .into_iter()
        .collect_vec();

    let jobs: HashMap<_, _> = StreamingJob::find()
        .filter(streaming_job::Column::JobId.is_in(job_ids))
        .all(txn)
        .await?
        .into_iter()
        .map(|job| (job.job_id, job))
        .collect();

    build_loaded_context(txn, ensembles, fragment_map, jobs).await
}

/// Sync render stage: uses loaded fragment context and current worker info
/// to produce actor-to-worker assignments and vnode bitmaps.
pub(crate) fn render_actor_assignments(
    actor_id_counter: &AtomicU32,
    worker_map: &HashMap<WorkerId, WorkerNode>,
    adaptive_parallelism_strategy: AdaptiveParallelismStrategy,
    loaded: &LoadedFragmentContext,
) -> MetaResult<RenderedGraph> {
    if loaded.is_empty() {
        return Ok(RenderedGraph::empty());
    }

    let render_context = RenderActorsContext {
        fragment_source_ids: &loaded.fragment_source_ids,
        fragment_splits: &loaded.fragment_splits,
        streaming_job_databases: &loaded.streaming_job_databases,
        database_map: &loaded.database_map,
    };

    let fragments = render_actors(
        actor_id_counter,
        &loaded.ensembles,
        &loaded.job_fragments,
        &loaded.job_map,
        worker_map,
        adaptive_parallelism_strategy,
        render_context,
    )?;

    Ok(RenderedGraph {
        fragments,
        ensembles: loaded.ensembles.clone(),
    })
}

async fn build_loaded_context<C>(
    txn: &C,
    ensembles: Vec<NoShuffleEnsemble>,
    fragment_models: HashMap<FragmentId, fragment::Model>,
    job_map: HashMap<JobId, streaming_job::Model>,
) -> MetaResult<LoadedFragmentContext>
where
    C: ConnectionTrait,
{
    if ensembles.is_empty() {
        return Ok(LoadedFragmentContext::default());
    }

    let mut job_fragments: HashMap<JobId, HashMap<FragmentId, LoadedFragment>> = HashMap::new();
    for (fragment_id, model) in fragment_models {
        job_fragments
            .entry(model.job_id)
            .or_default()
            .try_insert(fragment_id, LoadedFragment::from(model))
            .expect("duplicate fragment id for job");
    }

    #[cfg(debug_assertions)]
    {
        debug_sanity_check(&ensembles, &job_fragments, &job_map);
    }

    let (fragment_source_ids, fragment_splits) =
        resolve_source_fragments(txn, &job_fragments).await?;

    let job_ids = job_map.keys().copied().collect_vec();

    let streaming_job_databases: HashMap<JobId, _> = StreamingJob::find()
        .select_only()
        .column(streaming_job::Column::JobId)
        .column(object::Column::DatabaseId)
        .join(JoinType::LeftJoin, streaming_job::Relation::Object.def())
        .filter(streaming_job::Column::JobId.is_in(job_ids))
        .into_tuple()
        .all(txn)
        .await?
        .into_iter()
        .collect();

    let database_map: HashMap<_, _> = Database::find()
        .filter(
            database::Column::DatabaseId
                .is_in(streaming_job_databases.values().copied().collect_vec()),
        )
        .all(txn)
        .await?
        .into_iter()
        .map(|db| (db.database_id, db))
        .collect();

    Ok(LoadedFragmentContext {
        ensembles,
        job_fragments,
        job_map,
        streaming_job_databases,
        database_map,
        fragment_source_ids,
        fragment_splits,
    })
}

// Only metadata resolved asynchronously lives here so the renderer stays synchronous
// and the call site keeps the runtime dependencies (maps, strategy, actor counter, etc.) explicit.
struct RenderActorsContext<'a> {
    fragment_source_ids: &'a HashMap<FragmentId, SourceId>,
    fragment_splits: &'a HashMap<FragmentId, Vec<SplitImpl>>,
    streaming_job_databases: &'a HashMap<JobId, DatabaseId>,
    database_map: &'a HashMap<DatabaseId, database::Model>,
}

fn render_actors(
    actor_id_counter: &AtomicU32,
    ensembles: &[NoShuffleEnsemble],
    job_fragments: &HashMap<JobId, HashMap<FragmentId, LoadedFragment>>,
    job_map: &HashMap<JobId, streaming_job::Model>,
    worker_map: &HashMap<WorkerId, WorkerNode>,
    adaptive_parallelism_strategy: AdaptiveParallelismStrategy,
    context: RenderActorsContext<'_>,
) -> MetaResult<FragmentRenderMap> {
    let RenderActorsContext {
        fragment_source_ids,
        fragment_splits: fragment_splits_map,
        streaming_job_databases,
        database_map,
    } = context;

    let mut all_fragments: FragmentRenderMap = HashMap::new();
    let fragment_lookup: HashMap<FragmentId, &LoadedFragment> = job_fragments
        .values()
        .flat_map(|fragments| fragments.iter())
        .map(|(fragment_id, fragment)| (*fragment_id, fragment))
        .collect();

    for NoShuffleEnsemble {
        entries,
        components,
    } in ensembles
    {
        tracing::debug!("rendering ensemble entries {:?}", entries);

        let entry_fragments = entries
            .iter()
            .map(|fragment_id| fragment_lookup.get(fragment_id).unwrap())
            .collect_vec();

        let entry_fragment_parallelism = entry_fragments
            .iter()
            .map(|fragment| fragment.parallelism.clone())
            .dedup()
            .exactly_one()
            .map_err(|_| {
                anyhow!(
                    "entry fragments {:?} have inconsistent parallelism settings",
                    entries.iter().copied().collect_vec()
                )
            })?;

        let (job_id, vnode_count) = entry_fragments
            .iter()
            .map(|f| (f.job_id, f.vnode_count))
            .dedup()
            .exactly_one()
            .map_err(|_| anyhow!("Multiple jobs found in no-shuffle ensemble"))?;

        let job = job_map
            .get(&job_id)
            .ok_or_else(|| anyhow!("streaming job {job_id} not found"))?;

        let database_resource_group = streaming_job_databases
            .get(&job_id)
            .and_then(|database_id| database_map.get(database_id))
            .unwrap()
            .resource_group
            .clone();

        let source_entry_fragment = entry_fragments.iter().find(|f| {
            let mask = f.fragment_type_mask;
            if mask.contains(FragmentTypeFlag::Source) {
                assert!(!mask.contains(FragmentTypeFlag::SourceScan))
            }
            mask.contains(FragmentTypeFlag::Source) && !mask.contains(FragmentTypeFlag::Dml)
        });

        let actor_template = EnsembleActorTemplate::render_new(
            job,
            worker_map,
            adaptive_parallelism_strategy,
            entry_fragment_parallelism,
            database_resource_group,
            vnode_count,
        )?;

        let source_splits = match source_entry_fragment {
            Some(entry_fragment) => {
                let source_id = fragment_source_ids
                    .get(&entry_fragment.fragment_id)
                    .ok_or_else(|| {
                        anyhow!(
                            "missing source id in source fragment {}",
                            entry_fragment.fragment_id
                        )
                    })?;

                let entry_fragment_id = entry_fragment.fragment_id;

                let splits = fragment_splits_map
                    .get(&entry_fragment_id)
                    .cloned()
                    .unwrap_or_default();

                let splits: std::collections::BTreeMap<_, _> =
                    splits.into_iter().map(|s| (s.id(), s)).collect();
                let splits = actor_template.assign_splits(entry_fragment_id, splits);
                Some((splits, *source_id))
            }
            None => None,
        };

        for component_fragment_id in components {
            let fragment = fragment_lookup.get(component_fragment_id).unwrap();
            let fragment_id = fragment.fragment_id;
            let job_id = fragment.job_id;
            let fragment_type_mask = fragment.fragment_type_mask;
            let distribution_type = fragment.distribution_type;
            let stream_node = &fragment.nodes;
            let state_table_ids = &fragment.state_table_ids;
            let vnode_count = fragment.vnode_count;
            let source_id = fragment_source_ids.get(&fragment_id).cloned();

            let aligner = ComponentFragmentAligner::new(&actor_template, actor_id_counter);
            let actors = aligner.align_component_actor(distribution_type);
            let mut splits = source_id
                .map(|source_id| {
                    let (fragment_splits, shared_source_id) = source_splits.as_ref().unwrap();
                    assert_eq!(*shared_source_id, source_id);
                    aligner.align_component_splits(fragment_splits)
                })
                .unwrap_or_default();

            let actors: HashMap<ActorId, InflightActorInfo> = actors
                .into_iter()
                .map(|(actor_id, (worker_id, vnode_bitmap))| {
                    (
                        actor_id,
                        InflightActorInfo {
                            worker_id,
                            vnode_bitmap,
                            splits: splits.remove(&actor_id).unwrap_or_default(),
                        },
                    )
                })
                .collect();

            let fragment = InflightFragmentInfo {
                fragment_id,
                distribution_type,
                fragment_type_mask,
                vnode_count,
                nodes: stream_node.clone(),
                actors,
                state_table_ids: state_table_ids.clone(),
            };

            let &database_id = streaming_job_databases.get(&job_id).ok_or_else(|| {
                anyhow!("streaming job {job_id} not found in streaming_job_databases")
            })?;

            all_fragments
                .entry(database_id)
                .or_default()
                .entry(job_id)
                .or_default()
                .insert(fragment_id, fragment);
        }
    }

    Ok(all_fragments)
}

pub(crate) struct EnsembleActorTemplate {
    assignment: BTreeMap<WorkerId, BTreeMap<u32, Vec<usize>>>,
    actor_count: u32,
    vnode_count: usize,
}

impl EnsembleActorTemplate {
    pub(crate) fn render_new(
        job: &streaming_job::Model,
        worker_map: &HashMap<WorkerId, WorkerNode>,
        adaptive_parallelism_strategy: AdaptiveParallelismStrategy,
        entry_fragment_parallelism: Option<StreamingParallelism>,
        database_resource_group: String,
        vnode_count: usize,
    ) -> MetaResult<Self> {
        let job_id = job.job_id;
        let job_strategy = job
            .stream_context()
            .adaptive_parallelism_strategy
            .unwrap_or(adaptive_parallelism_strategy);

        let resource_group = match &job.specific_resource_group {
            None => database_resource_group,
            Some(resource_group) => resource_group.clone(),
        };

        let available_workers: BTreeMap<WorkerId, NonZeroUsize> = worker_map
            .iter()
            .filter_map(|(worker_id, worker)| {
                if worker
                    .resource_group()
                    .as_deref()
                    .unwrap_or(DEFAULT_RESOURCE_GROUP)
                    == resource_group.as_str()
                {
                    Some((
                        *worker_id,
                        worker
                            .parallelism()
                            .expect("should have parallelism for compute node")
                            .try_into()
                            .expect("parallelism for compute node"),
                    ))
                } else {
                    None
                }
            })
            .collect();

        let total_parallelism = available_workers.values().map(|w| w.get()).sum::<usize>();

        let effective_job_parallelism = if job.job_status != JobStatus::Created {
            job.backfill_parallelism
                .as_ref()
                .unwrap_or(&job.parallelism)
        } else {
            &job.parallelism
        };

        let actual_parallelism = match entry_fragment_parallelism
            .as_ref()
            .unwrap_or(effective_job_parallelism)
        {
            StreamingParallelism::Adaptive | StreamingParallelism::Custom => {
                job_strategy.compute_target_parallelism(total_parallelism)
            }
            StreamingParallelism::Fixed(n) => *n,
        }
        .min(vnode_count)
        .min(job.max_parallelism as usize);

        tracing::debug!(
            "job {}, final {} parallelism {:?} total_parallelism {} job_max {} vnode count {} fragment_override {:?}",
            job_id,
            actual_parallelism,
            job.parallelism,
            total_parallelism,
            job.max_parallelism,
            vnode_count,
            entry_fragment_parallelism
        );

        let assigner = AssignerBuilder::new(job_id).build();

        let actors = (0..(actual_parallelism as u32)).collect_vec();
        let vnodes = (0..vnode_count).collect_vec();

        let assignment = assigner.assign_hierarchical(&available_workers, &actors, &vnodes)?;

        let actor_count = u32::try_from(actors.len()).expect("actor parallelism exceeds u32::MAX");

        Ok(Self {
            assignment,
            actor_count,
            vnode_count,
        })
    }

    fn assign_splits(
        &self,
        entry_fragment_id: FragmentId,
        splits: BTreeMap<SplitId, SplitImpl>,
    ) -> HashMap<u32, Vec<SplitImpl>> {
        {
            {
                let empty_actor_splits: HashMap<_, _> = self
                    .assignment
                    .values()
                    .flat_map(|actors| actors.keys())
                    .map(|actor_id| (*actor_id, vec![]))
                    .collect();

                crate::stream::source_manager::reassign_splits(
                    entry_fragment_id,
                    empty_actor_splits,
                    &splits,
                    SplitDiffOptions::default(),
                )
                .unwrap_or_default()
            }
        }
    }
}

pub(crate) struct ComponentFragmentAligner<'a> {
    actor_template: &'a EnsembleActorTemplate,
    actor_id_base: ActorId,
}

impl<'a> ComponentFragmentAligner<'a> {
    pub(crate) fn new(
        actor_template: &'a EnsembleActorTemplate,
        actor_id_counter: &AtomicU32,
    ) -> Self {
        let actor_id_base = actor_id_counter
            .fetch_add(actor_template.actor_count, Ordering::Relaxed)
            .into();
        Self {
            actor_template,
            actor_id_base,
        }
    }

    pub(crate) fn align_component_actor(
        &self,
        distribution_type: DistributionType,
    ) -> HashMap<ActorId, (WorkerId, Option<Bitmap>)> {
        let EnsembleActorTemplate {
            assignment,
            vnode_count,
            actor_count,
        } = &self.actor_template;
        let actor_id_base = self.actor_id_base;
        let vnode_count = *vnode_count;
        {
            assignment
                .iter()
                .flat_map(|(worker_id, actors)| {
                    actors
                        .iter()
                        .map(move |(actor_id, vnodes)| (worker_id, actor_id, vnodes))
                })
                .map(|(&worker_id, &actor_idx, vnodes)| {
                    let vnode_bitmap = match distribution_type {
                        DistributionType::Single => {
                            assert_eq!(*actor_count, 1);
                            None
                        }
                        DistributionType::Hash => Some(Bitmap::from_indices(vnode_count, vnodes)),
                    };

                    let actor_id = actor_id_base + actor_idx;

                    (actor_id, (worker_id, vnode_bitmap))
                })
                .collect()
        }
    }

    pub(crate) fn align_component_splits(
        &self,
        split_assignment: &HashMap<u32, Vec<SplitImpl>>,
    ) -> HashMap<ActorId, Vec<SplitImpl>> {
        (0..self.actor_template.actor_count)
            .filter_map(|actor_idx| {
                split_assignment
                    .get(&actor_idx)
                    .map(|splits| ((self.actor_id_base + actor_idx), splits.clone()))
            })
            .collect()
    }
}

#[cfg(debug_assertions)]
fn debug_sanity_check(
    ensembles: &[NoShuffleEnsemble],
    job_fragments: &HashMap<JobId, HashMap<FragmentId, LoadedFragment>>,
    jobs: &HashMap<JobId, streaming_job::Model>,
) {
    let fragment_lookup: HashMap<FragmentId, (&LoadedFragment, JobId)> = job_fragments
        .iter()
        .flat_map(|(job_id, fragments)| {
            fragments
                .iter()
                .map(move |(fragment_id, fragment)| (*fragment_id, (fragment, *job_id)))
        })
        .collect();

    // Debug-only assertions to catch inconsistent ensemble metadata early.
    debug_assert!(
        ensembles
            .iter()
            .all(|ensemble| ensemble.entries.is_subset(&ensemble.components)),
        "entries must be subset of components"
    );

    let mut missing_fragments = BTreeSet::new();
    let mut missing_jobs = BTreeSet::new();

    for fragment_id in ensembles
        .iter()
        .flat_map(|ensemble| ensemble.components.iter())
    {
        match fragment_lookup.get(fragment_id) {
            Some((fragment, job_id)) => {
                if !jobs.contains_key(&fragment.job_id) {
                    missing_jobs.insert(*job_id);
                }
            }
            None => {
                missing_fragments.insert(*fragment_id);
            }
        }
    }

    debug_assert!(
        missing_fragments.is_empty(),
        "missing fragments in fragment_map: {:?}",
        missing_fragments
    );

    debug_assert!(
        missing_jobs.is_empty(),
        "missing jobs for fragments' job_id: {:?}",
        missing_jobs
    );

    for ensemble in ensembles {
        let unique_vnode_counts: Vec<_> = ensemble
            .components
            .iter()
            .flat_map(|fragment_id| {
                fragment_lookup
                    .get(fragment_id)
                    .map(|(fragment, _)| fragment.vnode_count)
            })
            .unique()
            .collect();

        debug_assert!(
            unique_vnode_counts.len() <= 1,
            "components in ensemble must share same vnode_count: ensemble={:?}, vnode_counts={:?}",
            ensemble.components,
            unique_vnode_counts
        );
    }
}

async fn resolve_source_fragments<C>(
    txn: &C,
    job_fragments: &HashMap<JobId, HashMap<FragmentId, LoadedFragment>>,
) -> MetaResult<(
    HashMap<FragmentId, SourceId>,
    HashMap<FragmentId, Vec<SplitImpl>>,
)>
where
    C: ConnectionTrait,
{
    let mut source_fragment_ids: HashMap<SourceId, BTreeSet<FragmentId>> = HashMap::new();
    for (fragment_id, fragment) in job_fragments.values().flatten() {
        let mask = fragment.fragment_type_mask;
        if mask.contains(FragmentTypeFlag::Source)
            && let Some(source_id) = fragment.nodes.find_stream_source()
        {
            source_fragment_ids
                .entry(source_id)
                .or_default()
                .insert(*fragment_id);
        }

        if mask.contains(FragmentTypeFlag::SourceScan)
            && let Some((source_id, _)) = fragment.nodes.find_source_backfill()
        {
            source_fragment_ids
                .entry(source_id)
                .or_default()
                .insert(*fragment_id);
        }
    }

    let fragment_source_ids: HashMap<_, _> = source_fragment_ids
        .iter()
        .flat_map(|(source_id, fragment_ids)| {
            fragment_ids
                .iter()
                .map(|fragment_id| (*fragment_id, *source_id))
        })
        .collect();

    let fragment_ids = fragment_source_ids.keys().copied().collect_vec();

    let fragment_splits: Vec<_> = FragmentSplits::find()
        .filter(fragment_splits::Column::FragmentId.is_in(fragment_ids))
        .all(txn)
        .await?;

    let fragment_splits: HashMap<_, _> = fragment_splits
        .into_iter()
        .flat_map(|model| {
            model.splits.map(|splits| {
                (
                    model.fragment_id,
                    splits
                        .to_protobuf()
                        .splits
                        .iter()
                        .flat_map(SplitImpl::try_from)
                        .collect_vec(),
                )
            })
        })
        .collect();

    Ok((fragment_source_ids, fragment_splits))
}

// Helper struct to make the function signature cleaner and to properly bundle the required data.
#[derive(Debug)]
pub struct ActorGraph<'a> {
    pub fragments: &'a HashMap<FragmentId, (Fragment, Vec<StreamActor>)>,
    pub locations: &'a HashMap<ActorId, WorkerId>,
}

#[derive(Debug, Clone)]
pub struct NoShuffleEnsemble {
    entries: HashSet<FragmentId>,
    components: HashSet<FragmentId>,
}

impl NoShuffleEnsemble {
    #[cfg(test)]
    pub fn for_test(
        entries: impl IntoIterator<Item = FragmentId>,
        components: impl IntoIterator<Item = FragmentId>,
    ) -> Self {
        let entries = entries.into_iter().collect();
        let components = components.into_iter().collect();
        Self {
            entries,
            components,
        }
    }

    pub fn fragments(&self) -> impl Iterator<Item = FragmentId> + '_ {
        self.components.iter().cloned()
    }

    pub fn entry_fragments(&self) -> impl Iterator<Item = FragmentId> + '_ {
        self.entries.iter().copied()
    }

    pub fn component_fragments(&self) -> impl Iterator<Item = FragmentId> + '_ {
        self.components.iter().copied()
    }

    pub fn contains_entry(&self, fragment_id: &FragmentId) -> bool {
        self.entries.contains(fragment_id)
    }
}

pub async fn find_fragment_no_shuffle_dags_detailed(
    db: &impl ConnectionTrait,
    initial_fragment_ids: &[FragmentId],
) -> MetaResult<Vec<NoShuffleEnsemble>> {
    let all_no_shuffle_relations: Vec<(_, _)> = FragmentRelation::find()
        .columns([
            fragment_relation::Column::SourceFragmentId,
            fragment_relation::Column::TargetFragmentId,
        ])
        .filter(fragment_relation::Column::DispatcherType.eq(DispatcherType::NoShuffle))
        .into_tuple()
        .all(db)
        .await?;

    let (forward_edges, backward_edges) =
        build_no_shuffle_fragment_graph_edges(all_no_shuffle_relations);

    find_no_shuffle_graphs(initial_fragment_ids, &forward_edges, &backward_edges)
}

pub(crate) fn build_no_shuffle_fragment_graph_edges(
    relations: impl IntoIterator<Item = (FragmentId, FragmentId)>,
) -> (
    HashMap<FragmentId, Vec<FragmentId>>,
    HashMap<FragmentId, Vec<FragmentId>>,
) {
    let mut forward_edges: HashMap<FragmentId, HashSet<FragmentId>> = HashMap::new();
    let mut backward_edges: HashMap<FragmentId, HashSet<FragmentId>> = HashMap::new();

    for (src, dst) in relations {
        forward_edges.entry(src).or_default().insert(dst);
        backward_edges.entry(dst).or_default().insert(src);
    }

    let forward_edges = forward_edges
        .into_iter()
        .map(|(src, dst_set)| (src, dst_set.into_iter().collect()))
        .collect();
    let backward_edges = backward_edges
        .into_iter()
        .map(|(dst, src_set)| (dst, src_set.into_iter().collect()))
        .collect();

    (forward_edges, backward_edges)
}

pub(crate) fn find_no_shuffle_graphs(
    initial_fragment_ids: &[impl Into<FragmentId> + Copy],
    forward_edges: &HashMap<FragmentId, Vec<FragmentId>>,
    backward_edges: &HashMap<FragmentId, Vec<FragmentId>>,
) -> MetaResult<Vec<NoShuffleEnsemble>> {
    let mut graphs: Vec<NoShuffleEnsemble> = Vec::new();
    let mut globally_visited: HashSet<FragmentId> = HashSet::new();

    for &init_id in initial_fragment_ids {
        let init_id = init_id.into();
        if globally_visited.contains(&init_id) {
            continue;
        }

        // Found a new component. Traverse it to find all its nodes.
        let mut components = HashSet::new();
        let mut queue: VecDeque<FragmentId> = VecDeque::new();

        queue.push_back(init_id);
        globally_visited.insert(init_id);

        while let Some(current_id) = queue.pop_front() {
            components.insert(current_id);
            let neighbors = forward_edges
                .get(&current_id)
                .into_iter()
                .flatten()
                .chain(backward_edges.get(&current_id).into_iter().flatten());

            for &neighbor_id in neighbors {
                if globally_visited.insert(neighbor_id) {
                    queue.push_back(neighbor_id);
                }
            }
        }

        // For the newly found component, identify its roots.
        let mut entries = HashSet::new();
        for &node_id in &components {
            let is_root = match backward_edges.get(&node_id) {
                Some(parents) => parents.iter().all(|p| !components.contains(p)),
                None => true,
            };
            if is_root {
                entries.insert(node_id);
            }
        }

        // Store the detailed DAG structure (roots, all nodes in this DAG).
        if !entries.is_empty() {
            graphs.push(NoShuffleEnsemble {
                entries,
                components,
            });
        }
    }

    Ok(graphs)
}

#[cfg(test)]
mod tests {
    use std::collections::{BTreeSet, HashMap, HashSet};
    use std::sync::Arc;

    use risingwave_connector::source::SplitImpl;
    use risingwave_connector::source::test_source::TestSourceSplit;
    use risingwave_meta_model::{CreateType, JobStatus};
    use risingwave_pb::common::WorkerType;
    use risingwave_pb::common::worker_node::Property as WorkerProperty;
    use risingwave_pb::stream_plan::StreamNode as PbStreamNode;

    use super::*;

    // Helper type aliases for cleaner test code
    // Using the actual FragmentId type from the module
    type Edges = (
        HashMap<FragmentId, Vec<FragmentId>>,
        HashMap<FragmentId, Vec<FragmentId>>,
    );

    /// A helper function to build forward and backward edge maps from a simple list of tuples.
    /// This reduces boilerplate in each test.
    fn build_edges(relations: &[(u32, u32)]) -> Edges {
        let mut forward_edges: HashMap<FragmentId, Vec<FragmentId>> = HashMap::new();
        let mut backward_edges: HashMap<FragmentId, Vec<FragmentId>> = HashMap::new();
        for &(src, dst) in relations {
            forward_edges
                .entry(src.into())
                .or_default()
                .push(dst.into());
            backward_edges
                .entry(dst.into())
                .or_default()
                .push(src.into());
        }
        (forward_edges, backward_edges)
    }

    /// Helper function to create a `HashSet` from a slice easily.
    fn to_hashset(ids: &[u32]) -> HashSet<FragmentId> {
        ids.iter().map(|id| (*id).into()).collect()
    }

    fn build_fragment(
        fragment_id: FragmentId,
        job_id: JobId,
        fragment_type_mask: i32,
        distribution_type: DistributionType,
        vnode_count: i32,
        parallelism: StreamingParallelism,
    ) -> LoadedFragment {
        LoadedFragment {
            fragment_id,
            job_id,
            fragment_type_mask: FragmentTypeMask::from(fragment_type_mask),
            distribution_type,
            vnode_count: vnode_count as usize,
            nodes: PbStreamNode::default(),
            state_table_ids: HashSet::new(),
            parallelism: Some(parallelism),
        }
    }

    type ActorState = (ActorId, WorkerId, Option<Vec<usize>>, Vec<String>);

    fn collect_actor_state(fragment: &InflightFragmentInfo) -> Vec<ActorState> {
        let base = fragment.actors.keys().copied().min().unwrap_or_default();

        let mut entries: Vec<_> = fragment
            .actors
            .iter()
            .map(|(&actor_id, info)| {
                let idx = actor_id.as_raw_id() - base.as_raw_id();
                let vnode_indices = info.vnode_bitmap.as_ref().map(|bitmap| {
                    bitmap
                        .iter()
                        .enumerate()
                        .filter_map(|(pos, is_set)| is_set.then_some(pos))
                        .collect::<Vec<_>>()
                });
                let splits = info
                    .splits
                    .iter()
                    .map(|split| split.id().to_string())
                    .collect::<Vec<_>>();
                (idx.into(), info.worker_id, vnode_indices, splits)
            })
            .collect();

        entries.sort_by_key(|(idx, _, _, _)| *idx);
        entries
    }

    fn build_worker_node(
        id: impl Into<WorkerId>,
        parallelism: usize,
        resource_group: &str,
    ) -> WorkerNode {
        WorkerNode {
            id: id.into(),
            r#type: WorkerType::ComputeNode as i32,
            property: Some(WorkerProperty {
                is_streaming: true,
                parallelism: u32::try_from(parallelism).expect("parallelism fits into u32"),
                resource_group: Some(resource_group.to_owned()),
                ..Default::default()
            }),
            ..Default::default()
        }
    }

    #[test]
    fn test_single_linear_chain() {
        // Scenario: A simple linear graph 1 -> 2 -> 3.
        // We start from the middle node (2).
        let (forward, backward) = build_edges(&[(1, 2), (2, 3)]);
        let initial_ids = &[2];

        // Act
        let result = find_no_shuffle_graphs(initial_ids, &forward, &backward);

        // Assert
        assert!(result.is_ok());
        let graphs = result.unwrap();

        assert_eq!(graphs.len(), 1);
        let graph = &graphs[0];
        assert_eq!(graph.entries, to_hashset(&[1]));
        assert_eq!(graph.components, to_hashset(&[1, 2, 3]));
    }

    #[test]
    fn test_two_disconnected_graphs() {
        // Scenario: Two separate graphs: 1->2 and 10->11.
        // We start with one node from each graph.
        let (forward, backward) = build_edges(&[(1, 2), (10, 11)]);
        let initial_ids = &[2, 10];

        // Act
        let mut graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert_eq!(graphs.len(), 2);

        // Sort results to make the test deterministic, as HashMap iteration order is not guaranteed.
        graphs.sort_by_key(|g| *g.components.iter().min().unwrap_or(&0.into()));

        // Graph 1
        assert_eq!(graphs[0].entries, to_hashset(&[1]));
        assert_eq!(graphs[0].components, to_hashset(&[1, 2]));

        // Graph 2
        assert_eq!(graphs[1].entries, to_hashset(&[10]));
        assert_eq!(graphs[1].components, to_hashset(&[10, 11]));
    }

    #[test]
    fn test_multiple_entries_in_one_graph() {
        // Scenario: A graph with two roots feeding into one node: 1->3, 2->3.
        let (forward, backward) = build_edges(&[(1, 3), (2, 3)]);
        let initial_ids = &[3];

        // Act
        let graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert_eq!(graphs.len(), 1);
        let graph = &graphs[0];
        assert_eq!(graph.entries, to_hashset(&[1, 2]));
        assert_eq!(graph.components, to_hashset(&[1, 2, 3]));
    }

    #[test]
    fn test_diamond_shape_graph() {
        // Scenario: A diamond shape: 1->2, 1->3, 2->4, 3->4
        let (forward, backward) = build_edges(&[(1, 2), (1, 3), (2, 4), (3, 4)]);
        let initial_ids = &[4];

        // Act
        let graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert_eq!(graphs.len(), 1);
        let graph = &graphs[0];
        assert_eq!(graph.entries, to_hashset(&[1]));
        assert_eq!(graph.components, to_hashset(&[1, 2, 3, 4]));
    }

    #[test]
    fn test_starting_with_multiple_nodes_in_same_graph() {
        // Scenario: Start with two different nodes (2 and 4) from the same component.
        // Should only identify one graph, not two.
        let (forward, backward) = build_edges(&[(1, 2), (2, 3), (3, 4)]);
        let initial_ids = &[2, 4];

        // Act
        let graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert_eq!(graphs.len(), 1);
        let graph = &graphs[0];
        assert_eq!(graph.entries, to_hashset(&[1]));
        assert_eq!(graph.components, to_hashset(&[1, 2, 3, 4]));
    }

    #[test]
    fn test_empty_initial_ids() {
        // Scenario: The initial ID list is empty.
        let (forward, backward) = build_edges(&[(1, 2)]);
        let initial_ids: &[u32] = &[];

        // Act
        let graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert!(graphs.is_empty());
    }

    #[test]
    fn test_isolated_node_as_input() {
        // Scenario: Start with an ID that has no relations.
        let (forward, backward) = build_edges(&[(1, 2)]);
        let initial_ids = &[100];

        // Act
        let graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert_eq!(graphs.len(), 1);
        let graph = &graphs[0];
        assert_eq!(graph.entries, to_hashset(&[100]));
        assert_eq!(graph.components, to_hashset(&[100]));
    }

    #[test]
    fn test_graph_with_a_cycle() {
        // Scenario: A graph with a cycle: 1 -> 2 -> 3 -> 1.
        // The algorithm should correctly identify all nodes in the component.
        // Crucially, NO node is a root because every node has a parent *within the component*.
        // Therefore, the `entries` set should be empty, and the graph should not be included in the results.
        let (forward, backward) = build_edges(&[(1, 2), (2, 3), (3, 1)]);
        let initial_ids = &[2];

        // Act
        let graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert!(
            graphs.is_empty(),
            "A graph with no entries should not be returned"
        );
    }
    #[test]
    fn test_custom_complex() {
        let (forward, backward) = build_edges(&[(1, 3), (1, 8), (2, 3), (4, 3), (3, 5), (6, 7)]);
        let initial_ids = &[1, 2, 4, 6];

        // Act
        let mut graphs = find_no_shuffle_graphs(initial_ids, &forward, &backward).unwrap();

        // Assert
        assert_eq!(graphs.len(), 2);
        // Sort results to make the test deterministic, as HashMap iteration order is not guaranteed.
        graphs.sort_by_key(|g| *g.components.iter().min().unwrap_or(&0.into()));

        // Graph 1
        assert_eq!(graphs[0].entries, to_hashset(&[1, 2, 4]));
        assert_eq!(graphs[0].components, to_hashset(&[1, 2, 3, 4, 5, 8]));

        // Graph 2
        assert_eq!(graphs[1].entries, to_hashset(&[6]));
        assert_eq!(graphs[1].components, to_hashset(&[6, 7]));
    }

    #[test]
    fn render_actors_increments_actor_counter() {
        let actor_id_counter = AtomicU32::new(100);
        let fragment_id: FragmentId = 1.into();
        let job_id: JobId = 10.into();
        let database_id: DatabaseId = DatabaseId::new(3);

        let fragment_model = build_fragment(
            fragment_id,
            job_id,
            0,
            DistributionType::Single,
            1,
            StreamingParallelism::Fixed(1),
        );

        let job_model = streaming_job::Model {
            job_id,
            job_status: JobStatus::Created,
            create_type: CreateType::Foreground,
            timezone: None,
            config_override: None,
            adaptive_parallelism_strategy: None,
            parallelism: StreamingParallelism::Fixed(1),
            backfill_parallelism: None,
            backfill_orders: None,
            max_parallelism: 1,
            specific_resource_group: None,
            is_serverless_backfill: false,
        };

        let database_model = database::Model {
            database_id,
            name: "test_db".into(),
            resource_group: "rg-a".into(),
            barrier_interval_ms: None,
            checkpoint_frequency: None,
        };

        let ensembles = vec![NoShuffleEnsemble {
            entries: HashSet::from([fragment_id]),
            components: HashSet::from([fragment_id]),
        }];

        let fragment_map = HashMap::from([(fragment_id, fragment_model)]);
        let job_fragments = HashMap::from([(job_id, fragment_map)]);
        let job_map = HashMap::from([(job_id, job_model)]);

        let worker_map: HashMap<WorkerId, WorkerNode> =
            HashMap::from([(1.into(), build_worker_node(1, 1, "rg-a"))]);

        let fragment_source_ids: HashMap<FragmentId, SourceId> = HashMap::new();
        let fragment_splits: HashMap<FragmentId, Vec<SplitImpl>> = HashMap::new();
        let streaming_job_databases = HashMap::from([(job_id, database_id)]);
        let database_map = HashMap::from([(database_id, database_model)]);

        let context = RenderActorsContext {
            fragment_source_ids: &fragment_source_ids,
            fragment_splits: &fragment_splits,
            streaming_job_databases: &streaming_job_databases,
            database_map: &database_map,
        };

        let result = render_actors(
            &actor_id_counter,
            &ensembles,
            &job_fragments,
            &job_map,
            &worker_map,
            AdaptiveParallelismStrategy::Auto,
            context,
        )
        .expect("actor rendering succeeds");

        let state = collect_actor_state(&result[&database_id][&job_id][&fragment_id]);
        assert_eq!(state.len(), 1);
        assert!(
            state[0].2.is_none(),
            "single distribution should not assign vnode bitmaps"
        );
        assert_eq!(actor_id_counter.load(Ordering::Relaxed), 101);
    }

    #[test]
    fn render_actors_aligns_hash_vnode_bitmaps() {
        let actor_id_counter = AtomicU32::new(0);
        let entry_fragment_id: FragmentId = 1.into();
        let downstream_fragment_id: FragmentId = 2.into();
        let job_id: JobId = 20.into();
        let database_id: DatabaseId = DatabaseId::new(5);

        let entry_fragment = build_fragment(
            entry_fragment_id,
            job_id,
            0,
            DistributionType::Hash,
            4,
            StreamingParallelism::Fixed(2),
        );

        let downstream_fragment = build_fragment(
            downstream_fragment_id,
            job_id,
            0,
            DistributionType::Hash,
            4,
            StreamingParallelism::Fixed(2),
        );

        let job_model = streaming_job::Model {
            job_id,
            job_status: JobStatus::Created,
            create_type: CreateType::Background,
            timezone: None,
            config_override: None,
            adaptive_parallelism_strategy: None,
            parallelism: StreamingParallelism::Fixed(2),
            backfill_parallelism: None,
            backfill_orders: None,
            max_parallelism: 2,
            specific_resource_group: None,
            is_serverless_backfill: false,
        };

        let database_model = database::Model {
            database_id,
            name: "test_db_hash".into(),
            resource_group: "rg-hash".into(),
            barrier_interval_ms: None,
            checkpoint_frequency: None,
        };

        let ensembles = vec![NoShuffleEnsemble {
            entries: HashSet::from([entry_fragment_id]),
            components: HashSet::from([entry_fragment_id, downstream_fragment_id]),
        }];

        let fragment_map = HashMap::from([
            (entry_fragment_id, entry_fragment),
            (downstream_fragment_id, downstream_fragment),
        ]);
        let job_fragments = HashMap::from([(job_id, fragment_map)]);
        let job_map = HashMap::from([(job_id, job_model)]);

        let worker_map: HashMap<WorkerId, WorkerNode> = HashMap::from([
            (1.into(), build_worker_node(1, 1, "rg-hash")),
            (2.into(), build_worker_node(2, 1, "rg-hash")),
        ]);

        let fragment_source_ids: HashMap<FragmentId, SourceId> = HashMap::new();
        let fragment_splits: HashMap<FragmentId, Vec<SplitImpl>> = HashMap::new();
        let streaming_job_databases = HashMap::from([(job_id, database_id)]);
        let database_map = HashMap::from([(database_id, database_model)]);

        let context = RenderActorsContext {
            fragment_source_ids: &fragment_source_ids,
            fragment_splits: &fragment_splits,
            streaming_job_databases: &streaming_job_databases,
            database_map: &database_map,
        };

        let result = render_actors(
            &actor_id_counter,
            &ensembles,
            &job_fragments,
            &job_map,
            &worker_map,
            AdaptiveParallelismStrategy::Auto,
            context,
        )
        .expect("actor rendering succeeds");

        let entry_state = collect_actor_state(&result[&database_id][&job_id][&entry_fragment_id]);
        let downstream_state =
            collect_actor_state(&result[&database_id][&job_id][&downstream_fragment_id]);

        assert_eq!(entry_state.len(), 2);
        assert_eq!(entry_state, downstream_state);

        let assigned_vnodes: BTreeSet<_> = entry_state
            .iter()
            .flat_map(|(_, _, vnodes, _)| {
                vnodes
                    .as_ref()
                    .expect("hash distribution should populate vnode bitmap")
                    .iter()
                    .copied()
            })
            .collect();
        assert_eq!(assigned_vnodes, BTreeSet::from([0, 1, 2, 3]));
        assert_eq!(actor_id_counter.load(Ordering::Relaxed), 4);
    }

    #[test]
    fn render_actors_propagates_source_splits() {
        let actor_id_counter = AtomicU32::new(0);
        let entry_fragment_id: FragmentId = 11.into();
        let downstream_fragment_id: FragmentId = 12.into();
        let job_id: JobId = 30.into();
        let database_id: DatabaseId = DatabaseId::new(7);
        let source_id: SourceId = 99.into();

        let source_mask = FragmentTypeFlag::raw_flag([FragmentTypeFlag::Source]) as i32;
        let source_scan_mask = FragmentTypeFlag::raw_flag([FragmentTypeFlag::SourceScan]) as i32;

        let entry_fragment = build_fragment(
            entry_fragment_id,
            job_id,
            source_mask,
            DistributionType::Hash,
            4,
            StreamingParallelism::Fixed(2),
        );

        let downstream_fragment = build_fragment(
            downstream_fragment_id,
            job_id,
            source_scan_mask,
            DistributionType::Hash,
            4,
            StreamingParallelism::Fixed(2),
        );

        let job_model = streaming_job::Model {
            job_id,
            job_status: JobStatus::Created,
            create_type: CreateType::Background,
            timezone: None,
            config_override: None,
            adaptive_parallelism_strategy: None,
            parallelism: StreamingParallelism::Fixed(2),
            backfill_parallelism: None,
            backfill_orders: None,
            max_parallelism: 2,
            specific_resource_group: None,
            is_serverless_backfill: false,
        };

        let database_model = database::Model {
            database_id,
            name: "split_db".into(),
            resource_group: "rg-source".into(),
            barrier_interval_ms: None,
            checkpoint_frequency: None,
        };

        let ensembles = vec![NoShuffleEnsemble {
            entries: HashSet::from([entry_fragment_id]),
            components: HashSet::from([entry_fragment_id, downstream_fragment_id]),
        }];

        let fragment_map = HashMap::from([
            (entry_fragment_id, entry_fragment),
            (downstream_fragment_id, downstream_fragment),
        ]);
        let job_fragments = HashMap::from([(job_id, fragment_map)]);
        let job_map = HashMap::from([(job_id, job_model)]);

        let worker_map: HashMap<WorkerId, WorkerNode> = HashMap::from([
            (1.into(), build_worker_node(1, 1, "rg-source")),
            (2.into(), build_worker_node(2, 1, "rg-source")),
        ]);

        let split_a = SplitImpl::Test(TestSourceSplit {
            id: Arc::<str>::from("split-a"),
            properties: HashMap::new(),
            offset: "0".into(),
        });
        let split_b = SplitImpl::Test(TestSourceSplit {
            id: Arc::<str>::from("split-b"),
            properties: HashMap::new(),
            offset: "0".into(),
        });

        let fragment_source_ids = HashMap::from([
            (entry_fragment_id, source_id),
            (downstream_fragment_id, source_id),
        ]);
        let fragment_splits =
            HashMap::from([(entry_fragment_id, vec![split_a.clone(), split_b.clone()])]);
        let streaming_job_databases = HashMap::from([(job_id, database_id)]);
        let database_map = HashMap::from([(database_id, database_model)]);

        let context = RenderActorsContext {
            fragment_source_ids: &fragment_source_ids,
            fragment_splits: &fragment_splits,
            streaming_job_databases: &streaming_job_databases,
            database_map: &database_map,
        };

        let result = render_actors(
            &actor_id_counter,
            &ensembles,
            &job_fragments,
            &job_map,
            &worker_map,
            AdaptiveParallelismStrategy::Auto,
            context,
        )
        .expect("actor rendering succeeds");

        let entry_state = collect_actor_state(&result[&database_id][&job_id][&entry_fragment_id]);
        let downstream_state =
            collect_actor_state(&result[&database_id][&job_id][&downstream_fragment_id]);

        assert_eq!(entry_state, downstream_state);

        let split_ids: BTreeSet<_> = entry_state
            .iter()
            .flat_map(|(_, _, _, splits)| splits.iter().cloned())
            .collect();
        assert_eq!(
            split_ids,
            BTreeSet::from([split_a.id().to_string(), split_b.id().to_string()])
        );
        assert_eq!(actor_id_counter.load(Ordering::Relaxed), 4);
    }

    /// Test that job-level strategy overrides global strategy for Adaptive parallelism.
    #[test]
    fn render_actors_job_strategy_overrides_global() {
        let actor_id_counter = AtomicU32::new(0);
        let fragment_id: FragmentId = 1.into();
        let job_id: JobId = 100.into();
        let database_id: DatabaseId = DatabaseId::new(10);

        // Fragment with Adaptive parallelism, vnode_count = 8
        let fragment_model = build_fragment(
            fragment_id,
            job_id,
            0,
            DistributionType::Hash,
            8,
            StreamingParallelism::Adaptive,
        );

        // Job has custom strategy: BOUNDED(2)
        let job_model = streaming_job::Model {
            job_id,
            job_status: JobStatus::Created,
            create_type: CreateType::Foreground,
            timezone: None,
            config_override: None,
            adaptive_parallelism_strategy: Some("BOUNDED(2)".to_owned()),
            parallelism: StreamingParallelism::Adaptive,
            backfill_parallelism: None,
            backfill_orders: None,
            max_parallelism: 8,
            specific_resource_group: None,
            is_serverless_backfill: false,
        };

        let database_model = database::Model {
            database_id,
            name: "test_db".into(),
            resource_group: "default".into(),
            barrier_interval_ms: None,
            checkpoint_frequency: None,
        };

        let ensembles = vec![NoShuffleEnsemble {
            entries: HashSet::from([fragment_id]),
            components: HashSet::from([fragment_id]),
        }];

        let fragment_map =
            HashMap::from([(job_id, HashMap::from([(fragment_id, fragment_model)]))]);
        let job_map = HashMap::from([(job_id, job_model)]);

        // 4 workers with 1 parallelism each = total 4 parallelism
        let worker_map = HashMap::from([
            (1.into(), build_worker_node(1, 1, "default")),
            (2.into(), build_worker_node(2, 1, "default")),
            (3.into(), build_worker_node(3, 1, "default")),
            (4.into(), build_worker_node(4, 1, "default")),
        ]);

        let fragment_source_ids: HashMap<FragmentId, SourceId> = HashMap::new();
        let fragment_splits: HashMap<FragmentId, Vec<SplitImpl>> = HashMap::new();
        let streaming_job_databases = HashMap::from([(job_id, database_id)]);
        let database_map = HashMap::from([(database_id, database_model)]);

        let context = RenderActorsContext {
            fragment_source_ids: &fragment_source_ids,
            fragment_splits: &fragment_splits,
            streaming_job_databases: &streaming_job_databases,
            database_map: &database_map,
        };

        // Global strategy is FULL (would give 4 actors), but job strategy is BOUNDED(2)
        let result = render_actors(
            &actor_id_counter,
            &ensembles,
            &fragment_map,
            &job_map,
            &worker_map,
            AdaptiveParallelismStrategy::Full,
            context,
        )
        .expect("actor rendering succeeds");

        let state = collect_actor_state(&result[&database_id][&job_id][&fragment_id]);
        // Job strategy BOUNDED(2) should limit to 2 actors, not 4 (global FULL)
        assert_eq!(
            state.len(),
            2,
            "Job strategy BOUNDED(2) should override global FULL"
        );
    }

    /// Test that global strategy is used when job has no custom strategy.
    #[test]
    fn render_actors_uses_global_strategy_when_job_has_none() {
        let actor_id_counter = AtomicU32::new(0);
        let fragment_id: FragmentId = 1.into();
        let job_id: JobId = 101.into();
        let database_id: DatabaseId = DatabaseId::new(11);

        let fragment_model = build_fragment(
            fragment_id,
            job_id,
            0,
            DistributionType::Hash,
            8,
            StreamingParallelism::Adaptive,
        );

        // Job has NO custom strategy (None)
        let job_model = streaming_job::Model {
            job_id,
            job_status: JobStatus::Created,
            create_type: CreateType::Foreground,
            timezone: None,
            config_override: None,
            adaptive_parallelism_strategy: None, // No custom strategy
            parallelism: StreamingParallelism::Adaptive,
            backfill_parallelism: None,
            backfill_orders: None,
            max_parallelism: 8,
            specific_resource_group: None,
            is_serverless_backfill: false,
        };

        let database_model = database::Model {
            database_id,
            name: "test_db".into(),
            resource_group: "default".into(),
            barrier_interval_ms: None,
            checkpoint_frequency: None,
        };

        let ensembles = vec![NoShuffleEnsemble {
            entries: HashSet::from([fragment_id]),
            components: HashSet::from([fragment_id]),
        }];

        let fragment_map =
            HashMap::from([(job_id, HashMap::from([(fragment_id, fragment_model)]))]);
        let job_map = HashMap::from([(job_id, job_model)]);

        // 4 workers = total 4 parallelism
        let worker_map = HashMap::from([
            (1.into(), build_worker_node(1, 1, "default")),
            (2.into(), build_worker_node(2, 1, "default")),
            (3.into(), build_worker_node(3, 1, "default")),
            (4.into(), build_worker_node(4, 1, "default")),
        ]);

        let fragment_source_ids: HashMap<FragmentId, SourceId> = HashMap::new();
        let fragment_splits: HashMap<FragmentId, Vec<SplitImpl>> = HashMap::new();
        let streaming_job_databases = HashMap::from([(job_id, database_id)]);
        let database_map = HashMap::from([(database_id, database_model)]);

        let context = RenderActorsContext {
            fragment_source_ids: &fragment_source_ids,
            fragment_splits: &fragment_splits,
            streaming_job_databases: &streaming_job_databases,
            database_map: &database_map,
        };

        // Global strategy is BOUNDED(3)
        let result = render_actors(
            &actor_id_counter,
            &ensembles,
            &fragment_map,
            &job_map,
            &worker_map,
            AdaptiveParallelismStrategy::Bounded(NonZeroUsize::new(3).unwrap()),
            context,
        )
        .expect("actor rendering succeeds");

        let state = collect_actor_state(&result[&database_id][&job_id][&fragment_id]);
        // Should use global strategy BOUNDED(3)
        assert_eq!(
            state.len(),
            3,
            "Should use global strategy BOUNDED(3) when job has no custom strategy"
        );
    }

    /// Test that Fixed parallelism ignores strategy entirely.
    #[test]
    fn render_actors_fixed_parallelism_ignores_strategy() {
        let actor_id_counter = AtomicU32::new(0);
        let fragment_id: FragmentId = 1.into();
        let job_id: JobId = 102.into();
        let database_id: DatabaseId = DatabaseId::new(12);

        // Fragment with FIXED parallelism
        let fragment_model = build_fragment(
            fragment_id,
            job_id,
            0,
            DistributionType::Hash,
            8,
            StreamingParallelism::Fixed(5),
        );

        // Job has custom strategy, but it should be ignored for Fixed parallelism
        let job_model = streaming_job::Model {
            job_id,
            job_status: JobStatus::Created,
            create_type: CreateType::Foreground,
            timezone: None,
            config_override: None,
            adaptive_parallelism_strategy: Some("BOUNDED(2)".to_owned()),
            parallelism: StreamingParallelism::Fixed(5),
            backfill_parallelism: None,
            backfill_orders: None,
            max_parallelism: 8,
            specific_resource_group: None,
            is_serverless_backfill: false,
        };

        let database_model = database::Model {
            database_id,
            name: "test_db".into(),
            resource_group: "default".into(),
            barrier_interval_ms: None,
            checkpoint_frequency: None,
        };

        let ensembles = vec![NoShuffleEnsemble {
            entries: HashSet::from([fragment_id]),
            components: HashSet::from([fragment_id]),
        }];

        let fragment_map =
            HashMap::from([(job_id, HashMap::from([(fragment_id, fragment_model)]))]);
        let job_map = HashMap::from([(job_id, job_model)]);

        // 6 workers = total 6 parallelism
        let worker_map = HashMap::from([
            (1.into(), build_worker_node(1, 1, "default")),
            (2.into(), build_worker_node(2, 1, "default")),
            (3.into(), build_worker_node(3, 1, "default")),
            (4.into(), build_worker_node(4, 1, "default")),
            (5.into(), build_worker_node(5, 1, "default")),
            (6.into(), build_worker_node(6, 1, "default")),
        ]);

        let fragment_source_ids: HashMap<FragmentId, SourceId> = HashMap::new();
        let fragment_splits: HashMap<FragmentId, Vec<SplitImpl>> = HashMap::new();
        let streaming_job_databases = HashMap::from([(job_id, database_id)]);
        let database_map = HashMap::from([(database_id, database_model)]);

        let context = RenderActorsContext {
            fragment_source_ids: &fragment_source_ids,
            fragment_splits: &fragment_splits,
            streaming_job_databases: &streaming_job_databases,
            database_map: &database_map,
        };

        let result = render_actors(
            &actor_id_counter,
            &ensembles,
            &fragment_map,
            &job_map,
            &worker_map,
            AdaptiveParallelismStrategy::Full,
            context,
        )
        .expect("actor rendering succeeds");

        let state = collect_actor_state(&result[&database_id][&job_id][&fragment_id]);
        // Fixed(5) should be used, ignoring both job strategy BOUNDED(2) and global FULL
        assert_eq!(
            state.len(),
            5,
            "Fixed parallelism should ignore all strategies"
        );
    }

    /// Test RATIO strategy calculation.
    #[test]
    fn render_actors_ratio_strategy() {
        let actor_id_counter = AtomicU32::new(0);
        let fragment_id: FragmentId = 1.into();
        let job_id: JobId = 103.into();
        let database_id: DatabaseId = DatabaseId::new(13);

        let fragment_model = build_fragment(
            fragment_id,
            job_id,
            0,
            DistributionType::Hash,
            16,
            StreamingParallelism::Adaptive,
        );

        // Job has RATIO(0.5) strategy
        let job_model = streaming_job::Model {
            job_id,
            job_status: JobStatus::Created,
            create_type: CreateType::Foreground,
            timezone: None,
            config_override: None,
            adaptive_parallelism_strategy: Some("RATIO(0.5)".to_owned()),
            parallelism: StreamingParallelism::Adaptive,
            backfill_parallelism: None,
            backfill_orders: None,
            max_parallelism: 16,
            specific_resource_group: None,
            is_serverless_backfill: false,
        };

        let database_model = database::Model {
            database_id,
            name: "test_db".into(),
            resource_group: "default".into(),
            barrier_interval_ms: None,
            checkpoint_frequency: None,
        };

        let ensembles = vec![NoShuffleEnsemble {
            entries: HashSet::from([fragment_id]),
            components: HashSet::from([fragment_id]),
        }];

        let fragment_map =
            HashMap::from([(job_id, HashMap::from([(fragment_id, fragment_model)]))]);
        let job_map = HashMap::from([(job_id, job_model)]);

        // 8 workers = total 8 parallelism
        let worker_map = HashMap::from([
            (1.into(), build_worker_node(1, 1, "default")),
            (2.into(), build_worker_node(2, 1, "default")),
            (3.into(), build_worker_node(3, 1, "default")),
            (4.into(), build_worker_node(4, 1, "default")),
            (5.into(), build_worker_node(5, 1, "default")),
            (6.into(), build_worker_node(6, 1, "default")),
            (7.into(), build_worker_node(7, 1, "default")),
            (8.into(), build_worker_node(8, 1, "default")),
        ]);

        let fragment_source_ids: HashMap<FragmentId, SourceId> = HashMap::new();
        let fragment_splits: HashMap<FragmentId, Vec<SplitImpl>> = HashMap::new();
        let streaming_job_databases = HashMap::from([(job_id, database_id)]);
        let database_map = HashMap::from([(database_id, database_model)]);

        let context = RenderActorsContext {
            fragment_source_ids: &fragment_source_ids,
            fragment_splits: &fragment_splits,
            streaming_job_databases: &streaming_job_databases,
            database_map: &database_map,
        };

        let result = render_actors(
            &actor_id_counter,
            &ensembles,
            &fragment_map,
            &job_map,
            &worker_map,
            AdaptiveParallelismStrategy::Full,
            context,
        )
        .expect("actor rendering succeeds");

        let state = collect_actor_state(&result[&database_id][&job_id][&fragment_id]);
        // RATIO(0.5) of 8 = 4
        assert_eq!(
            state.len(),
            4,
            "RATIO(0.5) of 8 workers should give 4 actors"
        );
    }
}