16728097825

Committed 04 Aug 2025 04:00PM UTC coverage: 48.355% (-35.1%) from 83.429%

Build # 16728097825

Build Type

Pull #4108

github

Committed by

web-flow

Commit Message

Merge 1b2d27fd8 into 649ba9576

Pull Request Pull Request #4108: perf[vortex-array]: use all_valid instead of `invalid_count() == 0`

Coverage Stats

1 of 1 new or added line in 1 file covered. (100.0%)

11596 existing lines in 378 files now uncovered.

18635 of 38538 relevant lines covered (48.35%)

151786.4 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.22

/vortex-scan/src/work_queue.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

//! A work-stealing iterator that supports dynamically adding tasks from task factories.

use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering::{Relaxed, SeqCst};
use std::{iter, thread};

use crossbeam_deque::{Steal, Stealer, Worker};
use crossbeam_queue::SegQueue;
use parking_lot::RwLock;
use vortex_error::VortexResult;

/// A factory that produces a vector of tasks.
pub type TaskFactory<T> = Box<dyn FnOnce() -> VortexResult<Vec<T>> + Send>;

/// A work-stealing queue that allows for dynamic task addition.
pub struct WorkStealingQueue<T> {
    state: Arc<State<T>>,
}

impl<T> Clone for WorkStealingQueue<T> {
    fn clone(&self) -> Self {
        Self {
            state: self.state.clone(),
        }
    }
}

impl<T> WorkStealingQueue<T> {
    /// Created with lazily constructed task factories.
    pub fn new<I>(factories: I) -> Self
    where
        I: IntoIterator<Item = TaskFactory<T>>,
    {
        let queue: SegQueue<TaskFactory<T>> = SegQueue::new();
        for factory in factories.into_iter() {
            queue.push(factory);
        }
        let num_factories = queue.len();

        Self {
            state: Arc::new(State {
                task_factories: queue,
                num_factories_constructed: AtomicUsize::new(0),
                num_factories,
                stealers: RwLock::new(Vec::new()),
                stealer_offset: Default::default(),
            }),
        }
    }

    pub fn new_iterator(self) -> WorkStealingIterator<T> {
        self.state.new_iterator()
    }
}

/// Shared state for the work queue.
struct State<T> {
    /// A queue of factories that lazily produce tasks of type `T`.
    task_factories: SegQueue<TaskFactory<T>>,

    /// The total number of task factories that need to be constructed.
    num_factories: usize,

    /// How many factories have been constructed and had their tasks completely pushed into
    /// a worker queue.
    num_factories_constructed: AtomicUsize,

    /// The vector of stealers, one for each worker.
    stealers: RwLock<Vec<Stealer<T>>>,

    /// An offset into the stealers vector, used to avoid skewed worker queues when stealing.
    stealer_offset: AtomicUsize,
}

impl<T> State<T> {
    /// Create a new iterator.
    fn new_iterator(self: Arc<Self>) -> WorkStealingIterator<T> {
        let worker = Worker::new_fifo();

        // Register the new worker with the shared state.
        self.stealers.write().push(worker.stealer());

        WorkStealingIterator {
            state: self,
            worker,
        }
    }

    /// Loads a factory and pushes its tasks into the given worker queue.
    ///
    /// Returns `true` if any tasks were pushed into the worker. Note that these tasks may have
    /// been stolen by the time the worker queue is checked.
    fn load_next_factory(&self, worker: &Worker<T>) -> VortexResult<bool> {
        loop {
            if let Some(factory_fn) = self.task_factories.pop() {
                let tasks = factory_fn().inspect_err(|_| {
                    // In case of an error, increment the counter such that all other workers are able to terminate.
                    // `num_factories_constructed` is part of the loop condition when workers attempt to steal work.
                    self.num_factories_constructed.fetch_add(1, SeqCst);
                })?;
                let is_empty = tasks.is_empty();

                // Tasks *must* be pushed before `num_factories_constructed` is incremented.
                for task in tasks {
                    worker.push(task);
                }

                self.num_factories_constructed.fetch_add(1, SeqCst);

                // Keep looping until we find a factory that has pushed tasks.
                if !is_empty {
                    return Ok(true);
                }
            } else {
                return Ok(false);
            }
        }
    }

