faucet_server/
cli.rs

1use std::{
2    ffi::OsString,
3    path::{Path, PathBuf},
4};
5
6use clap::{Parser, Subcommand};
7
8use crate::client::{load_balancing, worker::WorkerType};
9
10fn is_plumber(dir: &Path) -> bool {
11    let plumber = dir.join("plumber.R");
12    let plumber_entrypoint = dir.join("entrypoint.R");
13    plumber.exists() || plumber_entrypoint.exists()
14}
15
16fn is_shiny(dir: &Path) -> bool {
17    let shiny_app = dir.join("app.R");
18    let shiny_ui = dir.join("ui.R");
19    let shiny_server = dir.join("server.R");
20    shiny_app.exists() || (shiny_ui.exists() && shiny_server.exists())
21}
22
23#[derive(clap::ValueEnum, Debug, Clone, Copy)]
24enum ServerType {
25    Plumber,
26    Shiny,
27    QuartoShiny,
28    Auto,
29}
30
31#[derive(clap::ValueEnum, Debug, Clone, Copy)]
32pub enum Strategy {
33    /// Sends requests to workers in a round-robin fashion.
34    RoundRobin,
35    /// Hashes the IP address of the client to determine which worker to send the request to.
36    IpHash,
37    /// Adds a cookie to the requests to identify the worker to send the
38    /// request to. This is useful for sticky sessions from within the same
39    /// network.
40    CookieHash,
41}
42
43impl From<Strategy> for load_balancing::Strategy {
44    fn from(value: Strategy) -> Self {
45        match value {
46            Strategy::RoundRobin => load_balancing::Strategy::RoundRobin,
47            Strategy::IpHash => load_balancing::Strategy::IpHash,
48            Strategy::CookieHash => load_balancing::Strategy::CookieHash,
49        }
50    }
51}
52
53#[derive(clap::ValueEnum, Debug, Clone, Copy)]
54pub enum IpFrom {
55    Client,
56    XForwardedFor,
57    XRealIp,
58}
59
60impl From<IpFrom> for load_balancing::IpExtractor {
61    fn from(value: IpFrom) -> Self {
62        match value {
63            IpFrom::Client => load_balancing::IpExtractor::ClientAddr,
64            IpFrom::XForwardedFor => load_balancing::IpExtractor::XForwardedFor,
65            IpFrom::XRealIp => load_balancing::IpExtractor::XRealIp,
66        }
67    }
68}
69
70#[derive(clap::ValueEnum, Debug, Clone, Copy, Default)]
71pub enum Shutdown {
72    Graceful,
73    #[default]
74    Immediate,
75}
76
77#[derive(Parser, Debug)]
78pub struct StartArgs {
79    /// The number of threads to use to handle requests.
80    #[arg(short, long, env = "FAUCET_WORKERS", default_value_t = num_cpus::get())]
81    pub workers: usize,
82
83    /// The load balancing strategy to use.
84    #[arg(short, long, env = "FAUCET_STRATEGY", default_value = "round-robin")]
85    pub strategy: Strategy,
86
87    /// The type of workers to spawn.
88    #[arg(short, long, env = "FAUCET_TYPE", default_value = "auto")]
89    type_: ServerType,
90
91    /// The directory to spawn workers in.
92    /// Defaults to the current directory.
93    #[arg(short, long, env = "FAUCET_DIR", default_value = ".")]
94    pub dir: PathBuf,
95
96    /// Argument passed on to `appDir` when running Shiny.
97    #[arg(long, short, env = "FAUCET_APP_DIR", default_value = None)]
98    pub app_dir: Option<String>,
99
100    /// Quarto Shiny file path.
101    #[arg(long, short, env = "FAUCET_QMD", default_value = None)]
102    pub qmd: Option<PathBuf>,
103}
104
105#[derive(Parser, Debug)]
106pub struct RouterArgs {
107    /// Router config file.
108    #[arg(
109        long,
110        short,
111        env = "FAUCET_ROUTER_CONF",
112        default_value = "./frouter.toml"
113    )]
114    pub conf: PathBuf,
115}
116
117#[derive(Subcommand, Debug)]
118pub enum Commands {
119    /// Start a simple faucet server.
120    #[command(name = "start")]
121    Start(StartArgs),
122    /// Runs faucet in "router" mode. (Experimental)
123    #[command(name = "router")]
124    Router(RouterArgs),
125}
126
127///
128/// ███████╗ █████╗ ██╗   ██╗ ██████╗███████╗████████╗
129/// ██╔════╝██╔══██╗██║   ██║██╔════╝██╔════╝╚══██╔══╝
130/// █████╗  ███████║██║   ██║██║     █████╗     ██║
131/// ██╔══╝  ██╔══██║██║   ██║██║     ██╔══╝     ██║
132/// ██║     ██║  ██║╚██████╔╝╚██████╗███████╗   ██║
133/// ╚═╝     ╚═╝  ╚═╝ ╚═════╝  ╚═════╝╚══════╝   ╚═╝
134/// Fast, async, and concurrent data applications.
135///
136#[derive(Parser)]
137#[command(author, version, verbatim_doc_comment)]
138pub struct Args {
139    #[command(subcommand)]
140    pub command: Commands,
141
142    /// The host to bind to.
143    #[arg(long, env = "FAUCET_HOST", default_value = "127.0.0.1:3838")]
144    pub host: String,
145
146    /// The IP address to extract from.
147    /// Defaults to client address.
148    #[arg(short, long, env = "FAUCET_IP_FROM", default_value = "client")]
149    pub ip_from: IpFrom,
150
151    /// Command, path, or executable to run Rscript.
152    #[arg(long, short, env = "FAUCET_RSCRIPT", default_value = "Rscript")]
153    pub rscript: OsString,
154
155    /// Command, path, or executable to run quarto.
156    #[arg(long, short, env = "FAUCET_QUARTO", default_value = "quarto")]
157    pub quarto: OsString,
158
159    /// Save logs to a file. Will disable colors!
160    #[arg(long, short, env = "FAUCET_LOG_FILE", default_value = None)]
161    pub log_file: Option<PathBuf>,
162
163    #[arg(long, short, env = "FAUCET_MAX_LOG_FILE_SIZE", default_value = None, value_parser = |s: &str| parse_size::parse_size(s))]
164    /// The maximum size of the log file. (Ex. 10M, 1GB)
165    pub max_log_file_size: Option<u64>,
166
167    /// The strategy for shutting down faucet
168    #[arg(long, env = "FAUCET_SHUTDOWN", default_value = "immediate")]
169    pub shutdown: Shutdown,
170
171    /// Connection string to a PostgreSQL database for saving HTTP events.
172    #[arg(long, env = "FAUCET_TELEMETRY_POSTGRES_STRING", default_value = None)]
173    pub pg_con_string: Option<String>,
174
175    /// Save HTTP events on PostgreSQL under a specific namespace.
176    #[arg(long, env = "FAUCET_TELEMETRY_NAMESPACE", default_value = "faucet")]
177    pub telemetry_namespace: String,
178
179    /// Represents the source code version of the service to run. This is useful for telemetry.
180    #[arg(long, env = "FAUCET_TELEMETRY_VERSION", default_value = None)]
181    pub telemetry_version: Option<String>,
182}
183
184impl StartArgs {
185    pub fn server_type(&self) -> WorkerType {
186        match self.type_ {
187            ServerType::Plumber => WorkerType::Plumber,
188            ServerType::Shiny => WorkerType::Shiny,
189            ServerType::QuartoShiny => WorkerType::QuartoShiny,
190            ServerType::Auto => {
191                if is_plumber(&self.dir) {
192                    WorkerType::Plumber
193                } else if is_shiny(&self.dir) {
194                    WorkerType::Shiny
195                } else {
196                    log::error!(target: "faucet", "Could not determine worker type. Please specify with --type.");
197                    std::process::exit(1);
198                }
199            }
200        }
201    }
202}