main.rs

  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                    false,
164                    cx,
165                ))
166                };
167                println!("Loading worktrees");
168                let worktrees = project.update(cx, |this, cx| {
169                    args.worktrees
170                        .into_iter()
171                        .map(|worktree| this.find_or_create_worktree(worktree, true, cx))
172                        .collect::<Vec<_>>()
173                });
174
175                let worktrees = futures::future::join_all(worktrees)
176                    .await
177                    .into_iter()
178                    .collect::<Result<Vec<_>, anyhow::Error>>()?;
179
180                for (worktree, _) in &worktrees {
181                    let scan_complete = worktree
182                        .update(cx, |this, _| {
183                            if let Some(local) = this.as_local() {
184                                Some(local.scan_complete())
185                            } else {
186                                None
187                            }
188                        });
189                    if let Some(scan_complete) = scan_complete {
190                        scan_complete.await;
191                    } else {
192                        cx.background_executor().timer(Duration::from_secs(10)).await;
193                    }
194
195                }
196                println!("Worktrees loaded");
197
198                println!("Starting a project search");
199                let timer = std::time::Instant::now();
200                let mut first_match = None;
201                let matches = project.update(cx, |this, cx| this.search(query, cx));
202                let mut matched_files = 0;
203                let mut matched_chunks = 0;
204                while let Ok(match_result) = matches.rx.recv().await {
205                    if first_match.is_none() {
206                        let time = timer.elapsed();
207                        first_match = Some(time);
208                        println!("First match found after {time:?}");
209                    }
210                    if let SearchResult::Buffer { ranges, .. } = match_result {
211                        matched_files += 1;
212                        matched_chunks += ranges.len();
213                    } else {
214                        break;
215                    }
216                }
217                let elapsed = timer.elapsed();
218                println!(
219                    "Finished project search after {elapsed:?}. Matched {matched_files} files and {matched_chunks} excerpts"
220                );
221                drop(project);
222                cx.update(|cx| cx.quit());
223
224                anyhow::Ok(())
225            })
226            .detach_and_log_err(cx);
227
228    });
229    Ok(())
230}