    /// Reports whether there is any work left to steal.
    fn stealers_have_work(&self) -> bool {
        self.stealers
            .read()
            .iter()
            .any(|stealer| !stealer.is_empty())
    }

    /// Attempts to steal work from other workers, returns `true` if work was stolen.
    fn steal_work(&self, worker: &Worker<T>) -> Steal<()> {
        // Repeatedly attempt to steal work from other workers until there are no retries.
        iter::repeat_with(|| {
            // This collect tries all stealers, exits early on the first successful steal,
            // or else tracks whether any steal requires a retry.
            let guard = self.stealers.read();
            let num_stealers = guard.len();
            guard
                .iter()
                .cycle()
                .skip(self.stealer_offset.fetch_add(1, Relaxed) % num_stealers)
                .take(num_stealers)
                .map(|stealer| stealer.steal_batch(worker))
                .collect::<Steal<()>>()
        })
        .find(|steal| !steal.is_retry())
        .unwrap_or(Steal::Empty)
    }
}

/// A work-stealing iterator that supports dynamically adding tasks from task factories.
///
/// Each task factory has affinity to a particular worker. After all factories have been
/// constructed, workers will attempt to steal tasks from each other until all tasks are processed.
///
/// Workers are constructed by cloning the iterator.
pub struct WorkStealingIterator<T> {
    state: Arc<State<T>>,
    worker: Worker<T>,
}

impl<T> Clone for WorkStealingIterator<T> {
    fn clone(&self) -> Self {
        self.state.clone().new_iterator()
    }
}

impl<T> Iterator for WorkStealingIterator<T> {
    type Item = VortexResult<T>;

