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