diff --git a/crates/project/src/buffer_store.rs b/crates/project/src/buffer_store.rs index 3f96a1da968..6e3a570a5e1 100644 --- a/crates/project/src/buffer_store.rs +++ b/crates/project/src/buffer_store.rs @@ -40,6 +40,8 @@ pub struct BufferStore { path_to_buffer_id: HashMap, downstream_client: Option<(AnyProtoClient, u64)>, shared_buffers: HashMap>, + pending_buffer_shares: HashMap>, + next_buffer_share_id: u64, non_searchable_buffers: HashSet, project_search: RemoteProjectSearchState, } @@ -850,6 +852,8 @@ impl BufferStore { opened_buffers: Default::default(), path_to_buffer_id: Default::default(), shared_buffers: Default::default(), + pending_buffer_shares: HashMap::default(), + next_buffer_share_id: 0, loading_buffers: Default::default(), non_searchable_buffers: Default::default(), worktree_store, @@ -877,6 +881,8 @@ impl BufferStore { path_to_buffer_id: Default::default(), loading_buffers: Default::default(), shared_buffers: Default::default(), + pending_buffer_shares: HashMap::default(), + next_buffer_share_id: 0, non_searchable_buffers: Default::default(), worktree_store, project_search: Default::default(), @@ -1635,12 +1641,11 @@ impl BufferStore { return Task::ready(Ok(())); }; - cx.spawn(async move |this, cx| { - let Some(buffer) = this.read_with(cx, |this, _| this.get(buffer_id))? else { - return anyhow::Ok(()); - }; - - let operations = buffer.update(cx, |b, cx| b.serialize_ops(None, cx)); + let buffer = buffer.clone(); + let share_id = util::post_inc(&mut self.next_buffer_share_id); + let (completion_tx, completion_rx) = oneshot::channel(); + let transfer = cx.spawn(async move |buffer_store, cx| { + let operations = buffer.update(cx, |buffer, cx| buffer.serialize_ops(None, cx)); let operations = operations.await; let state = buffer.update(cx, |buffer, cx| buffer.to_proto(cx)); @@ -1673,6 +1678,17 @@ impl BufferStore { .await .log_err(); } + + completion_tx.send(()).ok(); + buffer_store + .update(cx, |store, _| { + store.pending_buffer_shares.remove(&share_id); + }) + .ok(); + }); + self.pending_buffer_shares.insert(share_id, transfer); + cx.background_spawn(async move { + completion_rx.await?; Ok(()) }) } diff --git a/crates/project/src/project_search.rs b/crates/project/src/project_search.rs index edb4b0adf59..ed7da579f06 100644 --- a/crates/project/src/project_search.rs +++ b/crates/project/src/project_search.rs @@ -11,7 +11,7 @@ use std::{ use anyhow::Context; use async_channel::{Receiver, Sender, bounded, unbounded}; -use collections::HashSet; +use collections::HashMap; use fs::{Fs, ReadSeek}; use futures::FutureExt as _; use futures::{SinkExt, StreamExt, select_biased, stream::FuturesOrdered}; @@ -20,6 +20,7 @@ use language::{Buffer, BufferSnapshot, Point}; use parking_lot::Mutex; use postage::oneshot; use rpc::{AnyProtoClient, proto}; +use text::Rope; use encoding_rs::Encoding; use file_content::{ @@ -167,7 +168,7 @@ impl Search { /// Prepares a project search run. The resulting [`SearchResultsHandle`] has to be used to specify whether you're interested in matching buffers /// or full search results. pub fn into_handle(mut self, query: SearchQuery, cx: &mut App) -> SearchResultsHandle { - let mut open_buffers = HashSet::default(); + let mut open_buffers = HashMap::default(); let mut unnamed_buffers = Vec::new(); let mut entryless_file_buffers = Vec::new(); const MAX_CONCURRENT_BUFFER_OPENS: usize = 64; @@ -183,7 +184,7 @@ impl Search { { continue; } else if let Some(entry_id) = buffer.entry_id(cx) { - open_buffers.insert(entry_id); + open_buffers.insert(entry_id, buffer.as_rope().clone()); } else if searches_all_unnamed_buffers { match (&self.kind, buffer.file()) { (SearchKind::Local { .. }, Some(file)) => { @@ -713,7 +714,7 @@ fn path_key_sort_key( struct Worker { query: Arc, - open_buffers: Arc>, + open_buffers: Arc>, candidates: FindSearchCandidates, /// Ok, we're back in background: run full scan & find all matches in a given buffer snapshot. /// Then, when you're done, share them via the channel you were given. @@ -793,7 +794,7 @@ impl Worker { struct RequestHandler<'worker> { query: &'worker SearchQuery, fs: Option<&'worker dyn Fs>, - open_entries: &'worker HashSet, + open_entries: &'worker HashMap, confirm_contents_will_match_tx: &'worker Sender, } @@ -929,9 +930,15 @@ impl RequestHandler<'_> { } } - if self.open_entries.contains(&entry.id) { - // The buffer is already in memory and that's the version we want to scan; - // hence skip the dilly-dally and look for all matches straight away. + if let Some(text) = self.open_entries.get(&entry.id) { + if self + .query + .detect(&mut text.bytes_in_range(0..text.len())) + .await? + .is_none() + { + return Ok(()); + } should_scan_tx .send(( ProjectPath { diff --git a/crates/project/tests/integration/project_tests.rs b/crates/project/tests/integration/project_tests.rs index 5cbbf8e115d..f4642a384b3 100644 --- a/crates/project/tests/integration/project_tests.rs +++ b/crates/project/tests/integration/project_tests.rs @@ -11988,6 +11988,27 @@ async fn test_search(cx: &mut gpui::TestAppContext) { (path!("dir/four.rs").to_string(), vec![25..28, 36..39]) ]) ); + + assert_eq!( + search( + &project, + SearchQuery::text( + "THREE", + false, + true, + false, + Default::default(), + Default::default(), + false, + None, + ) + .unwrap(), + cx + ) + .await + .unwrap(), + HashMap::from_iter([(path!("dir/three.rs").to_string(), vec![6..11])]) + ); } #[gpui::test] diff --git a/crates/remote_server/src/headless_project.rs b/crates/remote_server/src/headless_project.rs index 586d40e7211..bfad457cf8d 100644 --- a/crates/remote_server/src/headless_project.rs +++ b/crates/remote_server/src/headless_project.rs @@ -9,7 +9,8 @@ use lsp::LanguageServerId; use extension::ExtensionHostProxy; use extension_host::headless_host::HeadlessExtensionStore; use fs::Fs; -use gpui::{App, AppContext as _, AsyncApp, Context, Entity, PromptLevel, TaskExt}; +use futures::{channel::oneshot, stream::StreamExt as _}; +use gpui::{App, AppContext as _, AsyncApp, Context, Entity, PromptLevel, Task, TaskExt}; use http_client::HttpClient; use language::{Buffer, BufferEvent, LanguageRegistry, proto::serialize_operation}; use node_runtime::NodeRuntime; @@ -39,6 +40,7 @@ use smol::process::Child; use settings::{Settings as _, SettingsLocation, initial_server_settings_content}; use std::{ + mem, num::NonZeroU64, path::{Path, PathBuf}, sync::{ @@ -74,6 +76,8 @@ pub struct HeadlessProject { // Local variant is used within LSP store, but that's a separate entity. pub _toolchain_store: Entity, pub kernels: HashMap, + search_buffer_batches: async_channel::Sender, + _search_buffer_sender: Task<()>, } pub struct HeadlessAppState { @@ -86,6 +90,17 @@ pub struct HeadlessAppState { pub startup_time: Instant, } +const MAX_SEARCH_CANDIDATES_PER_BATCH: usize = 64; +const MAX_UNACKNOWLEDGED_SEARCH_BYTES: usize = 1024 * 1024; + +struct SearchBufferBatch { + buffers: Vec>, + handle: u64, + peer_id: proto::PeerId, + project_id: u64, + completion: oneshot::Sender>, +} + impl HeadlessProject { pub fn init(cx: &mut App) { settings::init(cx); @@ -343,6 +358,59 @@ impl HeadlessProject { AgentServerStore::init_headless(&session); ContextServerStore::init_headless(&session); + let (search_buffer_batches, batches) = async_channel::bounded::(1); + let search_buffer_sender = cx.spawn({ + let buffer_store = buffer_store.clone(); + let client = session.clone(); + async move |_, cx| { + while let Ok(batch) = batches.recv().await { + let result = async { + let mut buffer_ids = Vec::new(); + let mut unacknowledged_bytes = 0; + for buffer in &batch.buffers { + if batch.completion.is_canceled() { + break; + } + let (buffer_id, transferred_bytes, transfer) = + buffer_store.update(cx, |store, cx| { + let buffer_id = buffer.read(cx).remote_id(); + let transferred_bytes = if store.is_shared(buffer_id, cx) { + 0 + } else { + buffer.read(cx).len() + }; + let transfer = store.create_buffer_for_peer( + buffer, + REMOTE_SERVER_PEER_ID, + cx, + ); + (buffer_id, transferred_bytes, transfer) + }); + transfer.await?; + buffer_ids.push(buffer_id.to_proto()); + unacknowledged_bytes += transferred_bytes; + if unacknowledged_bytes >= MAX_UNACKNOWLEDGED_SEARCH_BYTES { + Self::send_search_candidates( + &client, + &batch, + mem::take(&mut buffer_ids), + ) + .await?; + unacknowledged_bytes = 0; + } + } + if !buffer_ids.is_empty() { + Self::send_search_candidates(&client, &batch, buffer_ids).await?; + } + anyhow::Ok(()) + } + .await; + + batch.completion.send(result).ok(); + } + } + }); + HeadlessProject { next_entry_id: Default::default(), session, @@ -363,6 +431,8 @@ impl HeadlessProject { profiling_collector: gpui::ProfilingCollector::new(startup_time), _toolchain_store: toolchain_store, kernels: Default::default(), + search_buffer_batches, + _search_buffer_sender: search_buffer_sender, } } @@ -1128,8 +1198,6 @@ impl HeadlessProject { envelope: TypedEnvelope, mut cx: AsyncApp, ) -> Result { - 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( @@ -1140,8 +1208,9 @@ impl HeadlessProject { 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 client = this.read_with(&cx, |project, _| project.session.clone()); + let search_buffer_batches = + this.read_with(&cx, |project, _| project.search_buffer_batches.clone()); let task = cx.spawn(async move |cx| { let results = this.update(cx, |this, cx| { project::Search::local( @@ -1154,44 +1223,25 @@ impl HeadlessProject { .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 Some((buffer, _)) = new_matches.next().await { - let _ = buffer_store - .update(cx, |this, cx| { - this.create_buffer_for_peer(&buffer, REMOTE_SERVER_PEER_ID, cx) + let mut batches = std::pin::pin!( + results + .rx + .map(|(buffer, _)| buffer) + .ready_chunks(MAX_SEARCH_CANDIDATES_PER_BATCH) + ); + while let Some(buffers) = batches.next().await { + let (completion, completed) = oneshot::channel(); + search_buffer_batches + .send(SearchBufferBatch { + buffers, + handle, + peer_id, + project_id, + completion, }) - .await; - let buffer_id = buffer.read_with(cx, |this, _| this.remote_id().to_proto()); - batcher.push(buffer_id).await; + .await?; + completed.await??; } - batcher.flush().await; - - sender_task.await?; client .request(proto::FindSearchCandidatesChunk { @@ -1205,13 +1255,31 @@ impl HeadlessProject { .await?; anyhow::Ok(()) }); - _buffer_store.update(&mut cx, |this, _| { - this.register_ongoing_project_search((peer_id, handle), task); + buffer_store.update(&mut cx, |store, _| { + store.register_ongoing_project_search((peer_id, handle), task); }); Ok(proto::Ack {}) } + async fn send_search_candidates( + client: &AnyProtoClient, + batch: &SearchBufferBatch, + buffer_ids: Vec, + ) -> Result<()> { + client + .request(proto::FindSearchCandidatesChunk { + handle: batch.handle, + peer_id: Some(batch.peer_id), + project_id: batch.project_id, + variant: Some(proto::find_search_candidates_chunk::Variant::Matches( + proto::FindSearchCandidatesMatches { buffer_ids }, + )), + }) + .await?; + Ok(()) + } + // Goes from client to host. async fn handle_find_search_candidates_cancel( this: Entity, @@ -1227,7 +1295,6 @@ impl HeadlessProject { envelope: TypedEnvelope, cx: AsyncApp, ) -> Result { - 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 diff --git a/crates/remote_server/src/remote_editing_tests.rs b/crates/remote_server/src/remote_editing_tests.rs index 3ef8e78748e..9b0deb09df6 100644 --- a/crates/remote_server/src/remote_editing_tests.rs +++ b/crates/remote_server/src/remote_editing_tests.rs @@ -20,6 +20,7 @@ use editor::{ }; use extension::ExtensionHostProxy; use fs::{FakeFs, Fs}; +use futures::FutureExt as _; use git::{ Oid, repository::{CommitData, GitCommitTemplate, RepoPath, Worktree as GitWorktree}, @@ -42,6 +43,7 @@ use node_runtime::NodeRuntime; use project::{ CompletionSource, LanguageServerLogType, ProgressToken, Project, ProjectPath, agent_server_store::AgentServerCommand, + buffer_store::BufferStoreEvent, image_store, lsp_store::log_store::{ GlobalLogStore, LanguageServerKind, LanguageServerLogKey, LogKind, LogStore, @@ -66,7 +68,12 @@ use std::{ }, }; use unindent::Unindent as _; -use util::{path, path_list::PathList, paths::PathMatcher, rel_path::rel_path}; +use util::{ + path, + path_list::PathList, + paths::{PathMatcher, PathStyle}, + rel_path::rel_path, +}; #[gpui::test] async fn test_basic_remote_editing(cx: &mut TestAppContext, server_cx: &mut TestAppContext) { @@ -598,6 +605,221 @@ async fn test_remote_project_search(cx: &mut TestAppContext, server_cx: &mut Tes .await; } +#[gpui::test] +async fn test_remote_project_search_backpressure_survives_cancellation( + cx: &mut TestAppContext, + server_cx: &mut TestAppContext, +) { + let fs = FakeFs::new(server_cx.executor()); + let files = (0..4) + .map(|directory| { + let files = (0..128) + .map(|index| (format!("{index:03}.txt"), json!("testing rust-analyzer"))) + .collect::>(); + (directory.to_string(), serde_json::Value::Object(files)) + }) + .collect::>(); + fs.insert_tree(path!("/code/project"), serde_json::Value::Object(files)) + .await; + let (project, headless) = init_test(&fs, cx, server_cx).await; + project + .update(cx, |project, cx| { + project.find_or_create_worktree(path!("/code/project"), true, cx) + }) + .await + .expect("remote worktree should open"); + cx.run_until_parked(); + + let client = project.read_with(cx, |project, cx| { + project.remote_client().unwrap().read(cx).proto_client() + }); + let pause_incoming = client.request(proto::Ping {}); + cx.run_until_parked(); + let received_buffers = Rc::new(RefCell::new(Vec::new())); + let _subscription = cx.update(|cx| { + let received_buffers = received_buffers.clone(); + let buffer_store = project.read(cx).buffer_store().clone(); + cx.subscribe(&buffer_store, move |_, event, cx| { + if let BufferStoreEvent::BufferAdded(buffer) = event { + received_buffers + .borrow_mut() + .push(buffer.read(cx).remote_id()); + } + }) + }); + + let shared_buffers = |cx: &mut TestAppContext| { + headless.read_with(cx, |headless, cx| { + let store = headless.buffer_store.read(cx); + let mut ids = store + .buffers() + .filter_map(|buffer| { + let id = buffer.read(cx).remote_id(); + store.is_shared(id, cx).then_some(id) + }) + .collect::>(); + ids.sort(); + ids + }) + }; + let mut first_batch = None; + let mut searches = Vec::new(); + for directory in 0..4 { + let query = SearchQuery::text( + "testing".to_string(), + false, + true, + false, + PathMatcher::new([format!("{directory}/**")], PathStyle::local()).unwrap(), + PathMatcher::default(), + false, + None, + ) + .unwrap(); + searches.push(project.update(cx, |project, cx| project.search(query, cx))); + cx.run_until_parked(); + let shared = shared_buffers(server_cx); + if let Some(first_batch) = &first_batch { + assert_eq!(&shared, first_batch); + } else { + assert!(!shared.is_empty()); + assert!( + shared.len() <= 64, + "queued {} buffers without an acknowledgement", + shared.len() + ); + first_batch = Some(shared); + } + if directory >= 2 { + searches.clear(); + cx.run_until_parked(); + } + } + + let mut replacement = std::pin::pin!(do_search_and_assert( + &project, + "rust-analyzer", + PathMatcher::new(["3/127.txt"], PathStyle::local()).unwrap(), + false, + &[path!("project/3/127.txt")], + cx.clone(), + )); + assert!(replacement.as_mut().now_or_never().is_none()); + cx.run_until_parked(); + let mut expected_buffers = first_batch.unwrap(); + assert_eq!(shared_buffers(server_cx), expected_buffers); + drop(pause_incoming); + let buffers = replacement.await; + expected_buffers.extend( + buffers + .iter() + .map(|buffer| buffer.read_with(cx, |buffer, _| buffer.remote_id())), + ); + expected_buffers.sort(); + cx.run_until_parked(); + received_buffers.borrow_mut().sort(); + assert_eq!(*received_buffers.borrow(), expected_buffers); +} + +#[gpui::test] +async fn test_remote_project_search_limits_unacknowledged_bytes( + cx: &mut TestAppContext, + server_cx: &mut TestAppContext, +) { + let fs = FakeFs::new(server_cx.executor()); + let text = format!("needle{}", "a".repeat(64 * 1024)); + let files = (0..64) + .map(|index| (format!("{index:03}.txt"), json!(text))) + .collect::>(); + fs.insert_tree(path!("/code/project"), serde_json::Value::Object(files)) + .await; + let (project, headless) = init_test(&fs, cx, server_cx).await; + project + .update(cx, |project, cx| { + project.find_or_create_worktree(path!("/code/project"), true, cx) + }) + .await + .expect("remote worktree should open"); + cx.run_until_parked(); + + let client = project.read_with(cx, |project, cx| { + project.remote_client().unwrap().read(cx).proto_client() + }); + let pause_incoming = client.request(proto::Ping {}); + cx.run_until_parked(); + let expected_paths = (0..64) + .map(|index| format!("project{}{index:03}.txt", std::path::MAIN_SEPARATOR)) + .collect::>(); + let expected_paths = expected_paths + .iter() + .map(String::as_str) + .collect::>(); + let mut search = std::pin::pin!(do_search_and_assert( + &project, + "needle", + PathMatcher::default(), + false, + &expected_paths, + cx.clone(), + )); + let shared_bytes = |cx: &mut TestAppContext| { + headless.read_with(cx, |headless, cx| { + let store = headless.buffer_store.read(cx); + store + .buffers() + .filter(|buffer| store.is_shared(buffer.read(cx).remote_id(), cx)) + .map(|buffer| buffer.read(cx).len()) + .sum::() + }) + }; + assert!(search.as_mut().now_or_never().is_none()); + cx.run_until_parked(); + let first_batch_bytes = shared_bytes(server_cx); + let pause_second_batch = client.request(proto::Ping {}); + cx.run_until_parked(); + drop(pause_incoming); + cx.run_until_parked(); + let second_batch_bytes = shared_bytes(server_cx) - first_batch_bytes; + assert!( + (text.len()..1024 * 1024 + text.len()).contains(&second_batch_bytes), + "shared {second_batch_bytes} bytes without an acknowledgement" + ); + drop(pause_second_batch); + search.await; +} + +#[gpui::test] +async fn test_remote_buffer_transfer_owned_by_buffer_store( + cx: &mut TestAppContext, + server_cx: &mut TestAppContext, +) { + let fs = FakeFs::new(server_cx.executor()); + let (project, headless) = init_test(&fs, cx, server_cx).await; + let buffer_store = headless.read_with(server_cx, |headless, _| headless.buffer_store.clone()); + let buffer = buffer_store.update(server_cx, |store, cx| { + store.create_local_buffer("needle", None, false, cx) + }); + let buffer_id = buffer.read_with(server_cx, |buffer, _| buffer.remote_id()); + buffer.update(server_cx, |buffer, cx| buffer.edit([(6..6, "!")], None, cx)); + let remote_buffer = project.update(cx, |project, cx| { + project + .buffer_store() + .update(cx, |store, cx| store.wait_for_remote_buffer(buffer_id, cx)) + }); + let transfer = buffer_store.update(server_cx, |store, cx| { + store.create_buffer_for_peer(&buffer, proto::REMOTE_SERVER_PEER_ID, cx) + }); + drop(transfer); + server_cx.update(|_| drop(buffer)); + cx.run_until_parked(); + + let remote_buffer = remote_buffer + .now_or_never() + .expect("the transfer should survive dropping its caller") + .expect("the transferred buffer should load"); + remote_buffer.read_with(cx, |buffer, _| assert_eq!(buffer.text(), "needle!")); +} + #[gpui::test] async fn test_remote_project_search_single_cpu( cx: &mut TestAppContext,