Skip to content

Commit 2c44db7

Browse files
omarabidomarabid
committed
wip: cli
1 parent 5170373 commit 2c44db7

File tree

7 files changed

+109
-156
lines changed

7 files changed

+109
-156
lines changed

‎pyroscope_cli/cli/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ edition = "2021"
99
utils = { path = "../utils" }
1010
core = { path = "../core" }
1111
clap_complete = "3.1"
12+
names= "0.13.0"
1213

1314
[dependencies.clap]
1415
version = "3.1"

‎pyroscope_cli/cli/src/lib.rs‎

Lines changed: 34 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -71,14 +71,14 @@ enum Commands {
7171
)]
7272
detect_subprocesses: bool,
7373
#[clap(
74+
arg_enum,
7475
name = "log_level",
7576
short,
7677
long = "log-level",
7778
value_name = "LOG_LEVEL",
78-
help = "",
79-
default_value = "info"
79+
help = "[default: info] log level for the application"
8080
)]
81-
log_level: LogLevel,
81+
log_level: Option<LogLevel>,
8282
#[clap(
8383
name = "no_logging",
8484
long = "no-logging",
@@ -112,26 +112,26 @@ enum Commands {
112112
name = "sample_rate",
113113
long = "sample-rate",
114114
value_name = "SAMPLE_RATE",
115-
help = "sample rate for the profiler in Hz. 100 means reading 100 times per second",
116-
default_value = "100"
115+
help = "[default: 100] sample rate for the profiler in Hz. 100 means reading 100 times per second"
117116
)]
118-
sample_rate: i32,
117+
sample_rate: Option<u32>,
119118
#[clap(
120119
name = "server_address",
121120
long = "server-address",
122121
value_name = "SERVER_ADDRESS",
123-
help = "Pyroscope server address",
124-
default_value = "http://localhost:4040"
122+
help = "[default: http://localhost:4040] Pyroscope server address"
125123
)]
126-
server_address: String,
124+
server_address: Option<String>,
127125
#[clap(
126+
arg_enum,
128127
name = "spy_name",
129128
long = "spy-name",
130129
value_name = "SPY_NAME",
131130
help = "name of the profiler to use"
132131
)]
133132
spy_name: Spy,
134133
#[clap(
134+
multiple = true,
135135
name = "tag",
136136
long = "tag",
137137
value_name = "TAG",
@@ -180,7 +180,7 @@ enum Commands {
180180
//value_name = "AUTH_TOKEN",
181181
//help = "authorization token used to upload profiling data"
182182
//)]
183-
auth_token: Option<String>,
183+
//auth_token: Option<String>,
184184
#[clap(
185185
name = "detect_subprocesses",
186186
long = "detect-subprocesses",
@@ -189,35 +189,28 @@ enum Commands {
189189
)]
190190
detect_subprocesses: bool,
191191
#[clap(
192-
name = "group_name",
193-
long = "group-name",
194-
value_name = "GROUP_NAME",
195-
help = "start process under specified group name"
196-
)]
197-
group_name: Option<String>,
198-
#[clap(
192+
arg_enum,
199193
name = "log_level",
200194
short,
201195
long = "log-level",
202196
value_name = "LOG_LEVEL",
203-
help = "",
204-
default_value = "info"
197+
help = "[default: info] log level for the application"
205198
)]
206-
log_level: LogLevel,
199+
log_level: Option<LogLevel>,
207200
#[clap(
208201
name = "no_logging",
209202
long = "no-logging",
210203
value_name = "NO_LOGGING",
211204
help = "disable logging from pyroscope"
212205
)]
213206
no_logging: bool,
214-
#[clap(
215-
name = "no_root_drop",
216-
long = "no-root-drop",
217-
value_name = "NO_ROOT_DROP",
218-
help = "disable permissions drop when ran under root. use this one if you want to run your command as root"
219-
)]
220-
no_root_drop: bool,
207+
//#[clap(
208+
//name = "no_root_drop",
209+
//long = "no-root-drop",
210+
//value_name = "NO_ROOT_DROP",
211+
//help = "disable permissions drop when ran under root. use this one if you want to run your command as root"
212+
//)]
213+
//no_root_drop: bool,
221214
#[clap(
222215
name = "pyspy_blocking",
223216
long = "pyspy-blocking",
@@ -236,19 +229,18 @@ enum Commands {
236229
name = "sample_rate",
237230
long = "sample-rate",
238231
value_name = "SAMPLE_RATE",
239-
help = "sample rate for the profiler in Hz. 100 means reading 100 times per second",
240-
default_value = "100"
232+
help = "[default: 100] sample rate for the profiler in Hz. 100 means reading 100 times per second"
241233
)]
242-
sample_rate: u32,
234+
sample_rate: Option<u32>,
243235
#[clap(
244236
name = "server_address",
245237
long = "server-address",
246238
value_name = "SERVER_ADDRESS",
247-
help = "Pyroscope server address",
248-
default_value = "http://localhost:4040"
239+
help = "[default: http://localhost:4040] Pyroscope server address"
249240
)]
250-
server_address: String,
241+
server_address: Option<String>,
251242
#[clap(
243+
arg_enum,
252244
name = "spy_name",
253245
long = "spy-name",
254246
value_name = "SPY_NAME",
@@ -287,6 +279,13 @@ enum Commands {
287279
help = "start process under specified user name"
288280
)]
289281
user_name: Option<String>,
282+
#[clap(
283+
name = "group_name",
284+
long = "group-name",
285+
value_name = "GROUP_NAME",
286+
help = "start process under specified group name"
287+
)]
288+
group_name: Option<String>,
290289
},
291290
#[clap(
292291
name = "config",
@@ -318,10 +317,11 @@ pub fn cli_match() -> Result<()> {
318317
AppConfig::merge_config(cli.config.as_deref())?;
319318

320319
let app = Cli::into_app();
320+
321+
// Merge clap args into config
321322
AppConfig::merge_args(app)?;
322323

323324
let mut app = Cli::into_app();
324-
325325
// Execute the subcommand
326326
match &cli.command {
327327
Commands::Exec { .. } => {

‎pyroscope_cli/core/src/commands.rs‎

Lines changed: 27 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use utils::app_config::AppConfig;
22
use utils::error::{Error, Result};
3+
use utils::types::{LogLevel, Spy};
34

45
use pyroscope::PyroscopeAgent;
56
use pyroscope_pyspy::{Pyspy, PyspyConfig};
@@ -13,13 +14,14 @@ use crate::profiler::Profiler;
1314

1415
/// exec command
1516
pub fn exec() -> Result<()> {
17+
// TODO: this is hacky
18+
set_application_name()?;
19+
set_tags()?;
20+
1621
// Get command to execute
1722
let command = AppConfig::get::<Option<String>>("command")?
1823
.ok_or_else(|| Error::new("command unwrap failed"))?;
1924

20-
//dbg!(AppConfig::fetch()?);
21-
//return Ok(());
22-
2325
// Get UID
2426
let uid = AppConfig::get::<Option<u32>>("user_name").unwrap_or(None);
2527
// Get GID
@@ -28,8 +30,6 @@ pub fn exec() -> Result<()> {
2830
// Create new executor and run it
2931
let executor = Executor::new(command.as_ref(), "", uid, gid).run()?;
3032

31-
println!("stopped here?");
32-
3333
// Set PID
3434
AppConfig::set("pid", executor.get_pid()?.to_string().as_str())?;
3535

@@ -59,7 +59,10 @@ pub fn exec() -> Result<()> {
5959

6060
/// connect command
6161
pub fn connect() -> Result<()> {
62-
// Initialize profiler
62+
// TODO: this is hacky
63+
set_application_name()?;
64+
set_tags()?;
65+
6366
let mut profiler = Profiler::default();
6467

6568
profiler.init()?;
@@ -90,86 +93,32 @@ pub fn config() -> Result<()> {
9093
Ok(())
9194
}
9295

93-
pub fn rbspy() -> Result<()> {
94-
let (tx, rx) = channel();
95-
96-
let pid: i32 = AppConfig::get("pid")?;
97-
let sample_rate: u32 = AppConfig::get("sample_rate")?;
98-
let lock_process: bool = AppConfig::get("rbspy_blocking")?;
99-
let with_subprocesses: bool = AppConfig::get("detect_subprocesses")?;
100-
101-
let config = RbspyConfig::new(pid)
102-
.sample_rate(sample_rate)
103-
.lock_process(lock_process)
104-
.with_subprocesses(with_subprocesses);
105-
106-
println!("Connecting to PID {}", pid);
107-
108-
let server_address: String = AppConfig::get("server_address")?;
109-
110-
let mut agent = PyroscopeAgent::builder(server_address, "rbspy.basic".to_string())
111-
.backend(Rbspy::new(config))
112-
.build()
113-
.unwrap();
96+
fn set_application_name() -> Result<()> {
97+
let pre_app_name: String = AppConfig::get::<String>("application_name").unwrap_or_else(|_| {
98+
names::Generator::default()
99+
.next()
100+
.unwrap_or("unassigned.app".to_string())
101+
.replace("-", ".")
102+
});
114103

115-
agent.start().unwrap();
104+
let pre = match AppConfig::get::<Spy>("spy_name")? {
105+
Spy::Pyspy => "pyspy",
106+
Spy::Rbspy => "rbspy",
107+
_ => "none",
108+
};
116109

117-
ctrlc::set_handler(move || {
118-
tx.send(()).unwrap();
119-
})
120-
.expect("Error setting Ctrl-C handler");
121-
122-
println!("Press Ctrl-C to exit.");
123-
124-
rx.recv().unwrap();
125-
126-
println!("Exiting.");
110+
// add pre to pre_app_name
111+
let app_name = format!("{}.{}", pre, pre_app_name);
127112

128-
agent.stop().unwrap();
129-
130-
drop(agent);
113+
AppConfig::set("application_name", app_name.as_str())?;
131114

132115
Ok(())
133116
}
134117

135-
pub fn pyspy() -> Result<()> {
136-
let (tx, rx) = channel();
137-
138-
let pid: i32 = AppConfig::get("pid")?;
139-
let sample_rate: u32 = AppConfig::get("sample_rate")?;
140-
let lock_process: bool = AppConfig::get("pyspy_blocking")?;
141-
let with_subprocesses: bool = AppConfig::get("detect_subprocesses")?;
142-
143-
let config = PyspyConfig::new(pid)
144-
.sample_rate(sample_rate)
145-
.lock_process(lock_process)
146-
.with_subprocesses(with_subprocesses);
147-
148-
println!("Connecting to PID {}", pid);
149-
150-
let server_address: String = AppConfig::get("server_address")?;
151-
152-
let mut agent = PyroscopeAgent::builder(server_address, "pyspy.basic".to_string())
153-
.backend(Pyspy::new(config))
154-
.build()
155-
.unwrap();
156-
157-
agent.start().unwrap();
158-
159-
ctrlc::set_handler(move || {
160-
tx.send(()).unwrap();
161-
})
162-
.expect("Error setting Ctrl-C handler");
163-
164-
println!("Press Ctrl-C to exit.");
165-
166-
rx.recv().unwrap();
167-
168-
println!("Exiting.");
169-
170-
agent.stop().unwrap();
118+
fn set_tags() -> Result<()> {
119+
let tag: String = AppConfig::get::<String>("tag").unwrap_or("".to_string());
171120

172-
drop(agent);
121+
AppConfig::set("tag", tag.as_str())?;
173122

174123
Ok(())
175124
}

0 commit comments

Comments
 (0)