Repository: risingwavelabs/risingwave
File: src/storage/src/hummock/compactor/iceberg_compaction/mod.rs (Rust)

Code before:
        if meta.required_parallelism > self.max_parallelism {
            return PushResult::RejectedTooLarge;
        }

        let new_total = self.inner.waiting_parallelism_sum + meta.required_parallelism;
        if new_total > self.pending_parallelism_budget {
            return PushResult::RejectedCapacity;
        }

        let key = meta.key();
        self.inner.id_map.insert(key, meta.required_parallelism);
        self.inner.waiting_parallelism_sum = new_total;

        if let Some(r) = runner {
            self.inner.runners.insert(key, r);
        }

        self.inner.deque.push_back(meta);

        self.notify_schedulable();
        PushResult::Added
    }

    /// Pop the next task if sufficient parallelism is available.
    ///
    /// Returns `None` if the queue is empty or the front task cannot fit
    /// within the available parallelism budget.
    pub fn pop(&mut self) -> Option<PoppedIcebergTask> {
        let front = self.inner.deque.front()?;
        if front.required_parallelism > self.available_parallelism() {
            return None;
        }

        let meta = self.inner.deque.pop_front()?;
        self.inner.waiting_parallelism_sum = self
            .inner
            .waiting_parallelism_sum
            .saturating_sub(meta.required_parallelism);
        self.inner.running_parallelism_sum = self
            .inner
            .running_parallelism_sum
            .saturating_add(meta.required_parallelism);

        let runner = self.inner.runners.remove(&meta.key());
        Some(PoppedIcebergTask { meta, runner })
    }

    /// Mark a task as finished, freeing its parallelism for other tasks.
    ///
    /// Returns `true` if the task was found and removed, `false` otherwise.
    pub fn finish_running(&mut self, task_key: TaskKey) -> bool {

Review comment (bug):
The metadata is pushed to the deque after inserting into the id_map and runners. If the push to deque succeeds but a later operation fails, the metadata will be in the maps but not the deque. Consider moving this line before the runner insertion (line 185-187) to maintain consistency, or ensure the deque push happens before any map updates that could potentially fail.

Code after:
            return PushResult::RejectedInvalidParallelism;
        }
        if meta.required_parallelism > self.max_parallelism {
            return PushResult::RejectedTooLarge;
        }

        let key = meta.key();

        // Reject duplicate keys to prevent inconsistent state between id_map and deque
        if self.inner.id_map.contains_key(&key) {
            return PushResult::RejectedDuplicate;
        }

        let new_total = self.inner.waiting_parallelism_sum + meta.required_parallelism;
        if new_total > self.pending_parallelism_budget {
            return PushResult::RejectedCapacity;
        }

        self.inner.id_map.insert(key, meta.required_parallelism);
        self.inner.waiting_parallelism_sum = new_total;

        if let Some(r) = runner {
            self.inner.runners.insert(key, r);
        }

        self.inner.deque.push_back(meta);

        self.notify_schedulable();
        PushResult::Added
    }

    /// Pop the next task if sufficient parallelism is available.
    ///
    /// Returns `None` if the queue is empty or the front task cannot fit
    /// within the available parallelism budget.
    pub fn pop(&mut self) -> Option<PoppedIcebergTask> {
        let front = self.inner.deque.front()?;
        if front.required_parallelism > self.available_parallelism() {
            return None;
        }

        let meta = self.inner.deque.pop_front()?;
        self.inner.waiting_parallelism_sum = self
            .inner
            .waiting_parallelism_sum
            .saturating_sub(meta.required_parallelism);
        self.inner.running_parallelism_sum = self
            .inner
            .running_parallelism_sum
            .saturating_add(meta.required_parallelism);