Detailed changes
@@ -12597,6 +12597,7 @@ name = "project_benchmarks"
version = "0.1.0"
dependencies = [
"anyhow",
+ "askpass",
"clap",
"client",
"futures 0.3.31",
@@ -12605,6 +12606,10 @@ dependencies = [
"language",
"node_runtime",
"project",
+ "release_channel",
+ "remote",
+ "rpassword",
+ "semver",
"settings",
"watch",
]
@@ -13880,6 +13885,17 @@ version = "0.20.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6c20b6793b5c2fa6553b250154b78d6d0db37e72700ae35fad9387a46f487c97"
+[[package]]
+name = "rpassword"
+version = "7.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "66d4c8b64f049c6721ec8ccec37ddfc3d641c4a7fca57e8f2a89de509c73df39"
+dependencies = [
+ "libc",
+ "rtoolbox",
+ "windows-sys 0.59.0",
+]
+
[[package]]
name = "rpc"
version = "0.1.0"
@@ -13925,6 +13941,16 @@ dependencies = [
"zeroize",
]
+[[package]]
+name = "rtoolbox"
+version = "0.0.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7cc970b249fbe527d6e02e0a227762c9108b2f49d81094fe357ffc6d14d7f6f"
+dependencies = [
+ "libc",
+ "windows-sys 0.52.0",
+]
+
[[package]]
name = "rtrb"
version = "0.3.2"
@@ -468,7 +468,8 @@ impl Server {
.add_request_handler(forward_mutating_project_request::<proto::ToggleLspLogs>)
.add_message_handler(broadcast_project_message_from_host::<proto::LanguageServerLog>)
.add_request_handler(share_agent_thread)
- .add_request_handler(get_shared_agent_thread);
+ .add_request_handler(get_shared_agent_thread)
+ .add_request_handler(forward_project_search_chunk);
Arc::new(server)
}
@@ -2426,6 +2427,20 @@ async fn update_context(message: proto::UpdateContext, session: MessageContext)
Ok(())
}
+async fn forward_project_search_chunk(
+ message: proto::FindSearchCandidatesChunk,
+ response: Response<proto::FindSearchCandidatesChunk>,
+ session: MessageContext,
+) -> Result<()> {
+ let peer_id = message.peer_id.context("missing peer_id")?;
+ let payload = session
+ .peer
+ .forward_request(session.connection_id, peer_id.into(), message)
+ .await?;
+ response.send(payload)?;
+ Ok(())
+}
+
/// Notify other participants that a project has been updated.
async fn broadcast_project_message_from_host<T: EntityMessage<Entity = ShareProject>>(
request: T,
@@ -19,7 +19,7 @@ use language::{
};
use rpc::{
AnyProtoClient, ErrorCode, ErrorExt as _, TypedEnvelope,
- proto::{self},
+ proto::{self, PeerId},
};
use settings::Settings;
@@ -39,6 +39,17 @@ pub struct BufferStore {
downstream_client: Option<(AnyProtoClient, u64)>,
shared_buffers: HashMap<proto::PeerId, HashMap<BufferId, SharedBuffer>>,
non_searchable_buffers: HashSet<BufferId>,
+ project_search: RemoteProjectSearchState,
+}
+
+#[derive(Default)]
+struct RemoteProjectSearchState {
+ // List of ongoing project search chunks from our remote host. Used by the side issuing a search RPC request.
+ chunks: HashMap<u64, smol::channel::Sender<BufferId>>,
+ // Monotonously-increasing handle to hand out to remote host in order to identify the project search result chunk.
+ next_id: u64,
+ // Used by the side running the actual search for match candidates to potentially cancel the search prematurely.
+ searches_in_progress: HashMap<(PeerId, u64), Task<Result<()>>>,
}
#[derive(Hash, Eq, PartialEq, Clone)]
@@ -771,6 +782,7 @@ impl BufferStore {
loading_buffers: Default::default(),
non_searchable_buffers: Default::default(),
worktree_store,
+ project_search: Default::default(),
}
}
@@ -796,6 +808,7 @@ impl BufferStore {
shared_buffers: Default::default(),
non_searchable_buffers: Default::default(),
worktree_store,
+ project_search: Default::default(),
}
}
@@ -1691,6 +1704,82 @@ impl BufferStore {
}
serialized_transaction
}
+
+ pub(crate) fn register_project_search_result_handle(
+ &mut self,
+ ) -> (u64, smol::channel::Receiver<BufferId>) {
+ let (tx, rx) = smol::channel::unbounded();
+ let handle = util::post_inc(&mut self.project_search.next_id);
+ let _old_entry = self.project_search.chunks.insert(handle, tx);
+ debug_assert!(_old_entry.is_none());
+ (handle, rx)
+ }
+
+ pub fn register_ongoing_project_search(
+ &mut self,
+ id: (PeerId, u64),
+ search: Task<anyhow::Result<()>>,
+ ) {
+ let _old = self.project_search.searches_in_progress.insert(id, search);
+ debug_assert!(_old.is_none());
+ }
+
+ pub async fn handle_find_search_candidates_cancel(
+ this: Entity<Self>,
+ envelope: TypedEnvelope<proto::FindSearchCandidatesCancelled>,
+ mut cx: AsyncApp,
+ ) -> Result<()> {
+ let id = (
+ envelope.original_sender_id.unwrap_or(envelope.sender_id),
+ envelope.payload.handle,
+ );
+ let _ = this.update(&mut cx, |this, _| {
+ this.project_search.searches_in_progress.remove(&id)
+ });
+ Ok(())
+ }
+
+ pub(crate) async fn handle_find_search_candidates_chunk(
+ this: Entity<Self>,
+ envelope: TypedEnvelope<proto::FindSearchCandidatesChunk>,
+ mut cx: AsyncApp,
+ ) -> Result<proto::Ack> {
+ use proto::find_search_candidates_chunk::Variant;
+ let handle = envelope.payload.handle;
+
+ let buffer_ids = match envelope
+ .payload
+ .variant
+ .context("Expected non-null variant")?
+ {
+ Variant::Matches(find_search_candidates_matches) => find_search_candidates_matches
+ .buffer_ids
+ .into_iter()
+ .filter_map(|buffer_id| BufferId::new(buffer_id).ok())
+ .collect::<Vec<_>>(),
+ Variant::Done(_) => {
+ this.update(&mut cx, |this, _| {
+ this.project_search.chunks.remove(&handle)
+ });
+ return Ok(proto::Ack {});
+ }
+ };
+ let Some(sender) = this.read_with(&mut cx, |this, _| {
+ this.project_search.chunks.get(&handle).cloned()
+ }) else {
+ return Ok(proto::Ack {});
+ };
+
+ for buffer_id in buffer_ids {
+ let Ok(_) = sender.send(buffer_id).await else {
+ this.update(&mut cx, |this, _| {
+ this.project_search.chunks.remove(&handle)
+ });
+ return Ok(proto::Ack {});
+ };
+ }
+ Ok(proto::Ack {})
+ }
}
impl OpenBuffer {
@@ -11,7 +11,7 @@ pub mod lsp_command;
pub mod lsp_store;
mod manifest_tree;
pub mod prettier_store;
-mod project_search;
+pub mod project_search;
pub mod project_settings;
pub mod search;
mod task_inventory;
@@ -1057,6 +1057,8 @@ impl Project {
client.add_entity_message_handler(Self::handle_create_buffer_for_peer);
client.add_entity_message_handler(Self::handle_toggle_lsp_logs);
client.add_entity_message_handler(Self::handle_create_image_for_peer);
+ client.add_entity_request_handler(Self::handle_find_search_candidates_chunk);
+ client.add_entity_message_handler(Self::handle_find_search_candidates_cancel);
WorktreeStore::init(&client);
BufferStore::init(&client);
@@ -1499,7 +1501,9 @@ impl Project {
remote_proto.add_entity_request_handler(Self::handle_update_buffer_from_remote_server);
remote_proto.add_entity_request_handler(Self::handle_trust_worktrees);
remote_proto.add_entity_request_handler(Self::handle_restrict_worktrees);
+ remote_proto.add_entity_request_handler(Self::handle_find_search_candidates_chunk);
+ remote_proto.add_entity_message_handler(Self::handle_find_search_candidates_cancel);
BufferStore::init(&remote_proto);
LspStore::init(&remote_proto);
SettingsObserver::init(&remote_proto);
@@ -4924,6 +4928,26 @@ impl Project {
Ok(proto::Ack {})
}
+ // Goes from host to client.
+ async fn handle_find_search_candidates_chunk(
+ this: Entity<Self>,
+ envelope: TypedEnvelope<proto::FindSearchCandidatesChunk>,
+ mut cx: AsyncApp,
+ ) -> Result<proto::Ack> {
+ let buffer_store = this.read_with(&mut cx, |this, _| this.buffer_store.clone());
+ BufferStore::handle_find_search_candidates_chunk(buffer_store, envelope, cx).await
+ }
+
+ // Goes from client to host.
+ async fn handle_find_search_candidates_cancel(
+ this: Entity<Self>,
+ envelope: TypedEnvelope<proto::FindSearchCandidatesCancelled>,
+ mut cx: AsyncApp,
+ ) -> Result<()> {
+ let buffer_store = this.read_with(&mut cx, |this, _| this.buffer_store.clone());
+ BufferStore::handle_find_search_candidates_cancel(buffer_store, envelope, cx).await
+ }
+
async fn handle_update_buffer(
this: Entity<Self>,
envelope: TypedEnvelope<proto::UpdateBuffer>,
@@ -5027,32 +5051,78 @@ impl Project {
Ok(response)
}
+ // Goes from client to host.
async fn handle_search_candidate_buffers(
this: Entity<Self>,
envelope: TypedEnvelope<proto::FindSearchCandidates>,
mut cx: AsyncApp,
- ) -> Result<proto::FindSearchCandidatesResponse> {
- let peer_id = envelope.original_sender_id()?;
+ ) -> Result<proto::Ack> {
+ let peer_id = envelope.original_sender_id.unwrap_or(envelope.sender_id);
let message = envelope.payload;
+ let project_id = message.project_id;
let path_style = this.read_with(&cx, |this, cx| this.path_style(cx));
let query =
SearchQuery::from_proto(message.query.context("missing query field")?, path_style)?;
- let results = this.update(&mut cx, |this, cx| {
- this.search_impl(query, cx).matching_buffers(cx)
- });
- let mut response = proto::FindSearchCandidatesResponse {
- buffer_ids: Vec::new(),
- };
-
- while let Ok(buffer) = results.rx.recv().await {
- this.update(&mut cx, |this, cx| {
- let buffer_id = this.create_buffer_for_peer(&buffer, peer_id, cx);
- response.buffer_ids.push(buffer_id.to_proto());
+ let handle = message.handle;
+ let buffer_store = this.read_with(&cx, |this, _| this.buffer_store().clone());
+ let client = this.read_with(&cx, |this, _| this.client());
+ let task = cx.spawn(async move |cx| {
+ let results = this.update(cx, |this, cx| {
+ this.search_impl(query, cx).matching_buffers(cx)
+ });
+ let (batcher, batches) = project_search::AdaptiveBatcher::new(cx.background_executor());
+ let mut new_matches = Box::pin(results.rx);
+
+ let sender_task = cx.background_executor().spawn({
+ let client = client.clone();
+ async move {
+ let mut batches = std::pin::pin!(batches);
+ while let Some(buffer_ids) = batches.next().await {
+ client
+ .request(proto::FindSearchCandidatesChunk {
+ handle,
+ peer_id: Some(peer_id),
+ project_id,
+ variant: Some(
+ proto::find_search_candidates_chunk::Variant::Matches(
+ proto::FindSearchCandidatesMatches { buffer_ids },
+ ),
+ ),
+ })
+ .await?;
+ }
+ anyhow::Ok(())
+ }
});
- }
- Ok(response)
+ while let Some(buffer) = new_matches.next().await {
+ let buffer_id = this.update(cx, |this, cx| {
+ this.create_buffer_for_peer(&buffer, peer_id, cx).to_proto()
+ });
+ batcher.push(buffer_id).await;
+ }
+ batcher.flush().await;
+
+ sender_task.await?;
+
+ let _ = client
+ .request(proto::FindSearchCandidatesChunk {
+ handle,
+ peer_id: Some(peer_id),
+ project_id,
+ variant: Some(proto::find_search_candidates_chunk::Variant::Done(
+ proto::FindSearchCandidatesDone {},
+ )),
+ })
+ .await?;
+ anyhow::Ok(())
+ });
+ buffer_store.update(&mut cx, |this, _| {
+ this.register_ongoing_project_search((peer_id, handle), task);
+ });
+
+ Ok(proto::Ack {})
}
async fn handle_open_buffer_by_id(
@@ -6,23 +6,21 @@ use std::{
path::{Path, PathBuf},
pin::pin,
sync::Arc,
+ time::Duration,
};
use anyhow::Context;
use collections::HashSet;
use fs::Fs;
+use futures::FutureExt as _;
use futures::{SinkExt, StreamExt, select_biased, stream::FuturesOrdered};
-use gpui::{App, AppContext, AsyncApp, Entity, Task};
+use gpui::{App, AppContext, AsyncApp, BackgroundExecutor, Entity, Priority, Task};
use language::{Buffer, BufferSnapshot};
use parking_lot::Mutex;
use postage::oneshot;
use rpc::{AnyProtoClient, proto};
-use smol::{
- channel::{Receiver, Sender, bounded, unbounded},
- future::FutureExt,
-};
+use smol::channel::{Receiver, Sender, bounded, unbounded};
-use text::BufferId;
use util::{ResultExt, maybe, paths::compare_rel_paths, rel_path::RelPath};
use worktree::{Entry, ProjectEntryId, Snapshot, Worktree, WorktreeSettings};
@@ -249,12 +247,26 @@ impl Search {
remote_id,
models,
} => {
+ let (handle, rx) = self
+ .buffer_store
+ .update(cx, |this, _| this.register_project_search_result_handle());
+
+ let cancel_ongoing_search = util::defer({
+ let client = client.clone();
+ move || {
+ _ = client.send(proto::FindSearchCandidatesCancelled {
+ project_id: remote_id,
+ handle,
+ });
+ }
+ });
let request = client.request(proto::FindSearchCandidates {
project_id: remote_id,
query: Some(query.to_proto()),
limit: self.limit as _,
+ handle,
});
- let weak_buffer_store = self.buffer_store.downgrade();
+
let buffer_store = self.buffer_store;
let guard = cx.update(|cx| {
Project::retain_remotely_created_models_impl(
@@ -268,18 +280,39 @@ impl Search {
let issue_remote_buffers_request = cx
.spawn(async move |cx| {
let _ = maybe!(async move {
- let response = request.await?;
- for buffer_id in response.buffer_ids {
- let buffer_id = BufferId::new(buffer_id)?;
- let buffer = weak_buffer_store
- .update(cx, |buffer_store, cx| {
- buffer_store.wait_for_remote_buffer(buffer_id, cx)
- })?
- .await?;
- let _ = grab_buffer_snapshot_tx.send(buffer).await;
- }
+ request.await?;
+
+ let (buffer_tx, buffer_rx) = bounded(24);
+
+ let wait_for_remote_buffers = cx.spawn(async move |cx| {
+ while let Ok(buffer_id) = rx.recv().await {
+ let buffer =
+ buffer_store.update(cx, |buffer_store, cx| {
+ buffer_store
+ .wait_for_remote_buffer(buffer_id, cx)
+ });
+ buffer_tx.send(buffer).await?;
+ }
+ anyhow::Ok(())
+ });
+
+ let forward_buffers = cx.background_spawn(async move {
+ while let Ok(buffer) = buffer_rx.recv().await {
+ let _ =
+ grab_buffer_snapshot_tx.send(buffer.await?).await;
+ }
+ anyhow::Ok(())
+ });
+ let (left, right) = futures::future::join(
+ wait_for_remote_buffers,
+ forward_buffers,
+ )
+ .await;
+ left?;
+ right?;
drop(guard);
+ cancel_ongoing_search.abort();
anyhow::Ok(())
})
.await
@@ -894,6 +927,104 @@ impl PathInclusionMatcher {
}
}
+type IsTerminating = bool;
+/// Adaptive batcher that starts eager (small batches) and grows batch size
+/// when items arrive quickly, reducing RPC overhead while preserving low latency
+/// for slow streams.
+pub struct AdaptiveBatcher<T> {
+ items: Sender<T>,
+ flush_batch: Sender<IsTerminating>,
+ _batch_task: Task<()>,
+}
+
+impl<T: 'static + Send> AdaptiveBatcher<T> {
+ pub fn new(cx: &BackgroundExecutor) -> (Self, Receiver<Vec<T>>) {
+ let (items, rx) = unbounded();
+ let (batch_tx, batch_rx) = unbounded();
+ let (flush_batch_tx, flush_batch_rx) = unbounded();
+ let flush_batch = flush_batch_tx.clone();
+ let executor = cx.clone();
+ let _batch_task = cx.spawn_with_priority(gpui::Priority::High, async move {
+ let mut current_batch = vec![];
+ let mut items_produced_so_far = 0_u64;
+
+ let mut _schedule_flush_after_delay: Option<Task<()>> = None;
+ let _time_elapsed_since_start_of_search = std::time::Instant::now();
+ let mut flush = pin!(flush_batch_rx);
+ let mut terminating = false;
+ loop {
+ select_biased! {
+ item = rx.recv().fuse() => {
+ match item {
+ Ok(new_item) => {
+ let is_fresh_batch = current_batch.is_empty();
+ items_produced_so_far += 1;
+ current_batch.push(new_item);
+ if is_fresh_batch {
+ // Chosen arbitrarily based on some experimentation with plots.
+ let desired_duration_ms = (20 * (items_produced_so_far + 2).ilog2() as u64).min(300);
+ let desired_duration = Duration::from_millis(desired_duration_ms);
+ let _executor = executor.clone();
+ let _flush = flush_batch_tx.clone();
+ let new_timer = executor.spawn_with_priority(Priority::High, async move {
+ _executor.timer(desired_duration).await;
+ _ = _flush.send(false).await;
+ });
+ _schedule_flush_after_delay = Some(new_timer);
+ }
+ }
+ Err(_) => {
+ // Items channel closed - send any remaining batch before exiting
+ if !current_batch.is_empty() {
+ _ = batch_tx.send(std::mem::take(&mut current_batch)).await;
+ }
+ break;
+ }
+ }
+ }
+ should_break_afterwards = flush.next() => {
+ if !current_batch.is_empty() {
+ _ = batch_tx.send(std::mem::take(&mut current_batch)).await;
+ _schedule_flush_after_delay = None;
+ }
+ if should_break_afterwards.unwrap_or_default() {
+ terminating = true;
+ }
+ }
+ complete => {
+ break;
+ }
+ }
+ if terminating {
+ // Drain any remaining items before exiting
+ while let Ok(new_item) = rx.try_recv() {
+ current_batch.push(new_item);
+ }
+ if !current_batch.is_empty() {
+ _ = batch_tx.send(std::mem::take(&mut current_batch)).await;
+ }
+ break;
+ }
+ }
+ });
+ let this = Self {
+ items,
+ _batch_task,
+ flush_batch,
+ };
+ (this, batch_rx)
+ }
+
+ pub async fn push(&self, item: T) {
+ _ = self.items.send(item).await;
+ }
+
+ pub async fn flush(self) {
+ _ = self.flush_batch.send(true).await;
+ self._batch_task.await;
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -6,6 +6,7 @@ edition.workspace = true
[dependencies]
anyhow.workspace = true
+askpass.workspace = true
clap.workspace = true
client.workspace = true
futures.workspace = true
@@ -14,6 +15,10 @@ http_client = { workspace = true, features = ["test-support"]}
language.workspace = true
node_runtime.workspace = true
project.workspace = true
+release_channel.workspace = true
+remote = { workspace = true, features = ["build-remote-server-binary"] }
+rpassword = "7.4"
+semver.workspace = true
settings.workspace = true
watch.workspace = true
@@ -1,7 +1,10 @@
-use std::sync::Arc;
+use std::{sync::Arc, time::Duration};
+use anyhow::anyhow;
+use askpass::EncryptedPassword;
use clap::Parser;
use client::{Client, UserStore};
+use futures::channel::oneshot;
use gpui::{AppContext as _, Application};
use http_client::FakeHttpClient;
use language::LanguageRegistry;
@@ -10,13 +13,19 @@ use project::{
Project, RealFs,
search::{SearchQuery, SearchResult},
};
+use release_channel::ReleaseChannel;
+use remote::{ConnectionIdentifier, RemoteClientDelegate, SshConnectionOptions};
+use semver::Version;
#[derive(Parser)]
struct Args {
/// List of worktrees to run the search against.
worktrees: Vec<String>,
#[clap(short)]
- query: String,
+ query: Option<String>,
+ /// Askpass socket for SSH authentication
+ #[clap(long)]
+ askpass: Option<String>,
/// Treat query as a regex.
#[clap(short, long)]
regex: bool,
@@ -29,13 +38,72 @@ struct Args {
/// Include gitignored files in the search.
#[clap(long)]
include_ignored: bool,
+ #[clap(long)]
+ ssh: Option<String>,
}
+struct BenchmarkRemoteClient;
+impl RemoteClientDelegate for BenchmarkRemoteClient {
+ fn ask_password(
+ &self,
+ prompt: String,
+ tx: oneshot::Sender<EncryptedPassword>,
+ _cx: &mut gpui::AsyncApp,
+ ) {
+ eprintln!("SSH asking for password: {}", prompt);
+ match rpassword::prompt_password(&prompt) {
+ Ok(password) => match EncryptedPassword::try_from(password.as_ref()) {
+ Ok(encrypted) => {
+ if tx.send(encrypted).is_err() {
+ eprintln!("Failed to send password");
+ }
+ }
+ Err(e) => eprintln!("Failed to encrypt password: {e}"),
+ },
+ Err(e) => eprintln!("Failed to read password: {e}"),
+ }
+ }
+
+ fn get_download_url(
+ &self,
+ _platform: remote::RemotePlatform,
+ _release_channel: ReleaseChannel,
+ _version: Option<Version>,
+ _cx: &mut gpui::AsyncApp,
+ ) -> gpui::Task<gpui::Result<Option<String>>> {
+ unimplemented!()
+ }
+
+ fn download_server_binary_locally(
+ &self,
+ _platform: remote::RemotePlatform,
+ _release_channel: ReleaseChannel,
+ _version: Option<Version>,
+ _cx: &mut gpui::AsyncApp,
+ ) -> gpui::Task<gpui::Result<std::path::PathBuf>> {
+ unimplemented!()
+ }
+
+ fn set_status(&self, status: Option<&str>, _: &mut gpui::AsyncApp) {
+ if let Some(status) = status {
+ println!("SSH status: {status}");
+ }
+ }
+}
fn main() -> Result<(), anyhow::Error> {
let args = Args::parse();
+
+ if let Some(socket) = &args.askpass {
+ askpass::main(socket);
+ return Ok(());
+ }
+
+ let query_str = args
+ .query
+ .ok_or_else(|| anyhow!("-q/--query is required"))?;
let query = if args.regex {
SearchQuery::regex(
- args.query,
+ query_str,
args.whole_word,
args.case_sensitive,
args.include_ignored,
@@ -47,7 +115,7 @@ fn main() -> Result<(), anyhow::Error> {
)
} else {
SearchQuery::text(
- args.query,
+ query_str,
args.whole_word,
args.case_sensitive,
args.include_ignored,
@@ -58,6 +126,7 @@ fn main() -> Result<(), anyhow::Error> {
)
}?;
Application::headless().run(|cx| {
+ release_channel::init_test(semver::Version::new(0, 0, 0), ReleaseChannel::Dev, cx);
settings::init(cx);
let client = Client::production(cx);
let http_client = FakeHttpClient::with_200_response();
@@ -66,19 +135,35 @@ fn main() -> Result<(), anyhow::Error> {
let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
let registry = Arc::new(LanguageRegistry::new(cx.background_executor().clone()));
let fs = Arc::new(RealFs::new(None, cx.background_executor().clone()));
- let project = Project::local(
- client,
- node,
- user_store,
- registry,
- fs,
- Some(Default::default()),
- false,
- cx,
- );
- project.clone().update(cx, move |_, cx| {
- cx.spawn(async move |_, cx| {
+
+
+ cx.spawn(async move |cx| {
+ let project = if let Some(ssh_target) = args.ssh {
+ println!("Setting up SSH connection for {}", &ssh_target);
+ let ssh_connection_options = SshConnectionOptions::parse_command_line(&ssh_target)?;
+
+ let connection_options = remote::RemoteConnectionOptions::from(ssh_connection_options);
+ let delegate = Arc::new(BenchmarkRemoteClient);
+ let remote_connection = remote::connect(connection_options.clone(), delegate.clone(), cx).await.unwrap();
+
+ let (_tx, rx) = oneshot::channel();
+ let remote_client = cx.update(|cx| remote::RemoteClient::new(ConnectionIdentifier::setup(), remote_connection, rx, delegate.clone(), cx )).await?.ok_or_else(|| anyhow!("ssh initialization returned None"))?;
+
+ cx.update(|cx| Project::remote(remote_client, client, node, user_store, registry, fs, false, cx))
+ } else {
+ println!("Setting up local project");
+ cx.update(|cx| Project::local(
+ client,
+ node,
+ user_store,
+ registry,
+ fs,
+ Some(Default::default()),
+ false,
+ cx,
+ ))
+ };
println!("Loading worktrees");
let worktrees = project.update(cx, |this, cx| {
args.worktrees
@@ -93,9 +178,20 @@ fn main() -> Result<(), anyhow::Error> {
.collect::<Result<Vec<_>, anyhow::Error>>()?;
for (worktree, _) in &worktrees {
- worktree
- .update(cx, |this, _| this.as_local().unwrap().scan_complete())
- .await;
+ let scan_complete = worktree
+ .update(cx, |this, _| {
+ if let Some(local) = this.as_local() {
+ Some(local.scan_complete())
+ } else {
+ None
+ }
+ });
+ if let Some(scan_complete) = scan_complete {
+ scan_complete.await;
+ } else {
+ cx.background_executor().timer(Duration::from_secs(10)).await;
+ }
+
}
println!("Worktrees loaded");
@@ -127,8 +223,8 @@ fn main() -> Result<(), anyhow::Error> {
anyhow::Ok(())
})
- .detach();
- });
+ .detach_and_log_err(cx);
+
});
Ok(())
}
@@ -308,8 +308,28 @@ message FindSearchCandidates {
uint64 project_id = 1;
SearchQuery query = 2;
uint64 limit = 3;
+ uint64 handle = 4;
}
-message FindSearchCandidatesResponse {
- repeated uint64 buffer_ids = 1;
+
+message FindSearchCandidatesDone {}
+
+message FindSearchCandidatesMatches {
+ repeated uint64 buffer_ids = 1;
+}
+
+message FindSearchCandidatesChunk {
+ uint64 project_id = 1;
+ PeerId peer_id = 2;
+ uint64 handle = 3;
+
+ oneof variant {
+ FindSearchCandidatesMatches matches = 4;
+ FindSearchCandidatesDone done = 5;
+ }
+}
+
+message FindSearchCandidatesCancelled {
+ uint64 project_id = 1;
+ uint64 handle = 2;
}
@@ -244,7 +244,6 @@ message Envelope {
LspExtSwitchSourceHeaderResponse lsp_ext_switch_source_header_response = 242;
FindSearchCandidates find_search_candidates = 243;
- FindSearchCandidatesResponse find_search_candidates_response = 244;
CloseBuffer close_buffer = 245;
@@ -448,10 +447,13 @@ message Envelope {
TrustWorktrees trust_worktrees = 404;
RestrictWorktrees restrict_worktrees = 405;
-
+
ShareAgentThread share_agent_thread = 406;
GetSharedAgentThread get_shared_agent_thread = 407;
- GetSharedAgentThreadResponse get_shared_agent_thread_response = 408; // current max
+ GetSharedAgentThreadResponse get_shared_agent_thread_response = 408;
+
+ FindSearchCandidatesChunk find_search_candidates_chunk = 409;
+ FindSearchCandidatesCancelled find_search_candidates_cancelled = 410; // current max
}
reserved 87 to 88;
@@ -469,6 +471,7 @@ message Envelope {
reserved 224 to 231;
reserved 234 to 236;
reserved 239 to 240;
+ reserved 244;
reserved 246 to 256;
reserved 259;
reserved 270;
@@ -68,7 +68,6 @@ messages!(
(Error, Foreground),
(ExpandProjectEntry, Foreground),
(ExpandProjectEntryResponse, Foreground),
- (FindSearchCandidatesResponse, Background),
(FindSearchCandidates, Background),
(FlushBufferedMessages, Foreground),
(ExpandAllForProjectEntry, Foreground),
@@ -345,7 +344,9 @@ messages!(
(GitCreateWorktree, Background),
(ShareAgentThread, Foreground),
(GetSharedAgentThread, Foreground),
- (GetSharedAgentThreadResponse, Foreground)
+ (GetSharedAgentThreadResponse, Foreground),
+ (FindSearchCandidatesChunk, Background),
+ (FindSearchCandidatesCancelled, Background),
);
request_messages!(
@@ -440,7 +441,7 @@ request_messages!(
(RespondToContactRequest, Ack),
(SaveBuffer, BufferSaved),
(Stage, Ack),
- (FindSearchCandidates, FindSearchCandidatesResponse),
+ (FindSearchCandidates, Ack),
(SendChannelMessage, SendChannelMessageResponse),
(SetChannelMemberRole, Ack),
(SetChannelVisibility, Ack),
@@ -535,6 +536,7 @@ request_messages!(
(GitCreateWorktree, Ack),
(TrustWorktrees, Ack),
(RestrictWorktrees, Ack),
+ (FindSearchCandidatesChunk, Ack),
);
lsp_messages!(
@@ -710,6 +712,8 @@ entity_messages!(
GitCreateWorktree,
TrustWorktrees,
RestrictWorktrees,
+ FindSearchCandidatesChunk,
+ FindSearchCandidatesCancelled,
);
entity_messages!(
@@ -32,7 +32,6 @@ use rpc::{
};
use settings::initial_server_settings_content;
-use smol::stream::StreamExt;
use std::{
num::NonZeroU64,
path::{Path, PathBuf},
@@ -288,6 +287,7 @@ impl HeadlessProject {
session.add_entity_request_handler(Self::handle_trust_worktrees);
session.add_entity_request_handler(Self::handle_restrict_worktrees);
+ session.add_entity_message_handler(Self::handle_find_search_candidates_cancel);
session.add_entity_request_handler(BufferStore::handle_update_buffer);
session.add_entity_message_handler(BufferStore::handle_close_buffer);
@@ -779,41 +779,99 @@ impl HeadlessProject {
this: Entity<Self>,
envelope: TypedEnvelope<proto::FindSearchCandidates>,
mut cx: AsyncApp,
- ) -> Result<proto::FindSearchCandidatesResponse> {
+ ) -> Result<proto::Ack> {
+ use futures::stream::StreamExt as _;
+
+ let peer_id = envelope.original_sender_id.unwrap_or(envelope.sender_id);
let message = envelope.payload;
let query = SearchQuery::from_proto(
message.query.context("missing query field")?,
PathStyle::local(),
)?;
- let results = this.update(&mut cx, |this, cx| {
- project::Search::local(
- this.fs.clone(),
- this.buffer_store.clone(),
- this.worktree_store.clone(),
- message.limit as _,
- cx,
- )
- .into_handle(query, cx)
- .matching_buffers(cx)
- });
-
- let mut response = proto::FindSearchCandidatesResponse {
- buffer_ids: Vec::new(),
- };
+ let project_id = message.project_id;
let buffer_store = this.read_with(&cx, |this, _| this.buffer_store.clone());
+ let handle = message.handle;
+ let _buffer_store = buffer_store.clone();
+ let client = this.read_with(&cx, |this, _| this.session.clone());
+ let task = cx.spawn(async move |cx| {
+ let results = this.update(cx, |this, cx| {
+ project::Search::local(
+ this.fs.clone(),
+ this.buffer_store.clone(),
+ this.worktree_store.clone(),
+ message.limit as _,
+ cx,
+ )
+ .into_handle(query, cx)
+ .matching_buffers(cx)
+ });
+ let (batcher, batches) =
+ project::project_search::AdaptiveBatcher::new(cx.background_executor());
+ let mut new_matches = Box::pin(results.rx);
+
+ let sender_task = cx.background_executor().spawn({
+ let client = client.clone();
+ async move {
+ let mut batches = std::pin::pin!(batches);
+ while let Some(buffer_ids) = batches.next().await {
+ client
+ .request(proto::FindSearchCandidatesChunk {
+ handle,
+ peer_id: Some(peer_id),
+ project_id,
+ variant: Some(
+ proto::find_search_candidates_chunk::Variant::Matches(
+ proto::FindSearchCandidatesMatches { buffer_ids },
+ ),
+ ),
+ })
+ .await?;
+ }
+ anyhow::Ok(())
+ }
+ });
- while let Ok(buffer) = results.rx.recv().await {
- let buffer_id = buffer.read_with(&cx, |this, _| this.remote_id());
- response.buffer_ids.push(buffer_id.to_proto());
- buffer_store
- .update(&mut cx, |buffer_store, cx| {
- buffer_store.create_buffer_for_peer(&buffer, REMOTE_SERVER_PEER_ID, cx)
+ while let Some(buffer) = new_matches.next().await {
+ let _ = buffer_store
+ .update(cx, |this, cx| {
+ this.create_buffer_for_peer(&buffer, REMOTE_SERVER_PEER_ID, cx)
+ })
+ .await;
+ let buffer_id = buffer.read_with(cx, |this, _| this.remote_id().to_proto());
+ batcher.push(buffer_id).await;
+ }
+ batcher.flush().await;
+
+ sender_task.await?;
+
+ client
+ .request(proto::FindSearchCandidatesChunk {
+ handle,
+ peer_id: Some(peer_id),
+ project_id,
+ variant: Some(proto::find_search_candidates_chunk::Variant::Done(
+ proto::FindSearchCandidatesDone {},
+ )),
})
.await?;
- }
+ anyhow::Ok(())
+ });
+ _buffer_store.update(&mut cx, |this, _| {
+ this.register_ongoing_project_search((peer_id, handle), task);
+ });
- Ok(response)
+ Ok(proto::Ack {})
+ }
+
+ // Goes from client to host.
+ async fn handle_find_search_candidates_cancel(
+ this: Entity<Self>,
+ envelope: TypedEnvelope<proto::FindSearchCandidatesCancelled>,
+ mut cx: AsyncApp,
+ ) -> Result<()> {
+ let buffer_store = this.read_with(&mut cx, |this, _| this.buffer_store.clone());
+ BufferStore::handle_find_search_candidates_cancel(buffer_store, envelope, cx).await
}
async fn handle_list_remote_directory(
@@ -821,6 +879,7 @@ impl HeadlessProject {
envelope: TypedEnvelope<proto::ListRemoteDirectory>,
cx: AsyncApp,
) -> Result<proto::ListRemoteDirectoryResponse> {
+ use smol::stream::StreamExt;
let fs = cx.read_entity(&this, |this, _| this.fs.clone());
let expanded = PathBuf::from(shellexpand::tilde(&envelope.payload.path).to_string());
let check_info = envelope
@@ -237,11 +237,12 @@ async fn test_remote_project_search(cx: &mut TestAppContext, server_cx: &mut Tes
assert!(headless.buffer_store.read(cx).has_shared_buffers())
});
do_search(&project, cx.clone()).await;
-
+ server_cx.run_until_parked();
cx.update(|_| {
drop(buffer);
});
cx.run_until_parked();
+ server_cx.run_until_parked();
headless.update(server_cx, |headless, cx| {
assert!(!headless.buffer_store.read(cx).has_shared_buffers())
});