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#[derive(Debug, Clone, Copy, clap::ValueEnum)]
135pub enum PgSslMode {
136    Disable,
137    Prefer,
138    Require,
139    VerifyCa,
140    VerifyFull,
141}
142
143impl PgSslMode {
144    pub fn as_str(self) -> &'static str {
145        match self {
146            Self::Disable => "disable",
147            Self::Prefer => "prefer",
148            Self::Require => "require",
149            Self::VerifyCa => "verify-ca",
150            Self::VerifyFull => "verify-full",
151        }
152    }
153}
154
155///
156/// ███████╗ █████╗ ██╗   ██╗ ██████╗███████╗████████╗
157/// ██╔════╝██╔══██╗██║   ██║██╔════╝██╔════╝╚══██╔══╝
158/// █████╗  ███████║██║   ██║██║     █████╗     ██║
159/// ██╔══╝  ██╔══██║██║   ██║██║     ██╔══╝     ██║
160/// ██║     ██║  ██║╚██████╔╝╚██████╗███████╗   ██║
161/// ╚═╝     ╚═╝  ╚═╝ ╚═════╝  ╚═════╝╚══════╝   ╚═╝
162/// Fast, async, and concurrent data applications.
163///
164#[derive(Parser)]
165#[command(author, version, verbatim_doc_comment)]
166pub struct Args {
167    #[command(subcommand)]
168    pub command: Commands,
169
170    /// The host to bind to.
171    #[arg(long, env = "FAUCET_HOST", default_value = "127.0.0.1:3838")]
172    pub host: String,
173
174    /// The IP address to extract from.
175    /// Defaults to client address.
176    #[arg(short, long, env = "FAUCET_IP_FROM", default_value = "client")]
177    pub ip_from: IpFrom,
178
179    /// Command, path, or executable to run Rscript.
180    #[arg(long, short, env = "FAUCET_RSCRIPT", default_value = "Rscript")]
181    pub rscript: OsString,
182
183    /// Command, path, or executable to run quarto.
184    #[arg(long, short, env = "FAUCET_QUARTO", default_value = "quarto")]
185    pub quarto: OsString,
186
187    /// Save logs to a file. Will disable colors!
188    #[arg(long, short, env = "FAUCET_LOG_FILE", default_value = None)]
189    pub log_file: Option<PathBuf>,
190
191    #[arg(long, short, env = "FAUCET_MAX_LOG_FILE_SIZE", default_value = None, value_parser = |s: &str| parse_size::parse_size(s))]
192    /// The maximum size of the log file. (Ex. 10M, 1GB)
193    pub max_log_file_size: Option<u64>,
194
195    /// The strategy for shutting down faucet
196    #[arg(long, env = "FAUCET_SHUTDOWN", default_value = "immediate")]
197    pub shutdown: Shutdown,
198
199    /// Maximum size of a WebSocket message. This is useful for DDOS prevention. Not set means no size limit.
200    #[arg(long, env = "FAUCET_MAX_MESSAGE_SIZE", default_value = None, value_parser = |s: &str| parse_size::parse_size(s))]
201    pub max_message_size: Option<u64>,
202
203    /// Connection string to a PostgreSQL database for saving HTTP events.
204    #[arg(long, env = "FAUCET_TELEMETRY_POSTGRES_STRING", default_value = None)]
205    pub pg_con_string: Option<String>,
206
207    /// Path to CA certificate for PostgreSQL SSL/TLS.
208    #[arg(long, env = "FAUCET_TELEMETRY_POSTGRES_SSLCERT", default_value = None)]
209    pub pg_sslcert: Option<PathBuf>,
210
211    /// SSL mode for PostgreSQL connection (disable, prefer, require, verify-ca, verify-full).
212    #[arg(
213        long,
214        env = "FAUCET_TELEMETRY_POSTGRES_SSLMODE",
215        default_value = "prefer"
216    )]
217    pub pg_sslmode: PgSslMode,
218
219    /// Save HTTP events on PostgreSQL under a specific namespace.
220    #[arg(long, env = "FAUCET_TELEMETRY_NAMESPACE", default_value = "faucet")]
221    pub telemetry_namespace: String,
222
223    /// Represents the source code version of the service to run. This is useful for telemetry.
224    #[arg(long, env = "FAUCET_TELEMETRY_VERSION", default_value = None)]
225    pub telemetry_version: Option<String>,
226}
227
228impl StartArgs {
229    pub fn server_type(&self) -> WorkerType {
230        match self.type_ {
231            ServerType::Plumber => WorkerType::Plumber,
232            ServerType::Shiny => WorkerType::Shiny,
233            ServerType::QuartoShiny => WorkerType::QuartoShiny,
234            ServerType::Auto => {
235                if is_plumber(&self.dir) {
236                    WorkerType::Plumber
237                } else if is_shiny(&self.dir) {
238                    WorkerType::Shiny
239                } else {
240                    log::error!(target: "faucet", "Could not determine worker type. Please specify with --type.");
241                    std::process::exit(1);
242                }
243            }
244        }
245    }
246}
faucet_server/cli.rs

faucet_server/
cli.rs
