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

faucet_server/
cli.rs