    fn next(&mut self) -> Option<VortexResult<T>> {
        if self.worker.is_empty() {
            let next_factory_loaded = match self.state.load_next_factory(&self.worker) {
                Ok(next_factory_loaded) => next_factory_loaded,
                Err(e) => return Some(Err(e)),
            };

            if !next_factory_loaded {
                // If there are no more factories to load, then there is at least one worker
                // constructing a factory and about to push some tasks.
                //
                // We sit in a loop trying to steal some of those tasks, or else bail out when
                // all scans have been constructed, and we didn't manage to steal anything. To avoid
                // spinning too hot, we yield the thread each time we fail to steal work.
                //
                // `steal_work` does have the side effect of stealing work, and we only want to loop
                // again if the result of an attempt of stealing results with `Retry`, for other cases
                // `Empty` and `Success` there is no point in trying again
                while self.state.num_factories_constructed.load(Relaxed) < self.state.num_factories
                    || self.state.stealers_have_work()
                {
                    if self.state.steal_work(&self.worker).is_success() {
                        break;
                    } else {
                        thread::yield_now();
                    }
                }
            }
        }

        // Attempt to pop a task from the worker queue.
        // Another worker may have stolen our tasks by this point. If that's the case, then we've
        // already finished loading the factories, and we're down to the last few tasks. Therefore,
        // it's ok for us to return `None` and terminate the iterator.
        Some(Ok(self.worker.pop()?))
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	//! A work-stealing iterator that supports dynamically adding tasks from task factories.
5
6	use std::sync::Arc;
7	use std::sync::atomic::AtomicUsize;
8	use std::sync::atomic::Ordering::{Relaxed, SeqCst};
9	use std::{iter, thread};
10
11	use crossbeam_deque::{Steal, Stealer, Worker};
12	use crossbeam_queue::SegQueue;
13	use parking_lot::RwLock;
14	use vortex_error::VortexResult;
15
16	/// A factory that produces a vector of tasks.
17	pub type TaskFactory<T> = Box<dyn FnOnce() -> VortexResult<Vec<T>> + Send>;
18
19	/// A work-stealing queue that allows for dynamic task addition.
20	pub struct WorkStealingQueue<T> {
21	state: Arc<State<T>>,
22	}
23
24	impl<T> Clone for WorkStealingQueue<T> {
25	fn clone(&self) -> Self {	1,376✔
26	Self {	1,376✔
27	state: self.state.clone(),	1,376✔
28	}	1,376✔
29	}	1,376✔
30	}
31
32	impl<T> WorkStealingQueue<T> {
33	/// Created with lazily constructed task factories.
34	pub fn new<I>(factories: I) -> Self	172✔
35	where	172✔
36	I: IntoIterator<Item = TaskFactory<T>>,	172✔
37	{
38	let queue: SegQueue<TaskFactory<T>> = SegQueue::new();	172✔
39	for factory in factories.into_iter() {	172✔
40	queue.push(factory);	172✔
41	}	172✔
42	let num_factories = queue.len();	172✔
43
44	Self {	172✔
45	state: Arc::new(State {	172✔
46	task_factories: queue,	172✔
47	num_factories_constructed: AtomicUsize::new(0),	172✔
48	num_factories,	172✔
49	stealers: RwLock::new(Vec::new()),	172✔
50	stealer_offset: Default::default(),	172✔
51	}),	172✔
52	}	172✔
53	}	172✔
54
55	pub fn new_iterator(self) -> WorkStealingIterator<T> {	1,376✔
56	self.state.new_iterator()	1,376✔
57	}	1,376✔
58	}
59
60	/// Shared state for the work queue.
61	struct State<T> {
62	/// A queue of factories that lazily produce tasks of type `T`.
63	task_factories: SegQueue<TaskFactory<T>>,
64
65	/// The total number of task factories that need to be constructed.
66	num_factories: usize,
67
68	/// How many factories have been constructed and had their tasks completely pushed into
69	/// a worker queue.
70	num_factories_constructed: AtomicUsize,
71
72	/// The vector of stealers, one for each worker.
73	stealers: RwLock<Vec<Stealer<T>>>,
74
75	/// An offset into the stealers vector, used to avoid skewed worker queues when stealing.
76	stealer_offset: AtomicUsize,
77	}
78
79	impl<T> State<T> {
80	/// Create a new iterator.
81	fn new_iterator(self: Arc<Self>) -> WorkStealingIterator<T> {	1,376✔
82	let worker = Worker::new_fifo();	1,376✔
83
84	// Register the new worker with the shared state.
85	self.stealers.write().push(worker.stealer());	1,376✔
86
87	WorkStealingIterator {	1,376✔
88	state: self,	1,376✔
89	worker,	1,376✔
90	}	1,376✔
91	}	1,376✔
92
93	/// Loads a factory and pushes its tasks into the given worker queue.
94	///
95	/// Returns `true` if any tasks were pushed into the worker. Note that these tasks may have
96	/// been stolen by the time the worker queue is checked.
97	fn load_next_factory(&self, worker: &Worker<T>) -> VortexResult<bool> {	1,926✔
98	loop {
99	if let Some(factory_fn) = self.task_factories.pop() {	1,926✔
100	let tasks = factory_fn().inspect_err(\|_\| {	172✔
101	// In case of an error, increment the counter such that all other workers are able to terminate.
102	// `num_factories_constructed` is part of the loop condition when workers attempt to steal work.
103	self.num_factories_constructed.fetch_add(1, SeqCst);	×
104	})?;	×
105	let is_empty = tasks.is_empty();	172✔
106
107	// Tasks must be pushed before `num_factories_constructed` is incremented.
108	for task in tasks {	748✔
109	worker.push(task);	576✔
110	}	576✔
111
112	self.num_factories_constructed.fetch_add(1, SeqCst);	172✔
113
114	// Keep looping until we find a factory that has pushed tasks.
115	if !is_empty {	172✔
116	return Ok(true);	172✔
UNCOV 117	}	×
118	} else {
119	return Ok(false);	1,754✔
120	}
121	}
122	}	1,926✔
123
124	/// Reports whether there is any work left to steal.
125	fn stealers_have_work(&self) -> bool {	1,608✔
126	self.stealers	1,608✔
127	.read()	1,608✔
128	.iter()	1,608✔
129	.any(\|stealer\| !stealer.is_empty())	10,996✔
130	}	1,608✔
131
132	/// Attempts to steal work from other workers, returns `true` if work was stolen.
133	fn steal_work(&self, worker: &Worker<T>) -> Steal<()> {	206,768✔
134	// Repeatedly attempt to steal work from other workers until there are no retries.
135	iter::repeat_with(\|\| {	207,206✔
136	// This collect tries all stealers, exits early on the first successful steal,
137	// or else tracks whether any steal requires a retry.
138	let guard = self.stealers.read();	207,206✔
139	let num_stealers = guard.len();	207,206✔
140	guard	207,206✔
141	.iter()	207,206✔
142	.cycle()	207,206✔
143	.skip(self.stealer_offset.fetch_add(1, Relaxed) % num_stealers)	207,206✔
144	.take(num_stealers)	207,206✔
145	.map(\|stealer\| stealer.steal_batch(worker))	1,527,590✔
146	.collect::<Steal<()>>()	207,206✔
147	})	207,206✔
148	.find(\|steal\| !steal.is_retry())	207,206✔
149	.unwrap_or(Steal::Empty)	206,768✔
150	}	206,768✔
151	}
152
153	/// A work-stealing iterator that supports dynamically adding tasks from task factories.
154	///
155	/// Each task factory has affinity to a particular worker. After all factories have been
156	/// constructed, workers will attempt to steal tasks from each other until all tasks are processed.
157	///
158	/// Workers are constructed by cloning the iterator.
159	pub struct WorkStealingIterator<T> {
160	state: Arc<State<T>>,
161	worker: Worker<T>,
162	}
163
164	impl<T> Clone for WorkStealingIterator<T> {
165	fn clone(&self) -> Self {	×
166	self.state.clone().new_iterator()	×
167	}	×
168	}
169
170	impl<T> Iterator for WorkStealingIterator<T> {
171	type Item = VortexResult<T>;
172
173	fn next(&mut self) -> Option<VortexResult<T>> {	1,952✔
174	if self.worker.is_empty() {	1,952✔
175	let next_factory_loaded = match self.state.load_next_factory(&self.worker) {	1,926✔
176	Ok(next_factory_loaded) => next_factory_loaded,	1,926✔
177	Err(e) => return Some(Err(e)),	×
178	};
179
180	if !next_factory_loaded {	1,926✔
181	// If there are no more factories to load, then there is at least one worker
182	// constructing a factory and about to push some tasks.
183	//
184	// We sit in a loop trying to steal some of those tasks, or else bail out when
185	// all scans have been constructed, and we didn't manage to steal anything. To avoid
186	// spinning too hot, we yield the thread each time we fail to steal work.
187	//
188	// `steal_work` does have the side effect of stealing work, and we only want to loop
189	// again if the result of an attempt of stealing results with `Retry`, for other cases
190	// `Empty` and `Success` there is no point in trying again
191	while self.state.num_factories_constructed.load(Relaxed) < self.state.num_factories	208,140✔
192	\|\| self.state.stealers_have_work()	1,608✔
193	{
194	if self.state.steal_work(&self.worker).is_success() {	206,768✔
195	break;	382✔
196	} else {	206,386✔
197	thread::yield_now();	206,386✔
198	}	206,386✔
199	}
200	}	172✔
201	}	26✔
202
203	// Attempt to pop a task from the worker queue.
204	// Another worker may have stolen our tasks by this point. If that's the case, then we've
205	// already finished loading the factories, and we're down to the last few tasks. Therefore,
206	// it's ok for us to return `None` and terminate the iterator.
207	Some(Ok(self.worker.pop()?))	1,952✔
208	}	1,952✔
209	}

vortex-data / vortex / 16728097825

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous