Repository: risingwavelabs/risingwave
File: src/meta/src/stream/scale.rs (Rust)

Code before:
                .then(id_a.cmp(id_b))
        };

    let builder_from_bitmap = |bitmap: &Bitmap| -> BitmapBuilder {
        let mut builder = BitmapBuilder::default();
        builder.append_bitmap(bitmap);
        builder
    };

    let (prev_expected, _) = vnode_count.div_rem(&actors.len());

    // Guard rail: if any existing actor (excluding the ones to remove) holds fewer than
    // `prev_expected` vnodes, the historical distribution is already invalid. Instead of
    // panicking, rebuild a fresh even distribution across the surviving + new actors so that
    // meta can proceed.
    let has_under_allocated = actors
        .iter()
        .filter(|actor| !actors_to_remove.contains(&actor.actor_id))
        .any(|actor| {
            let count = actor
                .vnode_bitmap
                .as_ref()
                .expect("vnode bitmap unset")
                .count_ones();
            count < prev_expected
        });

    if has_under_allocated {
        tracing::warn!(
            ?prev_expected,
            actors = actors.len(),
            actors_to_remove = actors_to_remove.len(),
            actors_to_create = actors_to_create.len(),
            "rebalance_actor_vnode detected under-allocated actors; rebuild distribution evenly"
        );

        let mut future_actors: Vec<_> = actors
            .iter()
            .filter(|actor| !actors_to_remove.contains(&actor.actor_id))
            .map(|actor| actor.actor_id as ActorId)
            .collect();
        future_actors.extend(actors_to_create.iter().copied());
        future_actors.sort_unstable();
        future_actors.dedup();

        assert!(
            !future_actors.is_empty(),
            "should have at least one actor after rebuild"
        );

        let mut builders: HashMap<ActorId, BitmapBuilder> = future_actors

Review comment (bug):
The guard rail only checks surviving actors for under-allocation, but the assertion at line 286 (below) still requires removed actors to have at least `prev_expected` vnodes. If an actor being removed has fewer vnodes than expected, the assertion will fail. Consider also checking actors in `actors_to_remove` for under-allocation to ensure complete recovery from invalid vnode distributions.

Code after:
                .then(id_a.cmp(id_b))
        };

    let builder_from_bitmap = |bitmap: &Bitmap| -> BitmapBuilder {
        let mut builder = BitmapBuilder::default();
        builder.append_bitmap(bitmap);
        builder
    };

    let (prev_expected, _) = vnode_count.div_rem(&actors.len());

    // Guard rail: if any actor (even one to be removed) holds fewer than `prev_expected` vnodes,
    // the historical distribution is already invalid. Instead of panicking later, rebuild a fresh
    // even distribution across the surviving + new actors so that meta can proceed.
    let has_under_allocated = actors.iter().any(|actor| {
        let count = actor
            .vnode_bitmap
            .as_ref()
            .expect("vnode bitmap unset")
            .count_ones();
        count < prev_expected
    });

    if has_under_allocated {
        tracing::warn!(
            ?prev_expected,
            actors = actors.len(),
            actors_to_remove = actors_to_remove.len(),
            actors_to_create = actors_to_create.len(),
            "rebalance_actor_vnode detected under-allocated actors; rebuild distribution evenly"
        );

        let mut future_actors: Vec<_> = actors
            .iter()
            .filter(|actor| !actors_to_remove.contains(&actor.actor_id))
            .map(|actor| actor.actor_id as ActorId)
            .collect();
        future_actors.extend(actors_to_create.iter().copied());
        future_actors.sort_unstable();
        future_actors.dedup();

        assert!(
            !future_actors.is_empty(),
            "should have at least one actor after rebuild"
        );

        let mut builders: HashMap<ActorId, BitmapBuilder> = future_actors
            .iter()
            .map(|id| (*id, BitmapBuilder::zeroed(vnode_count)))
            .collect();