Skip to main content

pubhubs/cli/
serve.rs

1use crate::cli;
2use crate::servers::Config;
3use crate::servers::Server as _;
4
5use anyhow::Result;
6
7#[derive(clap::Args, Debug)]
8pub struct ServeArgs {
9    #[command(flatten)]
10    common: cli::CommonArgs,
11
12    /// Run not all servers specified in the configuration file, but only these
13    #[arg(value_enum, short, long, value_name = "SERVERS")]
14    only: Option<Vec<crate::servers::Name>>,
15}
16
17impl ServeArgs {
18    pub fn run(self, _spec: &mut clap::Command) -> Result<()> {
19        if std::env::var("RUST_LOG").is_ok() {
20            env_logger::init();
21        }
22
23        let config = self.adjust_config(self.common.load_config()?)?;
24
25        log::info!("version: {}", crate::servers::version::VERSION);
26
27        tokio::runtime::Builder::new_multi_thread()
28            .enable_all()
29            .build()?
30            .block_on(async {
31                let (set, shutdown_sender) = crate::servers::Set::new(&config)?;
32
33                tokio::spawn(async move {
34                    tokio::signal::ctrl_c()
35                        .await
36                        .expect("failed to await ctrl+c");
37                    log::info!("ctrl+c received; shutting down server(s)");
38
39                    drop(shutdown_sender);
40
41                    tokio::signal::ctrl_c()
42                        .await
43                        .expect("failed to await ctrl+c");
44                    log::warn!("second ctrl+c received; aborting process...");
45                    std::process::abort();
46                });
47
48                let err_count = set.wait().await;
49                if err_count > 0 {
50                    anyhow::bail!("{} servers did not shutdown cleanly", err_count);
51                }
52
53                Ok(())
54            })
55    }
56
57    /// Adjust config based on the arguments (`--only`, ...) in `self`.
58    fn adjust_config(&self, config: Config) -> Result<Config> {
59        Ok(self.apply_only(config))
60    }
61
62    /// Filters servers from config that were not specified to run
63    fn apply_only(&self, mut config: Config) -> Config {
64        if self.only.is_none() {
65            // if no --only was specified, run all servers
66            return config;
67        }
68
69        let only: std::collections::HashSet<&crate::servers::Name> =
70            self.only.as_ref().unwrap().iter().collect();
71
72        macro_rules! remove_server_if_needed {
73            ($server:ident) => {
74                if !only.contains(&crate::servers::$server::Server::NAME) {
75                    config.$server = None;
76                }
77            };
78        }
79
80        crate::servers::for_all_servers!(remove_server_if_needed);
81
82        config
83    }
84}