606 lines
23 KiB
Rust
606 lines
23 KiB
Rust
use std::any::Any;
|
|
use std::collections::{HashMap, HashSet};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
use std::time::Duration;
|
|
|
|
use async_channel::Sender;
|
|
use async_trait::async_trait;
|
|
use futures_util::stream::AbortHandle;
|
|
use itertools::Itertools;
|
|
use warp_core::r#async::debounce;
|
|
use warp_core::send_telemetry_from_ctx;
|
|
use warpui_core::r#async::Timer;
|
|
use warpui_core::{Action, AppContext, Entity, ModelContext};
|
|
|
|
use super::data_source::{Query, QueryFilter, QueryResult};
|
|
use crate::telemetry::TelemetryEvent;
|
|
|
|
/// Maximum time to wait for matching data sources to return results before showing
|
|
/// partial results.
|
|
///
|
|
/// This is a UX tradeoff: waiting briefly reduces flicker in UIs that mix sync and async
|
|
/// sources (e.g. command palette file search), but we still want to show something quickly
|
|
/// if an async source is slow.
|
|
const INITIAL_RESULTS_TIMEOUT: Duration = Duration::from_millis(500);
|
|
|
|
pub use warpui_core::r#async::BoxFuture;
|
|
|
|
/// A structure that combines results from various data sources to produce a
|
|
/// single, ordered, heterogeneous list of search results.
|
|
#[derive(Default)]
|
|
pub struct SearchMixer<T: Action + Clone> {
|
|
/// The set of sources to be used to run a query against.
|
|
sources: HashMap<DataSourceId, RegisteredDataSource<T>>,
|
|
|
|
/// The latest set of search results produced by the latest `query`.
|
|
results: Vec<QueryResult<T>>,
|
|
|
|
/// The latest query that was used to search against, if any.
|
|
query: Option<Query>,
|
|
|
|
/// The set of sources that have finished running for the latest query.
|
|
finished_sources: HashSet<DataSourceId>,
|
|
|
|
/// Monotonically increasing counter incremented on each `run_query`. Used to discard stale
|
|
/// async callbacks and timeout callbacks whose futures completed before the abort took effect.
|
|
query_generation: u64,
|
|
|
|
/// Results buffered for the current query that haven't been committed to results yet.
|
|
/// `Some(vec)` means we're actively buffering (old results remain visible).
|
|
/// `None` means results have been committed; late-arriving results go directly to `results`.
|
|
pending_results: Option<Vec<QueryResult<T>>>,
|
|
|
|
/// Tracks whether the current query has emitted its initial set of visible results yet.
|
|
initial_results_emitted: bool,
|
|
}
|
|
|
|
impl<T: Action + Clone> Entity for SearchMixer<T> {
|
|
type Event = SearchMixerEvent;
|
|
}
|
|
|
|
/// A unique identifier for a DataSource.
|
|
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
|
|
struct DataSourceId(usize);
|
|
impl DataSourceId {
|
|
/// Constructs a new globally-unique entity ID.
|
|
#[allow(clippy::new_without_default)]
|
|
pub fn new() -> DataSourceId {
|
|
static NEXT_ID: AtomicUsize = AtomicUsize::new(0);
|
|
let raw = NEXT_ID.fetch_add(1, Ordering::Relaxed);
|
|
DataSourceId(raw)
|
|
}
|
|
}
|
|
|
|
pub enum SearchMixerEvent {
|
|
ResultsChanged,
|
|
}
|
|
|
|
pub struct AddAsyncSourceOptions {
|
|
pub debounce_interval: Option<Duration>,
|
|
/// Whether to run this source when the query text is empty
|
|
/// (i.e. the user hasn't typed anything yet).
|
|
pub run_in_zero_state: bool,
|
|
pub run_when_unfiltered: bool,
|
|
}
|
|
|
|
impl<T: Action + Clone> SearchMixer<T> {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
sources: HashMap::new(),
|
|
finished_sources: HashSet::new(),
|
|
results: vec![],
|
|
query: None,
|
|
query_generation: 0,
|
|
pending_results: None,
|
|
initial_results_emitted: false,
|
|
}
|
|
}
|
|
|
|
/// Resets the mixer's state.
|
|
pub fn reset(&mut self, ctx: &mut ModelContext<Self>) {
|
|
self.abort_in_flight_async_queries();
|
|
self.query_generation = self.query_generation.wrapping_add(1);
|
|
|
|
self.sources.clear();
|
|
self.finished_sources.clear();
|
|
self.results.clear();
|
|
self.pending_results = None;
|
|
self.query.take();
|
|
self.initial_results_emitted = false;
|
|
ctx.emit(SearchMixerEvent::ResultsChanged);
|
|
}
|
|
|
|
/// Abort the current in-flight query to avoid stale searches
|
|
/// continuing and passing back results when they are no longer wanted.
|
|
fn abort_in_flight_async_queries(&mut self) {
|
|
for registered_source in self.sources.values_mut() {
|
|
if let DataSource::AsyncDataSource {
|
|
latest_run_abort_handle,
|
|
..
|
|
} = &mut registered_source.source
|
|
&& let Some(abort_handle) = latest_run_abort_handle.take()
|
|
{
|
|
abort_handle.abort();
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Resets the mixer's results. Use the all-encompassing [`reset`] API
|
|
/// to clear _all_ of the mixer's state.
|
|
pub fn reset_results(&mut self, ctx: &mut ModelContext<Self>) {
|
|
self.abort_in_flight_async_queries();
|
|
self.query_generation = self.query_generation.wrapping_add(1);
|
|
|
|
self.results.clear();
|
|
self.pending_results = None;
|
|
self.query.take();
|
|
self.initial_results_emitted = false;
|
|
ctx.emit(SearchMixerEvent::ResultsChanged);
|
|
}
|
|
|
|
/// Adds a [`SyncDataSource`] to produce results when the mixer is queried. Query results will
|
|
/// be produced from this source if there are no filters provided or if one of the filters
|
|
/// within a [`Query`] is equal to this filter.
|
|
pub fn add_sync_source(
|
|
&mut self,
|
|
source: impl SyncDataSource<Action = T>,
|
|
filters: impl Into<HashSet<QueryFilter>>,
|
|
) {
|
|
self.sources.insert(
|
|
DataSourceId::new(),
|
|
RegisteredDataSource::new(
|
|
DataSource::SyncDataSource {
|
|
source: Arc::new(source),
|
|
},
|
|
filters.into(),
|
|
),
|
|
);
|
|
}
|
|
|
|
/// Adds an [`AsyncDataSource`] to produce results when the mixer is queried.
|
|
/// The results will be produced asynchronously and the mixer will notify its
|
|
/// subscribers whenever the result set changes.
|
|
///
|
|
/// A debounce interval can be provided to only query the data source in a debounced fashion.
|
|
///
|
|
/// By default, async sources only run when the query's filters explicitly match. Set
|
|
/// `run_when_unfiltered` to `true` so the source also runs when `query.filters` is empty.
|
|
/// Only enable this when the source's work is cheap (e.g. local fuzzy matching) — expensive
|
|
/// operations like network requests should not run on every unfiltered keystroke.
|
|
pub fn add_async_source(
|
|
&mut self,
|
|
source: impl AsyncDataSource<Action = T>,
|
|
filters: impl Into<HashSet<QueryFilter>>,
|
|
options: AddAsyncSourceOptions,
|
|
ctx: &mut ModelContext<Self>,
|
|
) {
|
|
let source = Arc::new(source);
|
|
let data_source_id = DataSourceId::new();
|
|
let debounce_tx = options.debounce_interval.map(|interval| {
|
|
self.start_debounce_stream_for_data_source(data_source_id, interval, ctx)
|
|
});
|
|
|
|
self.sources.insert(
|
|
data_source_id,
|
|
RegisteredDataSource::new(
|
|
DataSource::AsyncDataSource {
|
|
source,
|
|
debounce_tx,
|
|
latest_run_abort_handle: None,
|
|
run_in_zero_state: options.run_in_zero_state,
|
|
run_when_unfiltered: options.run_when_unfiltered,
|
|
},
|
|
filters.into(),
|
|
),
|
|
);
|
|
}
|
|
|
|
pub fn current_query(&self) -> Option<&Query> {
|
|
self.query.as_ref()
|
|
}
|
|
|
|
/// Runs a query against the registered data sources using the provided Query configuration.
|
|
/// On completion, the mixer emits an event to subscribers to indicate the result set has changed.
|
|
///
|
|
/// Old results remain visible while new results are buffered. The visible result set is
|
|
/// replaced atomically once all sources finish, or after [`INITIAL_RESULTS_TIMEOUT`] elapses.
|
|
/// Late-arriving async results are placed at the low-priority edge without reordering existing results.
|
|
pub fn run_query(&mut self, query: Query, ctx: &mut ModelContext<Self>) {
|
|
self.pending_results = Some(Vec::new());
|
|
self.finished_sources.clear();
|
|
self.query = Some(query.clone());
|
|
self.query_generation = self.query_generation.wrapping_add(1);
|
|
self.initial_results_emitted = false;
|
|
let query = &query;
|
|
|
|
// We want to run the queries in the order that the data sources were added.
|
|
let data_source_ids_to_run = self.ordered_data_source_ids_for_query(query).collect_vec();
|
|
for id in data_source_ids_to_run {
|
|
self.run_query_internal(id, false, ctx);
|
|
}
|
|
|
|
// Sync sources (and skipped async sources) will have already finished
|
|
// inside the loop. If everything is done, commit immediately.
|
|
if self.pending_results.is_some() {
|
|
if !self.is_loading() {
|
|
self.commit_pending_results_for_current_query(ctx);
|
|
} else {
|
|
let query_generation = self.query_generation;
|
|
let _ = ctx.spawn(
|
|
async move { Timer::after(INITIAL_RESULTS_TIMEOUT).await },
|
|
move |mixer, _, ctx| {
|
|
mixer.commit_pending_results_after_timeout(query_generation, ctx);
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn results(&self) -> &Vec<QueryResult<T>> {
|
|
&self.results
|
|
}
|
|
|
|
pub fn are_results_empty(&self) -> bool {
|
|
self.results.is_empty()
|
|
}
|
|
|
|
/// Returns all the filters that are currently registered.
|
|
pub fn registered_filters(&self) -> impl Iterator<Item = QueryFilter> + '_ {
|
|
self.sources
|
|
.values()
|
|
.flat_map(|source| source.filters.clone())
|
|
}
|
|
|
|
/// Returns the query filter for the first data source that hasn't completed.
|
|
pub fn loading_query_filters(&self) -> Option<HashSet<QueryFilter>> {
|
|
if self.initial_results_emitted {
|
|
return None;
|
|
}
|
|
let query = self.query.as_ref()?;
|
|
self.ordered_data_source_ids_for_query(query)
|
|
.find(|id| !self.finished_sources.contains(id))
|
|
.and_then(|id| self.sources.get(&id))
|
|
.map(|data_source| data_source.filters.clone())
|
|
}
|
|
|
|
/// Returns true iff there is at least one loading data source.
|
|
/// Helper that computes over `loading_query_filter`.
|
|
pub fn is_loading(&self) -> bool {
|
|
self.loading_query_filters().is_some()
|
|
}
|
|
|
|
/// Returns the first error found from running the data sources against the query, if any.
|
|
pub fn first_data_source_error(
|
|
&self,
|
|
) -> Option<(HashSet<QueryFilter>, &DataSourceRunErrorWrapper)> {
|
|
let query = self.query.as_ref()?;
|
|
self.ordered_data_source_ids_for_query(query)
|
|
.find_map(|id| {
|
|
self.sources
|
|
.get(&id)
|
|
.and_then(|s| Some(s.filters.clone()).zip(s.latest_run_error.as_ref()))
|
|
})
|
|
}
|
|
|
|
/// Returns an ordered list of data source IDs in the order that the corresponding
|
|
/// data sources were registered in.
|
|
/// We could use a map that respects insertion order but that will likely be
|
|
// overkill since the number of data sources is usually minute.
|
|
fn ordered_data_source_ids_for_query<'a>(
|
|
&'a self,
|
|
query: &'a Query,
|
|
) -> impl Iterator<Item = DataSourceId> + 'a {
|
|
self.sources
|
|
.keys()
|
|
.sorted()
|
|
.filter(|id| {
|
|
self.sources
|
|
.get(id)
|
|
.is_some_and(|registered_source| registered_source.matches_query(query))
|
|
})
|
|
.copied()
|
|
}
|
|
|
|
/// Runs the query for the [`DataSource`] identified by the provided `data_source_id`.
|
|
/// If `skip_debounce` is true, then the query is started immediately even if queries
|
|
/// against the data source are meant to be debounced.
|
|
fn run_query_internal(
|
|
&mut self,
|
|
data_source_id: DataSourceId,
|
|
skip_debounce: bool,
|
|
ctx: &mut ModelContext<Self>,
|
|
) {
|
|
let Some(registered_source) = self.sources.get_mut(&data_source_id) else {
|
|
return;
|
|
};
|
|
let Some(query) = self.query.clone() else {
|
|
return;
|
|
};
|
|
|
|
// Clear the latest run error, if any, because we're about to run a new query.
|
|
registered_source.latest_run_error = None;
|
|
|
|
match &mut registered_source.source {
|
|
DataSource::SyncDataSource { source } => {
|
|
let new_results = source.run_query(&query, ctx);
|
|
self.add_new_results(data_source_id, new_results, ctx);
|
|
}
|
|
DataSource::AsyncDataSource {
|
|
source,
|
|
debounce_tx,
|
|
latest_run_abort_handle,
|
|
run_in_zero_state,
|
|
run_when_unfiltered: _,
|
|
} => {
|
|
// Abort any existing run before starting a new one.
|
|
// This is necessary to do even if we end up debouncing
|
|
// because there might already be a running query that's taking long.
|
|
if let Some(abort_handle) = latest_run_abort_handle.take() {
|
|
abort_handle.abort();
|
|
}
|
|
|
|
// Only run async sources in the zero state if the async source indicated it should run in the
|
|
// zero state when registered. It can be costly to run async sources on blank queries so we don't
|
|
// do this by default.
|
|
if query.text.is_empty() && !*run_in_zero_state {
|
|
self.mark_source_as_finished(data_source_id);
|
|
if self.pending_results.is_some() && !self.is_loading() {
|
|
self.commit_pending_results_for_current_query(ctx);
|
|
}
|
|
return;
|
|
}
|
|
|
|
// Check if we should just be debouncing the query rather than running it right now.
|
|
if let Some(debounce_tx) = debounce_tx
|
|
&& !skip_debounce
|
|
{
|
|
let _ = debounce_tx.try_send(DataSourceDebounceArg {});
|
|
return;
|
|
}
|
|
|
|
// If we get here, then we should run the query against the data source right now.
|
|
let query_generation = self.query_generation;
|
|
let source = source.clone();
|
|
let filters = registered_source.filters.to_owned();
|
|
let new_abort_handle = ctx.spawn(
|
|
source.run_query(&query, ctx),
|
|
move |mixer, new_results, ctx| {
|
|
// Discard results from a previous query whose future completed before
|
|
// the abort took effect.
|
|
if mixer.query_generation != query_generation {
|
|
source.on_query_finished(ctx);
|
|
return;
|
|
}
|
|
let error_payload =
|
|
new_results.as_ref().err().map(|e| e.telemetry_payload());
|
|
send_telemetry_from_ctx!(
|
|
TelemetryEvent::CommandSearchAsyncQueryCompleted {
|
|
filters,
|
|
error_payload,
|
|
},
|
|
ctx
|
|
);
|
|
mixer.add_new_results(data_source_id, new_results, ctx);
|
|
source.on_query_finished(ctx);
|
|
},
|
|
);
|
|
*latest_run_abort_handle = Some(new_abort_handle.abort_handle());
|
|
}
|
|
}
|
|
}
|
|
|
|
fn add_new_results(
|
|
&mut self,
|
|
data_source_id: DataSourceId,
|
|
new_results: Result<Vec<QueryResult<T>>, DataSourceRunErrorWrapper>,
|
|
ctx: &mut ModelContext<Self>,
|
|
) {
|
|
if self.finished_sources.contains(&data_source_id) {
|
|
log::warn!(
|
|
"Ignoring duplicate results for source {data_source_id:?} that was already marked finished"
|
|
);
|
|
return;
|
|
}
|
|
self.mark_source_as_finished(data_source_id);
|
|
|
|
match new_results {
|
|
Ok(results) => {
|
|
let results_with_order = results
|
|
.into_iter()
|
|
.map(|mut result| {
|
|
result.source_order = data_source_id.0;
|
|
result
|
|
})
|
|
.collect_vec();
|
|
|
|
if let Some(pending) = &mut self.pending_results {
|
|
pending.extend(results_with_order);
|
|
if !self.is_loading() {
|
|
self.commit_pending_results_for_current_query(ctx);
|
|
}
|
|
} else if self.initial_results_emitted {
|
|
let mut late_results = results_with_order;
|
|
late_results.sort_by_key(|r| (r.priority_tier(), r.score(), r.source_order));
|
|
let mut existing_results = std::mem::take(&mut self.results);
|
|
self.results = late_results;
|
|
self.results.append(&mut existing_results);
|
|
|
|
ctx.emit(SearchMixerEvent::ResultsChanged);
|
|
} else {
|
|
self.results.extend(results_with_order);
|
|
self.sort_results();
|
|
ctx.emit(SearchMixerEvent::ResultsChanged);
|
|
}
|
|
}
|
|
Err(e) => {
|
|
if let Some(source) = self.sources.get_mut(&data_source_id) {
|
|
source.latest_run_error = Some(e);
|
|
}
|
|
|
|
if self.pending_results.is_some() && !self.is_loading() {
|
|
self.commit_pending_results_for_current_query(ctx);
|
|
} else if self.pending_results.is_none() {
|
|
ctx.emit(SearchMixerEvent::ResultsChanged);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Commits buffered results from the current query, replacing the visible result set.
|
|
/// After this, any late-arriving results are added directly to the low-priority edge of
|
|
/// `results`.
|
|
fn commit_pending_results(&mut self, ctx: &mut ModelContext<Self>) {
|
|
let Some(pending) = self.pending_results.take() else {
|
|
return;
|
|
};
|
|
self.results = pending;
|
|
self.sort_results();
|
|
ctx.emit(SearchMixerEvent::ResultsChanged);
|
|
}
|
|
|
|
fn commit_pending_results_for_current_query(&mut self, ctx: &mut ModelContext<Self>) {
|
|
self.initial_results_emitted = true;
|
|
self.commit_pending_results(ctx);
|
|
}
|
|
|
|
/// Sort by (priority_tier, score, source_order) so that equal-scored results
|
|
/// from earlier-registered sources appear first, regardless of async completion order.
|
|
fn sort_results(&mut self) {
|
|
self.results
|
|
.sort_by_key(|r| (r.priority_tier(), r.score(), r.source_order));
|
|
}
|
|
|
|
fn mark_source_as_finished(&mut self, data_source_id: DataSourceId) {
|
|
self.finished_sources.insert(data_source_id);
|
|
}
|
|
|
|
fn commit_pending_results_after_timeout(
|
|
&mut self,
|
|
query_generation: u64,
|
|
ctx: &mut ModelContext<Self>,
|
|
) {
|
|
if query_generation != self.query_generation || self.pending_results.is_none() {
|
|
return;
|
|
}
|
|
self.commit_pending_results_for_current_query(ctx);
|
|
}
|
|
|
|
fn start_debounce_stream_for_data_source(
|
|
&mut self,
|
|
data_source_id: DataSourceId,
|
|
interval: Duration,
|
|
ctx: &mut ModelContext<Self>,
|
|
) -> Sender<DataSourceDebounceArg> {
|
|
let (debounce_tx, debounce_rx) = async_channel::unbounded();
|
|
let _ = ctx.spawn_stream_local(
|
|
debounce(interval, debounce_rx),
|
|
move |mixer, _, ctx| {
|
|
mixer.run_query_internal(data_source_id, true, ctx);
|
|
},
|
|
|_, _| {},
|
|
);
|
|
debounce_tx
|
|
}
|
|
}
|
|
|
|
/// A trait representing a set of data that can be queried for search results synchronously.
|
|
pub trait SyncDataSource: 'static {
|
|
/// The action that is dispatched when a result produced by this data source is
|
|
/// accepted.
|
|
type Action: Action + Clone;
|
|
|
|
fn run_query(
|
|
&self,
|
|
query: &Query,
|
|
app: &AppContext,
|
|
) -> Result<Vec<QueryResult<Self::Action>>, DataSourceRunErrorWrapper>;
|
|
}
|
|
|
|
/// A trait representing a set of data that can be queried for search results asynchronously.
|
|
#[cfg_attr(not(target_family = "wasm"), async_trait)]
|
|
#[cfg_attr(target_family = "wasm", async_trait(?Send))]
|
|
pub trait AsyncDataSource: 'static + Send + Sync {
|
|
/// The action that is dispatched when a result produced by this data source is
|
|
/// accepted.
|
|
type Action: Action + Clone;
|
|
|
|
fn run_query(
|
|
&self,
|
|
query: &Query,
|
|
app: &AppContext,
|
|
) -> BoxFuture<'static, Result<Vec<QueryResult<Self::Action>>, DataSourceRunErrorWrapper>>;
|
|
|
|
/// Function that should be run in the callback after `run_query` finishes.
|
|
fn on_query_finished(&self, _ctx: &mut AppContext) {}
|
|
}
|
|
|
|
/// Helper type alias for a DataSourceRunError.
|
|
pub type DataSourceRunErrorWrapper = Box<dyn DataSourceRunError>;
|
|
|
|
pub trait DataSourceRunError: 'static + Send + Sync + std::fmt::Debug {
|
|
fn user_facing_error(&self) -> String;
|
|
fn telemetry_payload(&self) -> serde_json::Value;
|
|
fn as_any(&self) -> &dyn Any;
|
|
}
|
|
|
|
struct DataSourceDebounceArg {}
|
|
|
|
enum DataSource<T: Action + Clone> {
|
|
SyncDataSource {
|
|
source: Arc<dyn SyncDataSource<Action = T>>,
|
|
},
|
|
AsyncDataSource {
|
|
latest_run_abort_handle: Option<AbortHandle>,
|
|
source: Arc<dyn AsyncDataSource<Action = T>>,
|
|
debounce_tx: Option<Sender<DataSourceDebounceArg>>,
|
|
run_in_zero_state: bool,
|
|
run_when_unfiltered: bool,
|
|
},
|
|
}
|
|
|
|
/// A registered [`DataSource`] for a [`SearchMixer`].
|
|
struct RegisteredDataSource<T: Action + Clone> {
|
|
source: DataSource<T>,
|
|
|
|
/// Corresponding filter for this data source.
|
|
filters: HashSet<QueryFilter>,
|
|
|
|
/// The error produced by this data source during its last run.
|
|
latest_run_error: Option<DataSourceRunErrorWrapper>,
|
|
}
|
|
|
|
impl<T: Action + Clone> RegisteredDataSource<T> {
|
|
/// Sync sources always run when the query has no filters. Async sources only run on
|
|
/// unfiltered queries when `run_when_unfiltered` is set, to avoid running expensive
|
|
/// operations (e.g. network requests) on every keystroke.
|
|
fn matches_query(&self, query: &Query) -> bool {
|
|
match &self.source {
|
|
DataSource::SyncDataSource { .. } => {
|
|
query.filters.is_empty() || query.filters.intersection(&self.filters).count() > 0
|
|
}
|
|
DataSource::AsyncDataSource {
|
|
run_when_unfiltered,
|
|
..
|
|
} => {
|
|
(*run_when_unfiltered && query.filters.is_empty())
|
|
|| query.filters.intersection(&self.filters).count() > 0
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<T: Action + Clone> RegisteredDataSource<T> {
|
|
fn new(source: DataSource<T>, filters: HashSet<QueryFilter>) -> Self {
|
|
Self {
|
|
source,
|
|
filters,
|
|
latest_run_error: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "mixer_tests.rs"]
|
|
mod mixer_test;
|