History

git_graph: Add remote support for search operations (#55167)

### Motivation This is the second of three PRs to add remote/collab support for the git graph and is a follow-up to #54468. I'm adding remote support for the search because it's not user accessible without the initial graph fetch having remote support, so it allows us to merge this without having to add full remote support. Collab guest support will be added in a follow-up PR. #### Summary For large repos, searching can take a while to fully stream in all matched results. For example, running a basic search on the Linux repo took over 10s for me. Because of that, we want to stream search results in chunks to downstream users to keep the time-to-first-match low. After this change, the first chunk gets sent back after ~50ms on the Linux repo from receiving the request. In order to accomplish that, I added a new proto client API that allows for a request to map to n responses. e.g. ```/dev/null/example.rs#L1-1 client.add_entity_stream_request_handler(Self::handle_search_commits); ``` Note: The proto API isn't supported over collab yet, that will be another PR Self-Review Checklist: - [x] I've reviewed my own diff for quality, security, and reliability - [x] Unsafe blocks (if any) have justifying comments - [x] The content is consistent with the [UI/UX checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist) - [x] Tests cover the new/changed behavior - [x] Performance impact has been considered and is acceptable Closes #ISSUE Release Notes: - N/A ---------

Anthony Eid committed 6 months ago

Commit 149cd4e1 parent 8aedcbf10 files · +1016 −60Browse files

Unified · Split

10 files
  1. crates/client/src/client.rs +29 −0
  2. crates/fs/src/fake_git_repo.rs +36 −3
  3. crates/project/src/git_store.rs +129 −3
  4. crates/proto/proto/git.proto +23 −0
  5. crates/proto/proto/zed.proto +3 −1
  6. crates/proto/src/proto.rs +4 −0
  7. crates/remote/src/remote_client.rs +225 −3
  8. crates/remote_server/src/remote_editing_tests.rs +108 −1
  9. crates/rpc/src/peer.rs +364 −48
  10. crates/rpc/src/proto_client.rs +95 −1
crates/client/src/client.rs+29−0
@@ -24,6 +24,7 @@ use futures::{ · 23 unmodified lines
2424 AsyncReadExt, FutureExt, SinkExt, Stream, StreamExt, TryFutureExt as _, TryStreamExt,
2525 channel::{mpsc, oneshot},
2626 future::BoxFuture,
27+ stream::BoxStream,
2728 };
2829 use gpui::{App, AsyncApp, Entity, Global, Task, WeakEntity, actions};
2930 use http_client::{HttpClient, HttpClientWithUrl, http, read_proxy_from_env};
@@ -1789,6 +1790,34 @@ impl ProtoClient for Client { · 1759 unmodified lines
17891790 self.request_dynamic(envelope, request_type).boxed()
17901791 }
17911792
1793+ fn request_stream(
1794+ &self,
1795+ envelope: proto::Envelope,
1796+ request_type: &'static str,
1797+ ) -> BoxFuture<'static, Result<BoxStream<'static, Result<proto::Envelope>>>> {
1798+ let client_id = self.id();
1799+ let response = self.connection_id().map(|connection_id| {
1800+ self.peer
1801+ .request_stream_dynamic(connection_id, envelope, request_type)
1802+ });
1803+
1804+ async move {
1805+ log::debug!(
1806+ "rpc stream request start. client_id:{}. name:{}",
1807+ client_id,
1808+ request_type
1809+ );
1810+ let response = response?.await;
1811+ log::debug!(
1812+ "rpc stream request opened. client_id:{}. name:{}",
1813+ client_id,
1814+ request_type
1815+ );
1816+ response
1817+ }
1818+ .boxed()
1819+ }
1820+
17921821 fn send(&self, envelope: proto::Envelope, message_type: &'static str) -> Result<()> {
17931822 log::debug!("rpc send. client_id:{}, name:{}", self.id(), message_type);
17941823 let connection_id = self.connection_id()?;
crates/fs/src/fake_git_repo.rs+36−3
@@ -1435,10 +1435,43 @@ impl GitRepository for FakeGitRepository { · 1434 unmodified lines
14351435 fn search_commits(
14361436 &self,
14371437 _log_source: LogSource,
1438− _search_args: SearchCommitArgs,
1439− _request_tx: Sender<Oid>,
1438+ search_args: SearchCommitArgs,
1439+ request_tx: Sender<Oid>,
14401440 ) -> BoxFuture<'_, Result<()>> {
1441− async { bail!("search_commits not supported for FakeGitRepository") }.boxed()
1441+ async move {
1442+ let query = if search_args.case_sensitive {
1443+ search_args.query.to_string()
1444+ } else {
1445+ search_args.query.to_lowercase()
1446+ };
1447+
1448+ let matching_shas = self.fs.with_git_state(&self.dot_git_path, false, |state| {
1449+ state
1450+ .commit_data
1451+ .iter()
1452+ .filter_map(|(sha, entry)| {
1453+ let FakeCommitDataEntry::Success(commit_data) = entry else {
1454+ return None;
1455+ };
1456+ let message = if search_args.case_sensitive {
1457+ commit_data.message.to_string()
1458+ } else {
1459+ commit_data.message.to_lowercase()
1460+ };
1461+ message.contains(&query).then_some(*sha)
1462+ })
1463+ .collect::<Vec<_>>()
1464+ })?;
1465+
1466+ for sha in matching_shas {
1467+ if request_tx.send(sha).await.is_err() {
1468+ break;
1469+ }
1470+ }
1471+
1472+ Ok(())
1473+ }
1474+ .boxed()
14421475 }
14431476
14441477 fn commit_data_reader(&self) -> Result<CommitDataReader> {
crates/project/src/git_store.rs+129−3
@@ -20,7 +20,7 @@ use collections::HashMap; · 19 unmodified lines
2020 pub use conflict_set::{ConflictRegion, ConflictSet, ConflictSetSnapshot, ConflictSetUpdate};
2121 use fs::{Fs, RemoveOptions};
2222 use futures::{
23− FutureExt, StreamExt,
23+ FutureExt, SinkExt, Stream, StreamExt,
2424 channel::{
2525 mpsc,
2626 oneshot::{self, Canceled},
@@ -680,6 +680,7 @@ impl GitStore { · 653 unmodified lines
680680 client.add_entity_request_handler(Self::handle_edit_ref);
681681 client.add_entity_request_handler(Self::handle_repair_worktrees);
682682 client.add_entity_request_handler(Self::handle_get_commit_data);
683+ client.add_entity_stream_request_handler(Self::handle_search_commits);
683684 }
684685
685686 pub fn is_local(&self) -> bool {
@@ -2669,6 +2670,63 @@ impl GitStore { · 1983 unmodified lines
26692670 Ok(proto::GetCommitDataResponse { commits })
26702671 }
26712672
2673+ async fn handle_search_commits(
2674+ this: Entity<Self>,
2675+ envelope: TypedEnvelope<proto::SearchCommits>,
2676+ mut cx: AsyncApp,
2677+ ) -> Result<impl Stream<Item = Result<proto::SearchCommitsResponse>>> {
2678+ const CHUNK_SIZE: usize = 100;
2679+
2680+ let repository_id = RepositoryId::from_proto(envelope.payload.repository_id);
2681+ let repository_handle = Self::repository_for_request(&this, repository_id, &mut cx)?;
2682+ let log_source = log_source_from_proto(
2683+ envelope
2684+ .payload
2685+ .log_source
2686+ .context("missing search commit log source")?,
2687+ )?;
2688+ let search_args = SearchCommitArgs {
2689+ query: SharedString::from(envelope.payload.query),
2690+ case_sensitive: envelope.payload.case_sensitive,
2691+ };
2692+
2693+ let (request_tx, request_rx) = async_channel::unbounded();
2694+ repository_handle.update(&mut cx, |repository, cx| {
2695+ repository.search_commits(log_source, search_args, request_tx, cx);
2696+ });
2697+
2698+ let (mut response_tx, response_rx) = mpsc::unbounded();
2699+ cx.background_spawn(async move {
2700+ let mut shas = Vec::new();
2701+
2702+ while let Ok(sha) = request_rx.recv().await {
2703+ shas.push(sha.to_string());
2704+
2705+ if shas.len() >= CHUNK_SIZE {
2706+ if response_tx
2707+ .send(Ok(proto::SearchCommitsResponse {
2708+ shas: mem::take(&mut shas),
2709+ }))
2710+ .await
2711+ .is_err()
2712+ {
2713+ return;
2714+ }
2715+ }
2716+ }
2717+
2718+ if !shas.is_empty() {
2719+ response_tx
2720+ .send(Ok(proto::SearchCommitsResponse { shas }))
2721+ .await
2722+ .ok();
2723+ }
2724+ })
2725+ .detach();
2726+
2727+ Ok(response_rx)
2728+ }
2729+
26722730 async fn handle_edit_ref(
26732731 this: Entity<Self>,
26742732 envelope: TypedEnvelope<proto::GitEditRef>,
@@ -4974,6 +5032,7 @@ impl Repository { · 2299 unmodified lines
49745032 cx: &mut Context<Self>,
49755033 ) {
49765034 let repository_state = self.repository_state.clone();
5035+ let repository_id = self.id;
49775036
49785037 cx.background_spawn(async move {
49795038 let repo_state = repository_state.await;
@@ -4985,8 +5044,50 @@ impl Repository { · 5 unmodified lines
49855044 .await
49865045 .log_err();
49875046 }
4988− Ok(RepositoryState::Remote(_)) => {}
4989− Err(_) => {}
5047+
5048+ Ok(RepositoryState::Remote(RemoteRepositoryState { client, project_id })) => {
5049+ let result = client
5050+ .request_stream(proto::SearchCommits {
5051+ project_id: project_id.to_proto(),
5052+ repository_id: repository_id.to_proto(),
5053+ log_source: Some(log_source_to_proto(&log_source)),
5054+ query: search_args.query.to_string(),
5055+ case_sensitive: search_args.case_sensitive,
5056+ })
5057+ .await;
5058+
5059+ let mut stream = match result {
5060+ Ok(stream) => stream,
5061+ Err(error) => {
5062+ log::error!("failed to search commits remotely: {error:?}");
5063+ return;
5064+ }
5065+ };
5066+
5067+ while let Some(response) = stream.next().await {
5068+ let response = match response {
5069+ Ok(response) => response,
5070+ Err(error) => {
5071+ log::error!(
5072+ "failed to receive remote commit search results: {error:?}"
5073+ );
5074+ return;
5075+ }
5076+ };
5077+
5078+ for sha in &response.shas {
5079+ let Ok(oid) = Oid::from_str(sha) else {
5080+ return;
5081+ };
5082+ if request_tx.send(oid).await.is_err() {
5083+ return;
5084+ }
5085+ }
5086+ }
5087+ }
5088+ Err(error) => {
5089+ log::error!("failed to get repository state for commit search: {error}");
5090+ }
49905091 };
49915092 })
49925093 .detach();
@@ -8119,6 +8220,31 @@ fn deserialize_blame_buffer_response( · 3126 unmodified lines
81198220 Some(Blame { entries, messages })
81208221 }
81218222
8223+fn log_source_to_proto(log_source: &LogSource) -> proto::GitLogSource {
8224+ proto::GitLogSource {
8225+ source: Some(match log_source {
8226+ LogSource::All => proto::git_log_source::Source::All(proto::GitLogSourceAll {}),
8227+ LogSource::Branch(branch) => proto::git_log_source::Source::Branch(branch.to_string()),
8228+ LogSource::Sha(sha) => proto::git_log_source::Source::Sha(sha.to_string()),
8229+ LogSource::Path(path) => proto::git_log_source::Source::Path(path.to_proto()),
8230+ }),
8231+ }
8232+}
8233+
8234+fn log_source_from_proto(log_source: proto::GitLogSource) -> Result<LogSource> {
8235+ match log_source
8236+ .source
8237+ .context("git log source is missing source")?
8238+ {
8239+ proto::git_log_source::Source::All(_) => Ok(LogSource::All),
8240+ proto::git_log_source::Source::Branch(branch) => Ok(LogSource::Branch(branch.into())),
8241+ proto::git_log_source::Source::Sha(sha) => Ok(LogSource::Sha(Oid::from_str(&sha)?)),
8242+ proto::git_log_source::Source::Path(path) => {
8243+ Ok(LogSource::Path(RepoPath::from_proto(&path)?))
8244+ }
8245+ }
8246+}
8247+
81228248 fn commit_data_to_proto(commit: &CommitData) -> proto::CommitData {
81238249 proto::CommitData {
81248250 sha: commit.sha.to_string(),
crates/proto/proto/git.proto+23−0
@@ -693,3 +693,26 @@ message CommitData { · 692 unmodified lines
693693 message GetCommitDataResponse {
694694 repeated CommitData commits = 1;
695695 }
696+
697+message GitLogSourceAll {}
698+
699+message GitLogSource {
700+ oneof source {
701+ GitLogSourceAll all = 1;
702+ string branch = 2;
703+ string sha = 3;
704+ string path = 4;
705+ }
706+}
707+
708+message SearchCommits {
709+ uint64 project_id = 1;
710+ uint64 repository_id = 2;
711+ GitLogSource log_source = 3;
712+ string query = 4;
713+ bool case_sensitive = 5;
714+}
715+
716+message SearchCommitsResponse {
717+ repeated string shas = 1;
718+}
crates/proto/proto/zed.proto+3−1
@@ -480,7 +480,9 @@ message Envelope { · 479 unmodified lines
480480 GitCreateArchiveCheckpointResponse git_create_archive_checkpoint_response = 445;
481481 GitRestoreArchiveCheckpoint git_restore_archive_checkpoint = 446;
482482 GetCommitData get_commit_data = 447;
483− GetCommitDataResponse get_commit_data_response = 448; // current max
483+ GetCommitDataResponse get_commit_data_response = 448;
484+ SearchCommits search_commits = 449;
485+ SearchCommitsResponse search_commits_response = 450; // current max
484486 }
485487
486488 reserved 87 to 88;
crates/proto/src/proto.rs+4−0
@@ -358,6 +358,8 @@ messages!( · 357 unmodified lines
358358 (GitRepairWorktrees, Background),
359359 (GetCommitData, Background),
360360 (GetCommitDataResponse, Background),
361+ (SearchCommits, Background),
362+ (SearchCommitsResponse, Background),
361363 (GitWorktreesResponse, Background),
362364 (GitCreateWorktree, Background),
363365 (GitRemoveWorktree, Background),
@@ -573,6 +575,7 @@ request_messages!( · 209 unmodified lines
573575 (GitEditRef, Ack),
574576 (GitRepairWorktrees, Ack),
575577 (GetCommitData, GetCommitDataResponse),
578+ (SearchCommits, SearchCommitsResponse),
576579 (GitCreateWorktree, Ack),
577580 (GitRemoveWorktree, Ack),
578581 (GitRenameWorktree, Ack),
@@ -767,6 +770,7 @@ entity_messages!( · 188 unmodified lines
767770 GitEditRef,
768771 GitRepairWorktrees,
769772 GetCommitData,
773+ SearchCommits,
770774 GitCreateArchiveCheckpoint,
771775 GitRestoreArchiveCheckpoint,
772776 GitCreateWorktree,
crates/remote/src/remote_client.rs+225−3
@@ -22,6 +22,7 @@ use futures::{ · 21 unmodified lines
2222 },
2323 future::{BoxFuture, Shared, WeakShared},
2424 select, select_biased,
25+ stream::BoxStream,
2526 };
2627 use gpui::{
2728 App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,
@@ -1320,6 +1321,8 @@ impl RemoteConnectionOptions { · 1292 unmodified lines
13201321 #[cfg(test)]
13211322 mod tests {
13221323 use super::*;
1324+ use gpui::TestAppContext;
1325+ use rpc::{ErrorCodeExt, proto::ErrorCode};
13231326
13241327 #[test]
13251328 fn test_ssh_display_name_prefers_nickname() {
@@ -1341,6 +1344,137 @@ mod tests { · 15 unmodified lines
13411344
13421345 assert_eq!(options.display_name(), "1.2.3.4");
13431346 }
1347+
1348+ #[gpui::test]
1349+ async fn test_channel_client_request_stream_terminates_on_error(cx: &mut TestAppContext) {
1350+ let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
1351+ let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
1352+
1353+ let client =
1354+ cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
1355+
1356+ // The client sends RemoteStarted on startup; drain the outgoing channel
1357+ // so it doesn't block.
1358+ let _drain_outgoing = cx
1359+ .executor()
1360+ .spawn(async move { while outgoing_rx.next().await.is_some() {} });
1361+
1362+ let mut stream = client
1363+ .request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
1364+ .await
1365+ .unwrap();
1366+
1367+ let request_id = 0;
1368+
1369+ incoming_tx
1370+ .unbounded_send(proto::Test { id: 1 }.into_envelope(100, Some(request_id), None))
1371+ .unwrap();
1372+
1373+ let first = stream.next().await.unwrap().unwrap();
1374+ assert_eq!(
1375+ proto::Test::from_envelope(first).unwrap(),
1376+ proto::Test { id: 1 }
1377+ );
1378+
1379+ // Send an Error without a trailing EndStream. The Error alone should
1380+ // terminate the stream.
1381+ incoming_tx
1382+ .unbounded_send(
1383+ ErrorCode::Internal
1384+ .message("boom".to_string())
1385+ .to_proto()
1386+ .into_envelope(101, Some(request_id), None),
1387+ )
1388+ .unwrap();
1389+
1390+ let second = stream.next().await.unwrap();
1391+ let error = second.unwrap_err();
1392+ assert!(
1393+ format!("{error}").contains("boom"),
1394+ "expected error to surface server message, got: {error}"
1395+ );
1396+
1397+ assert!(stream.next().await.is_none());
1398+ assert_eq!(client.stream_response_channels.lock().len(), 0);
1399+ }
1400+
1401+ #[gpui::test]
1402+ async fn test_channel_client_dropping_stream_request_before_response_cleans_up_channel(
1403+ cx: &mut TestAppContext,
1404+ ) {
1405+ let (_incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
1406+ let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
1407+
1408+ let client =
1409+ cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
1410+
1411+ let _drain_outgoing = cx
1412+ .executor()
1413+ .spawn(async move { while outgoing_rx.next().await.is_some() {} });
1414+
1415+ let stream = client
1416+ .request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
1417+ .await
1418+ .unwrap();
1419+
1420+ assert_eq!(client.stream_response_channels.lock().len(), 1);
1421+
1422+ drop(stream);
1423+ cx.run_until_parked();
1424+
1425+ assert_eq!(
1426+ client.stream_response_channels.lock().len(),
1427+ 0,
1428+ "dropping a stream before any responses arrive should remove response channel bookkeeping"
1429+ );
1430+ }
1431+
1432+ #[gpui::test]
1433+ async fn test_channel_client_dropping_stream_request_before_completion(
1434+ cx: &mut TestAppContext,
1435+ ) {
1436+ let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
1437+ let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();
1438+
1439+ let client =
1440+ cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "test-client", false));
1441+
1442+ let _drain_outgoing = cx
1443+ .executor()
1444+ .spawn(async move { while outgoing_rx.next().await.is_some() {} });
1445+
1446+ let mut stream = client
1447+ .request_stream_dynamic(proto::Test { id: 0 }.into_envelope(0, None, None), "Test")
1448+ .await
1449+ .unwrap();
1450+
1451+ let request_id = 0;
1452+
1453+ incoming_tx
1454+ .unbounded_send(proto::Test { id: 1 }.into_envelope(100, Some(request_id), None))
1455+ .unwrap();
1456+ let _ = stream.next().await.unwrap().unwrap();
1457+
1458+ assert_eq!(client.stream_response_channels.lock().len(), 1);
1459+
1460+ drop(stream);
1461+
1462+ // Inject an orphaned non-terminal response. The read loop should detect
1463+ // that the consumer has been dropped and clean up its bookkeeping (no
1464+ // EndStream sent here on purpose, otherwise the cleanup would happen
1465+ // via the terminal-response path and mask the bug under test).
1466+ incoming_tx
1467+ .unbounded_send(proto::Test { id: 2 }.into_envelope(101, Some(request_id), None))
1468+ .unwrap();
1469+
1470+ cx.run_until_parked();
1471+
1472+ assert_eq!(
1473+ client.stream_response_channels.lock().len(),
1474+ 0,
1475+ "stream channel should be removed once the consumer has dropped the stream"
1476+ );
1477+ }
13441478 }
13451479
13461480 impl From<SshConnectionOptions> for RemoteConnectionOptions {
@@ -1418,6 +1552,8 @@ pub trait RemoteConnection: Send + Sync { · 71 unmodified lines
14181552 }
14191553
14201554 type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
1555+type StreamResponseChannels =
1556+ Arc<Mutex<HashMap<MessageId, UnboundedSender<(Result<Envelope>, oneshot::Sender<()>)>>>>;
14211557
14221558 struct Signal<T> {
14231559 tx: Mutex<Option<oneshot::Sender<T>>>,
@@ -1455,6 +1591,7 @@ pub(crate) struct ChannelClient { · 31 unmodified lines
14551591 outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,
14561592 buffer: Mutex<VecDeque<Envelope>>,
14571593 response_channels: ResponseChannels,
1594+ stream_response_channels: StreamResponseChannels,
14581595 message_handlers: Mutex<ProtoMessageHandlerSet>,
14591596 max_received: AtomicU32,
14601597 name: &'static str,
@@ -1477,6 +1614,7 @@ impl ChannelClient { · 16 unmodified lines
14771614 next_message_id: AtomicU32::new(0),
14781615 max_received: AtomicU32::new(0),
14791616 response_channels: ResponseChannels::default(),
1617+ stream_response_channels: StreamResponseChannels::default(),
14801618 message_handlers: Default::default(),
14811619 buffer: Mutex::new(VecDeque::new()),
14821620 name,
@@ -1550,13 +1688,40 @@ impl ChannelClient { · 67 unmodified lines
15501688
15511689 if let Some(request_id) = incoming.responding_to {
15521690 let request_id = MessageId(request_id);
1691+ // An incoming response with no payload is malformed; drop
1692+ // it. The request future and any stream consumers will
1693+ // remain pending until either a real response arrives or
1694+ // the connection is torn down.
1695+ if incoming.payload.is_none() {
1696+ continue;
1697+ }
15531698 let sender = this.response_channels.lock().remove(&request_id);
15541699 if let Some(sender) = sender {
15551700 let (tx, rx) = oneshot::channel();
1556− if incoming.payload.is_some() {
1557− sender.send((incoming, tx)).ok();
1558− }
1701+ sender.send((incoming, tx)).ok();
15591702 rx.await.ok();
1703+ } else {
1704+ let terminal_stream_response = matches!(
1705+ &incoming.payload,
1706+ Some(proto::envelope::Payload::Error(_))
1707+ | Some(proto::envelope::Payload::EndStream(_))
1708+ );
1709+ let sender = if terminal_stream_response {
1710+ this.stream_response_channels.lock().remove(&request_id)
1711+ } else {
1712+ this.stream_response_channels
1713+ .lock()
1714+ .get(&request_id)
1715+ .cloned()
1716+ };
1717+ if let Some(sender) = sender {
1718+ let (tx, rx) = oneshot::channel();
1719+ if sender.unbounded_send((Ok(incoming), tx)).is_err() {
1720+ this.stream_response_channels.lock().remove(&request_id);
1721+ continue;
1722+ }
1723+ rx.await.ok();
1724+ }
15601725 }
15611726 } else if let Some(envelope) =
15621727 build_typed_envelope(peer_id, Instant::now(), incoming)
@@ -1721,6 +1886,55 @@ impl ChannelClient { · 158 unmodified lines
17211886 }
17221887 }
17231888
1889+ fn request_stream_dynamic(
1890+ &self,
1891+ mut envelope: proto::Envelope,
1892+ type_name: &'static str,
1893+ ) -> impl 'static + Future<Output = Result<BoxStream<'static, Result<proto::Envelope>>>> {
1894+ envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1895+ let message_id = MessageId(envelope.id);
1896+ let (tx, rx) = mpsc::unbounded();
1897+ let stream_response_channels = self.stream_response_channels.clone();
1898+ stream_response_channels.lock().insert(message_id, tx);
1899+
1900+ let result = self.send_buffered(envelope);
1901+ async move {
1902+ if let Err(error) = &result {
1903+ log::error!("failed to send message: {error}");
1904+ anyhow::bail!("failed to send message: {error}");
1905+ }
1906+
1907+ let cleanup_stream_response_channel = util::defer({
1908+ let stream_response_channels = stream_response_channels.clone();
1909+ move || {
1910+ stream_response_channels.lock().remove(&message_id);
1911+ }
1912+ });
1913+
1914+ Ok(rx
1915+ .filter_map(move |(response, _barrier)| {
1916+ // Keep the cleanup guard alive until the returned stream is dropped.
1917+ let _keep_cleanup_guard_alive = &cleanup_stream_response_channel;
1918+ futures::future::ready(match response {
1919+ Ok(response) => {
1920+ if let Some(proto::envelope::Payload::Error(error)) = &response.payload
1921+ {
1922+ Some(Err(RpcError::from_proto(error, type_name)))
1923+ } else if let Some(proto::envelope::Payload::EndStream(_)) =
1924+ &response.payload
1925+ {
1926+ None
1927+ } else {
1928+ Some(Ok(response))
1929+ }
1930+ }
1931+ Err(error) => Some(Err(error)),
1932+ })
1933+ })
1934+ .boxed())
1935+ }
1936+ }
1937+
17241938 pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
17251939 envelope.id = self.next_message_id.fetch_add(1, SeqCst);
17261940 self.send_buffered(envelope)
@@ -1751,6 +1965,14 @@ impl ProtoClient for ChannelClient { · 24 unmodified lines
17511965 self.request_dynamic(envelope, request_type, true).boxed()
17521966 }
17531967
1968+ fn request_stream(
1969+ &self,
1970+ envelope: proto::Envelope,
1971+ request_type: &'static str,
1972+ ) -> BoxFuture<'static, Result<BoxStream<'static, Result<proto::Envelope>>>> {
1973+ self.request_stream_dynamic(envelope, request_type).boxed()
1974+ }
1975+
17541976 fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
17551977 self.send_dynamic(envelope)
17561978 }
crates/remote_server/src/remote_editing_tests.rs+108−1
@@ -11,7 +11,10 @@ use languages::rust_lang; · 10 unmodified lines
1111
1212 use extension::ExtensionHostProxy;
1313 use fs::{FakeFs, Fs};
14−use git::repository::Worktree as GitWorktree;
14+use git::{
15+ Oid,
16+ repository::{CommitData, Worktree as GitWorktree},
17+};
1518 use gpui::{AppContext as _, Entity, SharedString, TestAppContext};
1619 use http_client::{BlockedHttpClient, FakeHttpClient};
1720 use language::{
@@ -29,11 +32,13 @@ use project::{ · 11 unmodified lines
2932 search::{SearchQuery, SearchResult},
3033 };
3134 use remote::RemoteClient;
35+use rpc::proto;
3236 use serde_json::json;
3337 use settings::{Settings, SettingsLocation, SettingsStore, initial_server_settings_content};
3438 use smol::stream::StreamExt;
3539 use std::{
3640 path::{Path, PathBuf},
41+ str::FromStr,
3742 sync::Arc,
3843 };
3944 use unindent::Unindent as _;
@@ -1626,6 +1631,108 @@ async fn test_remote_root_repo_common_dir(cx: &mut TestAppContext, server_cx: &m · 1586 unmodified lines
16261631 assert_eq!(common_dir, None);
16271632 }
16281633
1634+#[gpui::test]
1635+async fn test_remote_search_commits_streams_proto_chunks(
1636+ cx: &mut TestAppContext,
1637+ server_cx: &mut TestAppContext,
1638+) {
1639+ const COMMIT_COUNT: usize = 900;
1640+ const RESPONSE_MAX_SIZE: usize = 100;
1641+
1642+ let fs = FakeFs::new(server_cx.executor());
1643+ fs.insert_tree(
1644+ path!("/code"),
1645+ json!({
1646+ "project1": {
1647+ ".git": {},
1648+ "file.txt": "content",
1649+ },
1650+ }),
1651+ )
1652+ .await;
1653+
1654+ let commit_data = (0..COMMIT_COUNT)
1655+ .map(|index| {
1656+ let sha = Oid::from_str(&format!("{:040x}", index + 1)).unwrap();
1657+ (
1658+ CommitData {
1659+ sha,
1660+ parents: Default::default(),
1661+ author_name: SharedString::from("Author"),
1662+ author_email: SharedString::from("author@example.com"),
1663+ commit_timestamp: index as i64,
1664+ subject: SharedString::from(format!("Subject {index}")),
1665+ message: SharedString::from(format!("needle commit {index}")),
1666+ },
1667+ false,
1668+ )
1669+ })
1670+ .collect::<Vec<_>>();
1671+ let expected_shas = commit_data
1672+ .iter()
1673+ .map(|(commit_data, _)| commit_data.sha.to_string())
1674+ .collect::<HashSet<_>>();
1675+ fs.set_commit_data(Path::new(path!("/code/project1/.git")), commit_data);
1676+
1677+ let (project, _headless) = init_test(&fs, cx, server_cx).await;
1678+ project
1679+ .update(cx, |project, cx| {
1680+ project.find_or_create_worktree(path!("/code/project1"), true, cx)
1681+ })
1682+ .await
1683+ .expect("should open remote worktree");
1684+ server_cx.run_until_parked();
1685+ cx.run_until_parked();
1686+ project
1687+ .update(cx, |project, cx| project.git_scans_complete(cx))
1688+ .await;
1689+
1690+ let (remote_client, repository_id) = project.read_with(cx, |project, cx| {
1691+ let repository = project
1692+ .active_repository(cx)
1693+ .expect("remote project should have an active repository");
1694+ let repository_id = repository.read(cx).snapshot().id;
1695+ let remote_client = project
1696+ .remote_client()
1697+ .expect("project should have a remote client");
1698+ (remote_client, repository_id)
1699+ });
1700+ let proto_client = remote_client.read_with(cx, |remote_client, _| remote_client.proto_client());
1701+ let mut stream = proto_client
1702+ .request_stream(proto::SearchCommits {
1703+ project_id: proto::REMOTE_SERVER_PROJECT_ID,
1704+ repository_id: repository_id.to_proto(),
1705+ log_source: Some(proto::GitLogSource {
1706+ source: Some(proto::git_log_source::Source::All(
1707+ proto::GitLogSourceAll {},
1708+ )),
1709+ }),
1710+ query: "needle".to_string(),
1711+ case_sensitive: true,
1712+ })
1713+ .await
1714+ .expect("search commits stream should start");
1715+
1716+ let mut chunks = Vec::new();
1717+ while let Some(response) = futures::StreamExt::next(&mut stream).await {
1718+ chunks.push(response.expect("search commits chunk should succeed").shas);
1719+ }
1720+
1721+ assert!(
1722+ chunks.len() > 1,
1723+ "expected search results to stream in multiple chunks"
1724+ );
1725+ for chunk in chunks.iter().take(chunks.len() - 1) {
1726+ assert!(
1727+ chunk.len() <= RESPONSE_MAX_SIZE,
1728+ "non-final chunks should meet the target byte size"
1729+ );
1730+ }
1731+
1732+ let actual_shas = chunks.into_iter().flatten().collect::<HashSet<_>>();
1733+ assert_eq!(actual_shas, expected_shas);
1734+}
1735+
16291736 #[gpui::test]
16301737 async fn test_remote_archive_git_operations_are_supported(
16311738 cx: &mut TestAppContext,
crates/rpc/src/peer.rs+364−48
@@ -8,7 +8,7 @@ use super::{ · 7 unmodified lines
88 use anyhow::{Context as _, Result, anyhow};
99 use collections::HashMap;
1010 use futures::{
11− FutureExt, SinkExt, Stream, StreamExt, TryFutureExt,
11+ FutureExt, SinkExt, StreamExt, TryFutureExt,
1212 channel::{mpsc, oneshot},
1313 stream::BoxStream,
1414 };
@@ -278,11 +278,23 @@ impl Peer { · 263 unmodified lines
278278 );
279279 let response_channel =
280280 response_channels.lock().as_mut()?.remove(&responding_to);
281− let stream_response_channel = stream_response_channels
282− .lock()
283− .as_ref()?
284− .get(&responding_to)
285− .cloned();
281+ let terminal_stream_response = matches!(
282+ &incoming.payload,
283+ Some(proto::envelope::Payload::Error(_))
284+ | Some(proto::envelope::Payload::EndStream(_))
285+ );
286+ let stream_response_channel = if terminal_stream_response {
287+ stream_response_channels
288+ .lock()
289+ .as_mut()?
290+ .remove(&responding_to)
291+ } else {
292+ stream_response_channels
293+ .lock()
294+ .as_ref()?
295+ .get(&responding_to)
296+ .cloned()
297+ };
286298
287299 if let Some(tx) = response_channel {
288300 let requester_resumed = oneshot::channel();
@@ -319,21 +331,15 @@ impl Peer { · 30 unmodified lines
319331 ?error,
320332 "incoming stream response: request future dropped",
321333 );
334+ // The consumer has gone away, so drop the bookkeeping
335+ // for this stream rather than letting it accumulate
336+ // every subsequent message until a terminal frame.
337+ if let Some(channels) = stream_response_channels.lock().as_mut() {
338+ channels.remove(&responding_to);
339+ }
340+ } else {
341+ let _ = requester_resumed.1.await;
322342 }
323−
324− tracing::debug!(
325− %connection_id,
326− message_id,
327− responding_to,
328− "incoming stream response: waiting to resume requester"
329− );
330− let _ = requester_resumed.1.await;
331− tracing::debug!(
332− %connection_id,
333− message_id,
334− responding_to,
335− "incoming stream response: requester resumed"
336− );
337343 } else {
338344 let message_type = proto::build_typed_envelope(
339345 connection_id.into(),
@@ -484,55 +490,96 @@ impl Peer { · 144 unmodified lines
484490 &self,
485491 receiver_id: ConnectionId,
486492 request: T,
487− ) -> impl Future<Output = Result<impl Unpin + Stream<Item = Result<T::Response>>>> {
493+ ) -> impl Future<Output = Result<BoxStream<'static, Result<T::Response>>>> {
494+ let stream =
495+ self.request_stream_dynamic(receiver_id, request.into_envelope(0, None, None), T::NAME);
496+
497+ async move {
498+ Ok(stream
499+ .await?
500+ .map(|response| {
501+ T::Response::from_envelope(response?)
502+ .context("received response of the wrong type")
503+ })
504+ .boxed())
505+ }
506+ }
507+
508+ pub fn request_stream_dynamic(
509+ &self,
510+ receiver_id: ConnectionId,
511+ mut envelope: proto::Envelope,
512+ request_type: &'static str,
513+ ) -> impl Future<Output = Result<BoxStream<'static, Result<proto::Envelope>>>> + use<> {
488514 let (tx, rx) = mpsc::unbounded();
489515 let send = self.connection_state(receiver_id).and_then(|connection| {
490516 let message_id = connection.next_message_id.fetch_add(1, SeqCst);
517+ envelope.id = message_id;
491518 let stream_response_channels = connection.stream_response_channels.clone();
492519 stream_response_channels
493520 .lock()
494521 .as_mut()
495522 .context("connection was closed")?
496523 .insert(message_id, tx);
497− connection
524+ if let Err(error) = connection
498525 .outgoing_tx
499− .unbounded_send(Message::Envelope(
500− request.into_envelope(message_id, None, None),
501− ))
502− .context("connection was closed")?;
526+ .unbounded_send(Message::Envelope(envelope))
527+ {
528+ if let Some(channels) = stream_response_channels.lock().as_mut() {
529+ channels.remove(&message_id);
530+ }
531+ return Err(error).context("connection was closed");
532+ }
503533 Ok((message_id, stream_response_channels))
504534 });
505535
506536 async move {
507537 let (message_id, stream_response_channels) = send?;
508538 let stream_response_channels = Arc::downgrade(&stream_response_channels);
509−
510− Ok(rx.filter_map(move |(response, _barrier)| {
539+ let cleanup_stream_response_channel = util::defer({
511540 let stream_response_channels = stream_response_channels.clone();
512− future::ready(match response {
513− Ok(response) => {
514− if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
515− Some(Err(RpcError::from_proto(error, T::NAME)))
516− } else if let Some(proto::envelope::Payload::EndStream(_)) =
517− &response.payload
518− {
519− // Remove the transmitting end of the response channel to end the stream.
520− if let Some(channels) = stream_response_channels.upgrade()
521− && let Some(channels) = channels.lock().as_mut()
541+ move || {
542+ if let Some(channels) = stream_response_channels.upgrade()
543+ && let Some(channels) = channels.lock().as_mut()
544+ {
545+ channels.remove(&message_id);
546+ }
547+ }
548+ });
549+
550+ Ok(rx
551+ .filter_map(move |(response, _barrier)| {
552+ let _keep_cleanup_guard_alive = &cleanup_stream_response_channel;
553+ let stream_response_channels = stream_response_channels.clone();
554+ future::ready(match response {
555+ Ok(response) => {
556+ if let Some(proto::envelope::Payload::Error(error)) = &response.payload
522557 {
523− channels.remove(&message_id);
558+ // Remove the transmitting end of the response channel to end the stream.
559+ if let Some(channels) = stream_response_channels.upgrade()
560+ && let Some(channels) = channels.lock().as_mut()
561+ {
562+ channels.remove(&message_id);
563+ }
564+ Some(Err(RpcError::from_proto(error, request_type)))
565+ } else if let Some(proto::envelope::Payload::EndStream(_)) =
566+ &response.payload
567+ {
568+ // Remove the transmitting end of the response channel to end the stream.
569+ if let Some(channels) = stream_response_channels.upgrade()
570+ && let Some(channels) = channels.lock().as_mut()
571+ {
572+ channels.remove(&message_id);
573+ }
574+ None
575+ } else {
576+ Some(Ok(response))
524577 }
525− None
526− } else {
527− Some(
528− T::Response::from_envelope(response)
529− .context("received response of the wrong type"),
530− )
531578 }
532− }
533− Err(error) => Some(Err(error)),
579+ Err(error) => Some(Err(error)),
580+ })
534581 })
535− }))
582+ .boxed())
536583 }
537584 }
538585
@@ -661,6 +708,13 @@ impl Peer { · 122 unmodified lines
661708 .with_context(|| format!("no such connection: {connection_id}"))?;
662709 Ok(connection.clone())
663710 }
711+
712+ #[cfg(any(test, feature = "test-support"))]
713+ pub fn pending_stream_request_count(&self, connection_id: ConnectionId) -> Option<usize> {
714+ let connection = self.connection_state(connection_id).ok()?;
715+ let channels = connection.stream_response_channels.lock();
716+ Some(channels.as_ref()?.len())
717+ }
664718 }
665719
666720 impl Serialize for Peer {
@@ -992,6 +1046,268 @@ mod tests { · 325 unmodified lines
9921046 );
9931047 }
9941048
1049+ #[gpui::test(iterations = 50)]
1050+ async fn test_request_stream(cx: &mut TestAppContext) {
1051+ init_logger();
1052+
1053+ let executor = cx.executor();
1054+ let server = Peer::new(0);
1055+ let client = Peer::new(0);
1056+
1057+ let (client_to_server_conn, server_to_client_conn, _kill) =
1058+ Connection::in_memory(executor.clone());
1059+ let (client_to_server_conn_id, io_task1, mut client_incoming) =
1060+ client.add_test_connection(client_to_server_conn, executor.clone());
1061+ let (_, io_task2, mut server_incoming) =
1062+ server.add_test_connection(server_to_client_conn, executor.clone());
1063+
1064+ executor.spawn(io_task1).detach();
1065+ executor.spawn(io_task2).detach();
1066+ executor
1067+ .spawn(async move { while client_incoming.next().await.is_some() {} })
1068+ .detach();
1069+
1070+ executor
1071+ .spawn({
1072+ let server = server.clone();
1073+ async move {
1074+ let request = server_incoming
1075+ .next()
1076+ .await
1077+ .unwrap()
1078+ .into_any()
1079+ .downcast::<TypedEnvelope<proto::Test>>()
1080+ .unwrap();
1081+ let receipt = request.receipt();
1082+ server.respond(receipt, proto::Test { id: 1 }).unwrap();
1083+ server.respond(receipt, proto::Test { id: 2 }).unwrap();
1084+ server.respond(receipt, proto::Test { id: 3 }).unwrap();
1085+ server.end_stream(receipt).unwrap();
1086+
1087+ // Prevent the connection from being dropped.
1088+ server_incoming.next().await;
1089+ }
1090+ })
1091+ .detach();
1092+
1093+ let mut stream = client
1094+ .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1095+ .await
1096+ .unwrap();
1097+
1098+ let mut received = Vec::new();
1099+ while let Some(item) = stream.next().await {
1100+ received.push(item.unwrap());
1101+ }
1102+
1103+ assert_eq!(
1104+ received,
1105+ vec![
1106+ proto::Test { id: 1 },
1107+ proto::Test { id: 2 },
1108+ proto::Test { id: 3 },
1109+ ]
1110+ );
1111+ assert_eq!(
1112+ client.pending_stream_request_count(client_to_server_conn_id),
1113+ Some(0)
1114+ );
1115+ }
1116+
1117+ #[gpui::test]
1118+ async fn test_request_stream_send_failure_cleans_up_response_channel(cx: &mut TestAppContext) {
1119+ init_logger();
1120+
1121+ let executor = cx.executor();
1122+ let client = Peer::new(0);
1123+
1124+ let (client_to_server_conn, _server_to_client_conn, _kill) =
1125+ Connection::in_memory(executor.clone());
1126+ let (client_to_server_conn_id, io_task, _client_incoming) =
1127+ client.add_test_connection(client_to_server_conn, executor.clone());
1128+
1129+ drop(io_task);
1130+
1131+ let result = client
1132+ .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1133+ .await;
1134+
1135+ assert!(
1136+ result.is_err(),
1137+ "stream request should fail when the connection write task has gone away"
1138+ );
1139+ assert_eq!(
1140+ client.pending_stream_request_count(client_to_server_conn_id),
1141+ Some(0),
1142+ "failed stream request should not leave response channel bookkeeping behind"
1143+ );
1144+ }
1145+
1146+ #[gpui::test(iterations = 50)]
1147+ async fn test_request_stream_terminates_on_error(cx: &mut TestAppContext) {
1148+ init_logger();
1149+
1150+ let executor = cx.executor();
1151+ let server = Peer::new(0);
1152+ let client = Peer::new(0);
1153+
1154+ let (client_to_server_conn, server_to_client_conn, _kill) =
1155+ Connection::in_memory(executor.clone());
1156+ let (client_to_server_conn_id, io_task1, mut client_incoming) =
1157+ client.add_test_connection(client_to_server_conn, executor.clone());
1158+ let (_, io_task2, mut server_incoming) =
1159+ server.add_test_connection(server_to_client_conn, executor.clone());
1160+
1161+ executor.spawn(io_task1).detach();
1162+ executor.spawn(io_task2).detach();
1163+ executor
1164+ .spawn(async move { while client_incoming.next().await.is_some() {} })
1165+ .detach();
1166+
1167+ executor
1168+ .spawn({
1169+ let server = server.clone();
1170+ async move {
1171+ let request = server_incoming
1172+ .next()
1173+ .await
1174+ .unwrap()
1175+ .into_any()
1176+ .downcast::<TypedEnvelope<proto::Test>>()
1177+ .unwrap();
1178+ let receipt = request.receipt();
1179+ server.respond(receipt, proto::Test { id: 1 }).unwrap();
1180+ // Send an Error without a trailing EndStream. The Error alone
1181+ // should be treated as a terminal stream response.
1182+ server
1183+ .respond_with_error(
1184+ receipt,
1185+ ErrorCode::Internal.message("boom".to_string()).to_proto(),
1186+ )
1187+ .unwrap();
1188+
1189+ // Prevent the connection from being dropped.
1190+ server_incoming.next().await;
1191+ }
1192+ })
1193+ .detach();
1194+
1195+ let mut stream = client
1196+ .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1197+ .await
1198+ .unwrap();
1199+
1200+ assert_eq!(stream.next().await.unwrap().unwrap(), proto::Test { id: 1 });
1201+
1202+ let error = stream.next().await.unwrap().unwrap_err();
1203+ assert!(
1204+ format!("{error}").contains("boom"),
1205+ "expected error to surface server message, got: {error}"
1206+ );
1207+
1208+ // The error alone (without an EndStream) should terminate the stream.
1209+ assert!(stream.next().await.is_none());
1210+ assert_eq!(
1211+ client.pending_stream_request_count(client_to_server_conn_id),
1212+ Some(0)
1213+ );
1214+ }
1215+
1216+ #[gpui::test(iterations = 50)]
1217+ async fn test_dropping_stream_request_before_completion(cx: &mut TestAppContext) {
1218+ init_logger();
1219+
1220+ let executor = cx.executor();
1221+ let server = Peer::new(0);
1222+ let client = Peer::new(0);
1223+
1224+ let (client_to_server_conn, server_to_client_conn, _kill) =
1225+ Connection::in_memory(executor.clone());
1226+ let (client_to_server_conn_id, io_task1, mut client_incoming) =
1227+ client.add_test_connection(client_to_server_conn, executor.clone());
1228+ let (_, io_task2, mut server_incoming) =
1229+ server.add_test_connection(server_to_client_conn, executor.clone());
1230+
1231+ executor.spawn(io_task1).detach();
1232+ executor.spawn(io_task2).detach();
1233+ executor
1234+ .spawn(async move { while client_incoming.next().await.is_some() {} })
1235+ .detach();
1236+
1237+ let (drop_signal_tx, drop_signal_rx) = oneshot::channel::<()>();
1238+ let server_task = executor.spawn({
1239+ let server = server.clone();
1240+ async move {
1241+ let request = server_incoming
1242+ .next()
1243+ .await
1244+ .unwrap()
1245+ .into_any()
1246+ .downcast::<TypedEnvelope<proto::Test>>()
1247+ .unwrap();
1248+ let receipt = request.receipt();
1249+ server.respond(receipt, proto::Test { id: 1 }).unwrap();
1250+
1251+ // Wait until the consumer has dropped the stream.
1252+ drop_signal_rx.await.ok();
1253+
1254+ // Send a non-terminal response after the consumer is gone. The
1255+ // peer should detect that the receiver has been dropped and clean
1256+ // up its bookkeeping. Crucially, we do NOT send EndStream here
1257+ // because that would clean up via the terminal-response path and
1258+ // mask the bug.
1259+ server.respond(receipt, proto::Test { id: 2 }).unwrap();
1260+
1261+ // A Ping/Ack round-trip after the response acts as a sync
1262+ // barrier: because messages over the in-memory connection are
1263+ // delivered in order, by the time the client observes the Ack,
1264+ // it has already processed the dropped response above.
1265+ let ping = server_incoming
1266+ .next()
1267+ .await
1268+ .unwrap()
1269+ .into_any()
1270+ .downcast::<TypedEnvelope<proto::Ping>>()
1271+ .unwrap();
1272+ server.respond(ping.receipt(), proto::Ack {}).unwrap();
1273+
1274+ // Prevent the connection from being dropped.
1275+ server_incoming.next().await;
1276+ }
1277+ });
1278+
1279+ let mut stream = client
1280+ .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1281+ .await
1282+ .unwrap();
1283+
1284+ assert_eq!(stream.next().await.unwrap().unwrap(), proto::Test { id: 1 });
1285+
1286+ // The stream is mid-flight, so the channel should be tracked.
1287+ assert_eq!(
1288+ client.pending_stream_request_count(client_to_server_conn_id),
1289+ Some(1)
1290+ );
1291+
1292+ drop(stream);
1293+ drop_signal_tx.send(()).ok();
1294+
1295+ // Synchronization barrier: once this Ack arrives, the read loop has
1296+ // already processed the orphaned stream response that came before it.
1297+ client
1298+ .request(client_to_server_conn_id, proto::Ping {})
1299+ .await
1300+ .unwrap();
1301+
1302+ assert_eq!(
1303+ client.pending_stream_request_count(client_to_server_conn_id),
1304+ Some(0),
1305+ "stream channel should be removed once the consumer has dropped the stream"
1306+ );
1307+
1308+ drop(server_task);
1309+ }
1310+
9951311 #[gpui::test(iterations = 50)]
9961312 async fn test_disconnect(cx: &mut TestAppContext) {
9971313 let executor = cx.executor();
crates/rpc/src/proto_client.rs+95−1
@@ -1,9 +1,10 @@
11 use anyhow::{Context, Result};
22 use collections::HashMap;
33 use futures::{
4− Future, FutureExt as _,
4+ Future, FutureExt as _, Stream, StreamExt as _,
55 channel::oneshot,
66 future::{BoxFuture, LocalBoxFuture},
7+ stream::BoxStream,
78 };
89 use gpui::{AnyEntity, AnyWeakEntity, AsyncApp, BackgroundExecutor, Entity, FutureExt as _};
910 use parking_lot::Mutex;
@@ -61,6 +62,20 @@ pub trait ProtoClient: Send + Sync { · 51 unmodified lines
6162 request_type: &'static str,
6263 ) -> BoxFuture<'static, Result<Envelope>>;
6364
65+ fn request_stream(
66+ &self,
67+ envelope: Envelope,
68+ request_type: &'static str,
69+ ) -> BoxFuture<'static, Result<BoxStream<'static, Result<Envelope>>>> {
70+ async move {
71+ anyhow::bail!(
72+ "stream requests are not supported for {request_type}: {:?}",
73+ envelope.payload
74+ )
75+ }
76+ .boxed()
77+ }
78+
6479 fn send(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;
6580
6681 fn send_response(&self, envelope: Envelope, message_type: &'static str) -> Result<()>;
@@ -223,6 +238,23 @@ impl AnyProtoClient { · 156 unmodified lines
223238 }
224239 }
225240
241+ pub fn request_stream<T: RequestMessage>(
242+ &self,
243+ request: T,
244+ ) -> impl Future<Output = Result<BoxStream<'static, Result<T::Response>>>> + use<T> {
245+ let envelope = request.into_envelope(0, None, None);
246+ let response_stream = self.0.client.request_stream(envelope, T::NAME);
247+ async move {
248+ Ok(response_stream
249+ .await?
250+ .map(|response| {
251+ T::Response::from_envelope(response?)
252+ .context("received response of the wrong type")
253+ })
254+ .boxed())
255+ }
256+ }
257+
226258 pub fn send<T: EnvelopedMessage>(&self, request: T) -> Result<()> {
227259 let envelope = request.into_envelope(0, None, None);
228260 self.0.client.send(envelope, T::NAME)
@@ -479,6 +511,68 @@ impl AnyProtoClient { · 250 unmodified lines
479511 );
480512 }
481513
514+ pub fn add_entity_stream_request_handler<M, E, H, F, S>(&self, handler: H)
515+ where
516+ M: EnvelopedMessage + RequestMessage + EntityMessage,
517+ E: 'static,
518+ H: 'static + Sync + Send + Fn(gpui::Entity<E>, TypedEnvelope<M>, AsyncApp) -> F,
519+ F: 'static + Future<Output = Result<S>>,
520+ S: 'static + Stream<Item = Result<M::Response>>,
521+ {
522+ let message_type_id = TypeId::of::<M>();
523+ let entity_type_id = TypeId::of::<E>();
524+ let entity_id_extractor = |envelope: &dyn AnyTypedEnvelope| {
525+ (envelope as &dyn Any)
526+ .downcast_ref::<TypedEnvelope<M>>()
527+ .unwrap()
528+ .payload
529+ .remote_entity_id()
530+ };
531+ self.0
532+ .client
533+ .message_handler_set()
534+ .lock()
535+ .add_entity_message_handler(
536+ message_type_id,
537+ entity_type_id,
538+ entity_id_extractor,
539+ Arc::new(move |entity, envelope, client, cx| {
540+ let entity = entity.downcast::<E>().unwrap();
541+ let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
542+ let request_id = envelope.message_id();
543+ let stream = handler(entity, *envelope, cx);
544+ async move {
545+ // An Error response is itself a terminal stream frame on
546+ // both transports (Peer and ChannelClient), so we don't
547+ // need to follow it with an EndStream.
548+ match stream.await {
549+ Ok(stream) => {
550+ futures::pin_mut!(stream);
551+ while let Some(result) = stream.next().await {
552+ match result {
553+ Ok(response) => {
554+ client.send_response(request_id, response)?
555+ }
556+ Err(error) => {
557+ client.send_response(request_id, error.to_proto())?;
558+ return Err(error);
559+ }
560+ }
561+ }
562+ client.send_response(request_id, proto::EndStream {})?;
563+ Ok(())
564+ }
565+ Err(error) => {
566+ client.send_response(request_id, error.to_proto())?;
567+ Err(error)
568+ }
569+ }
570+ }
571+ .boxed_local()
572+ }),
573+ );
574+ }
575+
482576 pub fn add_entity_message_handler<M, E, H, F>(&self, handler: H)
483577 where
484578 M: EnvelopedMessage + EntityMessage,