1use std::{sync::Arc, time::Duration};
2
3use anyhow::anyhow;
4use askpass::EncryptedPassword;
5use clap::Parser;
6use client::{Client, UserStore};
7use futures::channel::oneshot;
8use gpui::{AppContext as _, Application};
9use http_client::FakeHttpClient;
10use language::LanguageRegistry;
11use node_runtime::NodeRuntime;
12use project::{
13 Project, RealFs,
14 search::{SearchQuery, SearchResult},
15};
16use release_channel::ReleaseChannel;
17use remote::{ConnectionIdentifier, RemoteClientDelegate, SshConnectionOptions};
18use semver::Version;
19
20#[derive(Parser)]
21struct Args {
22 /// List of worktrees to run the search against.
23 worktrees: Vec<String>,
24 #[clap(short)]
25 query: Option<String>,
26 /// Askpass socket for SSH authentication
27 #[clap(long)]
28 askpass: Option<String>,
29 /// Treat query as a regex.
30 #[clap(short, long)]
31 regex: bool,
32 /// Matches have to be standalone words.
33 #[clap(long)]
34 whole_word: bool,
35 /// Make matching case-sensitive.
36 #[clap(long, default_value_t = false)]
37 case_sensitive: bool,
38 /// Include gitignored files in the search.
39 #[clap(long)]
40 include_ignored: bool,
41 #[clap(long)]
42 ssh: Option<String>,
43}
44
45struct BenchmarkRemoteClient;
46impl RemoteClientDelegate for BenchmarkRemoteClient {
47 fn ask_password(
48 &self,
49 prompt: String,
50 tx: oneshot::Sender<EncryptedPassword>,
51 _cx: &mut gpui::AsyncApp,
52 ) {
53 eprintln!("SSH asking for password: {}", prompt);
54 match rpassword::prompt_password(&prompt) {
55 Ok(password) => match EncryptedPassword::try_from(password.as_ref()) {
56 Ok(encrypted) => {
57 if tx.send(encrypted).is_err() {
58 eprintln!("Failed to send password");
59 }
60 }
61 Err(e) => eprintln!("Failed to encrypt password: {e}"),
62 },
63 Err(e) => eprintln!("Failed to read password: {e}"),
64 }
65 }
66
67 fn get_download_url(
68 &self,
69 _platform: remote::RemotePlatform,
70 _release_channel: ReleaseChannel,
71 _version: Option<Version>,
72 _cx: &mut gpui::AsyncApp,
73 ) -> gpui::Task<gpui::Result<Option<String>>> {
74 unimplemented!()
75 }
76
77 fn download_server_binary_locally(
78 &self,
79 _platform: remote::RemotePlatform,
80 _release_channel: ReleaseChannel,
81 _version: Option<Version>,
82 _cx: &mut gpui::AsyncApp,
83 ) -> gpui::Task<gpui::Result<std::path::PathBuf>> {
84 unimplemented!()
85 }
86
87 fn set_status(&self, status: Option<&str>, _: &mut gpui::AsyncApp) {
88 if let Some(status) = status {
89 println!("SSH status: {status}");
90 }
91 }
92}
93fn main() -> Result<(), anyhow::Error> {
94 let args = Args::parse();
95
96 if let Some(socket) = &args.askpass {
97 askpass::main(socket);
98 return Ok(());
99 }
100
101 let query_str = args
102 .query
103 .ok_or_else(|| anyhow!("-q/--query is required"))?;
104 let query = if args.regex {
105 SearchQuery::regex(
106 query_str,
107 args.whole_word,
108 args.case_sensitive,
109 args.include_ignored,
110 false,
111 Default::default(),
112 Default::default(),
113 false,
114 None,
115 )
116 } else {
117 SearchQuery::text(
118 query_str,
119 args.whole_word,
120 args.case_sensitive,
121 args.include_ignored,
122 Default::default(),
123 Default::default(),
124 false,
125 None,
126 )
127 }?;
128 Application::headless().run(|cx| {
129 release_channel::init_test(semver::Version::new(0, 0, 0), ReleaseChannel::Dev, cx);
130 settings::init(cx);
131 let client = Client::production(cx);
132 let http_client = FakeHttpClient::with_200_response();
133 let (_, rx) = watch::channel(None);
134 let node = NodeRuntime::new(http_client, None, rx);
135 let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
136 let registry = Arc::new(LanguageRegistry::new(cx.background_executor().clone()));
137 let fs = Arc::new(RealFs::new(None, cx.background_executor().clone()));
138
139
140
141 cx.spawn(async move |cx| {
142 let project = if let Some(ssh_target) = args.ssh {
143 println!("Setting up SSH connection for {}", &ssh_target);
144 let ssh_connection_options = SshConnectionOptions::parse_command_line(&ssh_target)?;
145
146 let connection_options = remote::RemoteConnectionOptions::from(ssh_connection_options);
147 let delegate = Arc::new(BenchmarkRemoteClient);
148 let remote_connection = remote::connect(connection_options.clone(), delegate.clone(), cx).await.unwrap();
149
150 let (_tx, rx) = oneshot::channel();
151 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"))?;
152
153 cx.update(|cx| Project::remote(remote_client, client, node, user_store, registry, fs, false, cx))
154 } else {
155 println!("Setting up local project");
156 cx.update(|cx| Project::local(
157 client,
158 node,
159 user_store,
160 registry,
161 fs,
162 Some(Default::default()),
163 project::LocalProjectFlags {
164 init_worktree_trust: false,
165 ..Default::default()
166 },
167 cx,
168 ))
169 };
170 println!("Loading worktrees");
171 let worktrees = project.update(cx, |this, cx| {
172 args.worktrees
173 .into_iter()
174 .map(|worktree| this.find_or_create_worktree(worktree, true, cx))
175 .collect::<Vec<_>>()
176 });
177
178 let worktrees = futures::future::join_all(worktrees)
179 .await
180 .into_iter()
181 .collect::<Result<Vec<_>, anyhow::Error>>()?;
182
183 for (worktree, _) in &worktrees {
184 let scan_complete = worktree
185 .update(cx, |this, _| {
186 if let Some(local) = this.as_local() {
187 Some(local.scan_complete())
188 } else {
189 None
190 }
191 });
192 if let Some(scan_complete) = scan_complete {
193 scan_complete.await;
194 } else {
195 cx.background_executor().timer(Duration::from_secs(10)).await;
196 }
197
198 }
199 println!("Worktrees loaded");
200
201 println!("Starting a project search");
202 let timer = std::time::Instant::now();
203 let mut first_match = None;
204 let matches = project.update(cx, |this, cx| this.search(query, cx));
205 let mut matched_files = 0;
206 let mut matched_chunks = 0;
207 while let Ok(match_result) = matches.rx.recv().await {
208 if first_match.is_none() {
209 let time = timer.elapsed();
210 first_match = Some(time);
211 println!("First match found after {time:?}");
212 }
213 if let SearchResult::Buffer { ranges, .. } = match_result {
214 matched_files += 1;
215 matched_chunks += ranges.len();
216 } else {
217 break;
218 }
219 }
220 let elapsed = timer.elapsed();
221 println!(
222 "Finished project search after {elapsed:?}. Matched {matched_files} files and {matched_chunks} excerpts"
223 );
224 drop(project);
225 cx.update(|cx| cx.quit());
226
227 anyhow::Ok(())
228 })
229 .detach_and_log_err(cx);
230
231 });
232 Ok(())
233}