functiontrace 0.4.4__tar.gz → 0.5.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -7,13 +7,19 @@ edition = "2021"
7
7
  color-eyre = { version = "0.6", default-features = false }
8
8
  dashmap = "6"
9
9
  functiontrace-server = { version = "0.8.4", default-features = false, features = ["serialize"] }
10
+ jiff = { version = "0.2", default-features = false, features = ["std", "serde"] }
10
11
  libc = "0.2"
12
+ owo-colors = "4.2"
11
13
  pyo3 = { version = "0.25", features = ["eyre"] }
14
+ reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
12
15
  rmp-serde = "1.3"
13
16
  serde = "1"
17
+ toml = "0.9"
18
+ tracing = { version = "0.1", features = ["log"] }
19
+ tracing-subscriber = { version = "0.3", default-features = false, features = ["ansi", "env-filter", "fmt", "std"] }
14
20
 
15
21
  [build-dependencies]
16
- toml = { version = "0.8", default-features = false, features = ["parse"] }
22
+ toml = { version = "0.9", default-features = false, features = ["parse", "serde"] }
17
23
  pyo3-build-config = { version = "0.25", features = ["resolve-config"] }
18
24
 
19
25
  [lib]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: functiontrace
3
- Version: 0.4.4
3
+ Version: 0.5.1
4
4
  Summary: A profiling tool that traces your application's execution
5
5
  Author: Alistair Dobke
6
6
  Author-email: Matt Bryant <functiontrace@programsareproofs.com>
@@ -5,14 +5,24 @@ fn main() -> std::io::Result<()> {
5
5
  let pyproject = std::fs::read_to_string("pyproject.toml")?
6
6
  .parse::<toml::Table>()
7
7
  .expect("pyproject.toml is valid toml");
8
- let version = pyproject
8
+ let project = pyproject
9
9
  .get("project")
10
10
  .and_then(|p| p.as_table())
11
- .and_then(|p| p.get("version"))
11
+ .expect("project field is present in pyproject.toml");
12
+
13
+ let version = project
14
+ .get("version")
12
15
  .and_then(|v| v.as_str())
13
16
  .expect("version field is present in pyproject.toml");
14
17
  println!("cargo::rustc-env=PACKAGE_VERSION={}", version);
15
18
 
19
+ let min_python = project
20
+ .get("requires-python")
21
+ .and_then(|v| v.as_str())
22
+ .and_then(|v| v.strip_prefix(">=3."))
23
+ .expect("requires-python field is present in pyproject.toml");
24
+ println!("cargo::rustc-env=MIN_PYTHON_VERSION={}", min_python);
25
+
16
26
  // Figure out which Python version we're building for
17
27
  pyo3_build_config::use_pyo3_cfgs();
18
28
 
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: functiontrace
3
- Version: 0.4.4
3
+ Version: 0.5.1
4
4
  Summary: A profiling tool that traces your application's execution
5
5
  Author: Alistair Dobke
6
6
  Author-email: Matt Bryant <functiontrace@programsareproofs.com>
@@ -1,6 +1,5 @@
1
1
  import argparse
2
2
  import os
3
- import shutil
4
3
  import sys
5
4
  import tempfile
6
5
  from importlib import machinery, util
@@ -20,19 +19,6 @@ exec $(which {python}) -m functiontrace "$@"
20
19
 
21
20
 
22
21
  def setup_dependencies():
23
- # We need the functiontrace-server installed and locatable in order to
24
- # trace anything.
25
- if shutil.which("functiontrace-server") is None:
26
- print(
27
- "Unable to find `functiontrace-server` in the current $PATH.",
28
- file=sys.stderr,
29
- )
30
- print(
31
- "See https://functiontrace.com#installation for installation instructions.",
32
- file=sys.stderr,
33
- )
34
- sys.exit(1)
35
-
36
22
  # Generate a temp directory to store our wrappers in. We'll temporarily
37
23
  # add this directory to our path.
38
24
  tempdir = tempfile.mkdtemp(prefix="py-functiontrace")
@@ -41,7 +27,8 @@ def setup_dependencies():
41
27
  # Generate wrappers for the various Python versions we support to ensure
42
28
  # they're included in our PATH.
43
29
  wrap_pythons = ["python", "python3"] + [
44
- "python3.{}".format(minor) for minor in range(6, 13)
30
+ "python3.{}".format(minor)
31
+ for minor in range(rust._min_python_version, rust._max_python_version + 1)
45
32
  ]
46
33
  for python in wrap_pythons:
47
34
  with open(os.path.join(tempdir, python), "w") as f:
@@ -79,6 +66,9 @@ def main():
79
66
  parser.print_help()
80
67
  sys.exit(1)
81
68
 
69
+ # Load our config and verify the license before we trace.
70
+ rust.load_config()
71
+
82
72
  # Ignore ourselves, keeping sys.argv looking reasonable as the child script
83
73
  # will expect it to be sane.
84
74
  sys.argv[:] = args.script
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "functiontrace"
3
- version = "0.4.4"
3
+ version = "0.5.1"
4
4
  description = "A profiling tool that traces your application's execution"
5
5
  authors = [
6
6
  { name = "Matt Bryant", email = "functiontrace@programsareproofs.com"},
@@ -1,12 +1,13 @@
1
1
  #[pyo3::pymodule]
2
2
  mod _functiontrace_rs {
3
- use color_eyre::eyre::{eyre, OptionExt, Result, WrapErr};
3
+ use color_eyre::eyre::{OptionExt, Result, WrapErr, eyre};
4
4
  use functiontrace_server::function_trace;
5
- use pyo3::prelude::*; // TODO: What's in here?
6
- use serde::Serialize;
5
+ use jiff::Timestamp;
6
+ use pyo3::prelude::*;
7
+ use serde::{Deserialize, Serialize};
7
8
  use std::borrow::Cow;
8
9
  use std::cell::Cell;
9
- use std::ffi::{c_void, CStr, OsString};
10
+ use std::ffi::{CStr, OsString, c_void};
10
11
  use std::io::Write;
11
12
  use std::os::unix::net::UnixStream;
12
13
  use std::sync::atomic::{AtomicBool, Ordering};
@@ -16,6 +17,18 @@ mod _functiontrace_rs {
16
17
  ////////////////////////////////////////////////////////////////////////////////////////////////
17
18
  // Core Types
18
19
  ////////////////////////////////////////////////////////////////////////////////////////////////
20
+
21
+ /// The configuration file for FunctionTrace.
22
+ ///
23
+ /// TODO: This is currently only used for licensing, but could reasonably be used to augment or
24
+ /// replace some of the env vars we use.
25
+ #[derive(Serialize, Deserialize, Clone, PartialEq)]
26
+ struct Config {
27
+ /// The user's license key. This is NONCOMMERCIAL where relevant, and otherwise must
28
+ /// contain a key that will be validated on startup.
29
+ license_key: Option<String>,
30
+ }
31
+
19
32
  /// The size of messagepack buffers. This is picked somewhat arbitrarily, but should be
20
33
  /// reasonably small since we'll need one per thread, and should be large enough that we
21
34
  /// can fit many messages in before needing to do expensive UnixStream syscalls.
@@ -72,7 +85,7 @@ mod _functiontrace_rs {
72
85
 
73
86
  /// Flush our local buffer to the socket.
74
87
  fn flush(&mut self) -> std::io::Result<()> {
75
- // We're about to send all of the data, so reset our buffer. Do this now to ensure we
88
+ // We're about to send all the data, so reset our buffer. Do this now to ensure we
76
89
  // don't miss it, since the buffer will be left in an invalid state if this doesn't
77
90
  // happen.
78
91
  let head = self.head;
@@ -182,12 +195,40 @@ mod _functiontrace_rs {
182
195
  /// Initialization code run on this module's import (mostly useful for adding constants)
183
196
  #[pymodule_init]
184
197
  fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
185
- // Arbitrary code to run at the module initialization
186
198
  m.add("__version__", env!("PACKAGE_VERSION"))
187
199
  .wrap_err("Failed to expose PACKAGE_VERSION")?;
188
200
 
201
+ // Expose the minimum and maximum Python 3 versions we support
202
+ m.add(
203
+ "_min_python_version",
204
+ env!("MIN_PYTHON_VERSION")
205
+ .parse::<u32>()
206
+ .wrap_err("Failed to convert MIN_PYTHON_VERSION to int")?,
207
+ )
208
+ .wrap_err("Failed to expose _min_python_version")?;
209
+ m.add("_max_python_version", 13)
210
+ .wrap_err("Failed to expose _max_python_version")?;
211
+
189
212
  color_eyre::install().wrap_err("Failed to install color_eyre")?;
190
213
 
214
+ // If someone else has already installed a tracing subscriber, we'll fail to initialize
215
+ // but get _something_ anyways.
216
+ // This may not be possible, but I suspect a different Rust extension could interfere with
217
+ // us so might as well be defensive.
218
+ if let Err(e) = tracing_subscriber::fmt::fmt()
219
+ .with_target(false)
220
+ .with_ansi(true)
221
+ .with_env_filter(tracing_subscriber::filter::EnvFilter::from_default_env())
222
+ .try_init()
223
+ {
224
+ tracing::error!(
225
+ error = e,
226
+ "Failed to install tracing_subscriber for FunctionTrace"
227
+ );
228
+ }
229
+
230
+ tracing::info!("Loaded `functiontrace` python module!");
231
+
191
232
  Ok(())
192
233
  }
193
234
 
@@ -213,6 +254,57 @@ mod _functiontrace_rs {
213
254
  return Err(eyre!("FunctionTrace is already initialized!"));
214
255
  }
215
256
 
257
+ // Store the processes' args, since threads will need it as part of the registration message.
258
+ ARGV.set({
259
+ // Parse args out of `sys.argv`, since it excludes the Python executable and shows the
260
+ // script's path first. This is much more complicated than using `std::env::args`, but
261
+ // is more practical.
262
+ Python::with_gil(|py| -> Result<String> {
263
+ let argv = py
264
+ .import("sys")
265
+ .and_then(|sys| sys.getattr("argv"))
266
+ .wrap_err("Failed to access sys.argv")?;
267
+
268
+ let args = argv
269
+ .downcast::<pyo3::types::PyList>()
270
+ .map_err(|e| eyre!("Failed to read sys.argv as list: {}", e))?;
271
+
272
+ let mut full_args = args.iter().fold(String::new(), |s, arg| {
273
+ // Fetch the next arg as a string if possible.
274
+ let arg = arg
275
+ .downcast::<pyo3::types::PyString>()
276
+ .map_err(|e| eyre!("Failed to read sys.argv[x] as string: {}", e))
277
+ .and_then(|arg| {
278
+ arg.to_str()
279
+ .wrap_err("Failed to read Rust string out of PyString")
280
+ })
281
+ .unwrap_or("<UNKNOWN>");
282
+
283
+ s + " " + arg
284
+ });
285
+
286
+ // We'll display this in the UI, so ensure it's under (the rather
287
+ // arbitrary) 100 chars.
288
+ if full_args.chars().count() > 95 {
289
+ full_args.truncate(95);
290
+ full_args.push_str(" ...");
291
+ }
292
+
293
+ // We'll be cloning this
294
+ full_args.shrink_to_fit();
295
+ Ok(full_args)
296
+ })
297
+ .unwrap_or_else(|_err| {
298
+ tracing::warn!("Failed to read ARGV");
299
+
300
+ "<UNKNOWN>".into()
301
+ })
302
+ })
303
+ .expect("ARGV is only set in `trace_initialization`");
304
+
305
+ /////////////////////////////////////////////////////////////////////////
306
+ // Server initialization
307
+ /////////////////////////////////////////////////////////////////////////
216
308
  // Check if we're in a subprocess of a command being run under functiontrace. If we are,
217
309
  // we should connect to the same socket. Otherwise, we should setup the full profiling
218
310
  // configuration ourselves.
@@ -222,140 +314,105 @@ mod _functiontrace_rs {
222
314
  SOCKADDR
223
315
  .set(existing_sockaddr)
224
316
  .expect("sockaddr is only set in `trace_initialization`");
225
- return Ok(());
226
- }
317
+ } else {
318
+ // Launch the functiontrace-server as a daemon, then record the address of the socket we'll
319
+ // need to communicate with it on.
320
+ let sockaddr = {
321
+ let server = match std::process::Command::new("functiontrace-server")
322
+ .arg("--directory")
323
+ .arg(output_directory)
324
+ .spawn()
325
+ {
326
+ Ok(server) => server,
327
+ Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
328
+ eprintln!("Unable to find `functiontrace-server` in the current $PATH.");
329
+ eprintln!(
330
+ "See https://functiontrace.com#installation for installation instructions."
331
+ );
332
+
333
+ // Exit here rather than returning an error, since that'll emit a confusing
334
+ // Python stacktrace.
335
+ std::process::exit(1);
336
+ }
337
+ e => e.wrap_err("Failed to launch functiontrace-server")?,
338
+ };
227
339
 
228
- // Launch the functiontrace-server as a daemon, then record the address of the socket we'll
229
- // need to communicate with it on.
230
- let sockaddr = {
231
- let server = std::process::Command::new("functiontrace-server")
232
- .arg("--directory")
233
- .arg(output_directory)
234
- .spawn()
235
- .wrap_err("Failed to launch functiontrace-server")?;
340
+ format!("/tmp/functiontrace-server.sock.{}", server.id()).into()
341
+ };
236
342
 
237
- format!("/tmp/functiontrace-server.sock.{}", server.id()).into()
238
- };
343
+ Python::with_gil(|_| {
344
+ // Register an env var so subprocesses know how to connect to our
345
+ // profile.
346
+ unsafe {
347
+ // SAFETY: `set_var` can only be done while no one else is reading or writing to the
348
+ // environment. We're holding the GIL, so this _should_ be true, but there's no real
349
+ // way to enforce this.
350
+ //
351
+ // This doesn't seem to blow up in C even when we don't have the GIL, so
352
+ // it's hopefully good enough in practice since there's not really an alternative.
353
+ std::env::set_var(BREADCRUMB, &sockaddr);
354
+ }
355
+ });
239
356
 
240
- Python::with_gil(|_| {
241
- // Register an env var so subprocesses know how to connect to our
242
- // profile.
243
- unsafe {
244
- // SAFETY: `set_var` can only be done while no one else is reading or writing to the
245
- // environment. We're holding the GIL, so this _should_ be true, but there's no real
246
- // way to enforce this.
247
- //
248
- // This doesn't seem to blow up in C even when we don't have the GIL, so
249
- // it's hopefully good enough in practice since there's not really an alternative.
250
- std::env::set_var(BREADCRUMB, &sockaddr);
251
- }
252
- });
357
+ // Store the server's address so our threads know how to talk to it.
358
+ SOCKADDR
359
+ .set(sockaddr)
360
+ .expect("sockaddr is only set in `trace_initialization`");
253
361
 
254
- // Store the server's address so our threads know how to talk to it.
255
- SOCKADDR
256
- .set(sockaddr)
257
- .expect("sockaddr is only set in `trace_initialization`");
258
-
259
- // Send an initialization message to the server with some information about us. This will
260
- // only be sent once, and all future threads/subprocesses will be associated underneath us.
261
- let init = {
262
- let argv = {
263
- // Parse args out of `sys.argv`, since it excludes the Python executable and shows
264
- // the script's path first. This is much more complicated than using
265
- // `std::env::args`, but is more practical.
266
- let args_prefix = Python::with_gil(|py| -> Result<String> {
267
- let argv = py
268
- .import("sys")
269
- .and_then(|sys| sys.getattr("argv"))
270
- .wrap_err("Failed to access sys.argv")?;
271
-
272
- let args = argv
273
- .downcast::<pyo3::types::PyList>()
274
- .map_err(|e| eyre!("Failed to read sys.argv as list: {}", e))?;
275
-
276
- let mut full_args = args.iter().fold(String::new(), |s, arg| {
277
- // Fetch the next arg as a string if possible.
278
- let arg = arg
279
- .downcast::<pyo3::types::PyString>()
280
- .map_err(|e| eyre!("Failed to read sys.argv[x] as string: {}", e))
281
- .and_then(|arg| {
282
- arg.to_str()
283
- .wrap_err("Failed to read Rust string out of PyString")
284
- })
285
- .unwrap_or("<UNKNOWN>");
286
-
287
- s + " " + arg
288
- });
289
-
290
- // We'll display this in the UI, so ensure it's under (the rather
291
- // arbitrary) 100 chars.
292
- if full_args.chars().count() > 95 {
293
- full_args.truncate(95);
294
- full_args.push_str(" ...");
295
- }
362
+ // Send an initialization message to the server with some information about us. This will
363
+ // only be sent once, and all future threads/subprocesses will be associated underneath us.
364
+ let init = {
365
+ let python_version = Python::with_gil(|py| {
366
+ let version = py.version_info();
296
367
 
297
- // We'll be cloning this
298
- full_args.shrink_to_fit();
299
- Ok(full_args)
368
+ format!(
369
+ "Python {}.{}.{}",
370
+ version.major, version.minor, version.patch
371
+ )
372
+ });
373
+
374
+ let platform = Python::with_gil(|py| -> Result<String> {
375
+ let platform = py
376
+ .import("sys")
377
+ .and_then(|sys| sys.getattr("platform"))
378
+ .wrap_err("Failed to access sys.platform")?;
379
+
380
+ Ok(platform
381
+ .downcast::<pyo3::types::PyString>()
382
+ .map_err(|e| eyre!("Failed to read sys.platform as string: {}", e))?
383
+ .to_str()
384
+ .wrap_err("Failed to read Rust string out of PyString")?
385
+ .into())
300
386
  })
301
387
  .unwrap_or_else(|_err| {
302
- // TODO: tracing log that this happened
388
+ tracing::warn!("Failed to read sys.platform");
389
+
303
390
  "<UNKNOWN>".into()
304
391
  });
305
392
 
306
- // Store the args, since threads will need it as part of the registration message.
307
- ARGV.set(args_prefix.clone())
308
- .expect("ARGV is only set in `trace_initialization`");
309
-
310
- args_prefix
393
+ function_trace::TraceInitialization {
394
+ program_name: ARGV
395
+ .get()
396
+ .expect("ARGV is set earlier in trace_initialization")
397
+ .clone(),
398
+ program_version: format!("py-functiontrace {}", env!("PACKAGE_VERSION")),
399
+ lang_version: python_version,
400
+ platform,
401
+ time: trace_time(),
402
+ }
311
403
  };
312
404
 
313
- let python_version = Python::with_gil(|py| {
314
- let version = py.version_info();
405
+ let mut socket = message_initialize();
315
406
 
316
- format!(
317
- "Python {}.{}.{}",
318
- version.major, version.minor, version.patch
319
- )
320
- });
407
+ init.serialize(&mut rmp_serde::encode::Serializer::new(&mut socket))
408
+ .wrap_err("Failed to emit TraceInitialization message")?;
321
409
 
322
- let platform = Python::with_gil(|py| -> Result<String> {
323
- let platform = py
324
- .import("sys")
325
- .and_then(|sys| sys.getattr("platform"))
326
- .wrap_err("Failed to access sys.platform")?;
327
-
328
- Ok(platform
329
- .downcast::<pyo3::types::PyString>()
330
- .map_err(|e| eyre!("Failed to read sys.platform as string: {}", e))?
331
- .to_str()
332
- .wrap_err("Failed to read Rust string out of PyString")?
333
- .into())
334
- })
335
- .unwrap_or_else(|_err| {
336
- // TODO: tracing log that this happened
337
- "<UNKNOWN>".into()
338
- });
339
-
340
- function_trace::TraceInitialization {
341
- program_name: argv,
342
- program_version: format!("py-functiontrace {}", env!("PACKAGE_VERSION")),
343
- lang_version: python_version,
344
- platform,
345
- time: trace_time(),
346
- }
347
- };
348
-
349
- let mut socket = message_initialize();
350
-
351
- init.serialize(&mut rmp_serde::encode::Serializer::new(&mut socket))
352
- .wrap_err("Failed to emit TraceInitialization message")?;
353
-
354
- // Close the initialization socket to trigger `functiontrace-server` to
355
- // read our message, since otherwise it'll wait forever for more data.
356
- socket
357
- .shutdown(std::net::Shutdown::Both)
358
- .wrap_err("Failed to close initialization socket")?;
410
+ // Close the initialization socket to trigger `functiontrace-server` to
411
+ // read our message, since otherwise it'll wait forever for more data.
412
+ socket
413
+ .shutdown(std::net::Shutdown::Both)
414
+ .wrap_err("Failed to close initialization socket")?;
415
+ }
359
416
 
360
417
  ////////////////////////////////////////////////////////////////////////////
361
418
  // Teardown Configuration
@@ -431,7 +488,7 @@ mod _functiontrace_rs {
431
488
  }
432
489
 
433
490
  /// Called when the process is shutting down via `atexit`, notifies us to stop tracing and
434
- /// flush the buffer.
491
+ /// flush the buffer.
435
492
  ///
436
493
  /// NOTE: This is always called on the main thread.
437
494
  fn trace_terminate(
@@ -596,7 +653,7 @@ mod _functiontrace_rs {
596
653
  }
597
654
 
598
655
  /// Enable memory tracing for the current program. This is implemented by attaching wrapper
599
- /// versions of malloc/free/etc to all of the current allocators.
656
+ /// versions of malloc/free/etc to all the current allocators.
600
657
  ///
601
658
  /// NOTE: Memory tracing may have up to 40% overhead on traces with many small allocations, so
602
659
  /// is not enabled by default.
@@ -825,7 +882,12 @@ mod _functiontrace_rs {
825
882
  if let Ok(module) = m_module.downcast::<pyo3::types::PyString>() {
826
883
  module.to_cow().unwrap_or(Cow::Borrowed("<invalid string>"))
827
884
  } else {
828
- // TODO: tracing log this
885
+ if let Ok(typ) = m_module.get_type().qualname().map(|x| x.to_string()) {
886
+ tracing::warn!(type = typ, "Failed to read module name due to unexpected type");
887
+ } else {
888
+ tracing::warn!("Failed to read module name due to indecipherable type");
889
+ }
890
+
829
891
  Cow::Borrowed("<unknown module>")
830
892
  };
831
893
  function_trace::FunctionTrace::NativeCall {
@@ -881,13 +943,452 @@ mod _functiontrace_rs {
881
943
  // We unfortunately can't get these from `PyEval_SetProfile`, so we'd need to use a
882
944
  // different method to capture them.
883
945
  }
884
- _ => {
885
- // TODO: this shouldn't be possible - tracing log it so we can handle later.
946
+ v => {
947
+ tracing::error!(event = v, "Unexpected PyTrace message!");
886
948
  }
887
949
  };
888
950
 
889
951
  0
890
952
  }
953
+
954
+ ////////////////////////////////////////////////////////////////////////////
955
+ // Commercial support
956
+ ////////////////////////////////////////////////////////////////////////////
957
+
958
+ /// Load our configuration, including:
959
+ /// - Displaying a first-use message
960
+ /// - Creating the configuration file if non-existent
961
+ /// - Validating license keys
962
+ ///
963
+ /// We default to being permissive, and want to make this as non-invasive as possible for
964
+ /// non-commercial users.
965
+ #[pyfunction]
966
+ fn load_config() -> Result<()> {
967
+ use owo_colors::OwoColorize;
968
+
969
+ ////////////////////////////////////////////////////////////////////////
970
+ // Load or initialize config file
971
+ ////////////////////////////////////////////////////////////////////////
972
+ // Load configuration from the standard directory, or `~/.config` if the XDG spec isn't
973
+ // being followed.
974
+ let config_path = {
975
+ if let Some(config_path) = std::env::var_os("FUNCTIONTRACE_CONFIGFILE") {
976
+ // Allow `FUNCTIONTRACE_CONFIGFILE` as an override (mostly for tests, but it could
977
+ // be useful for Nix users or others).
978
+ std::path::PathBuf::from(config_path)
979
+ } else if let Some(mut config_dir) = std::env::var_os("XDG_CONFIG_HOME")
980
+ .map(std::path::PathBuf::from)
981
+ .or_else(|| {
982
+ std::env::home_dir().map(|mut home| {
983
+ home.push(".config");
984
+ home
985
+ })
986
+ })
987
+ {
988
+ config_dir.push("functiontrace.toml");
989
+ config_dir
990
+ } else {
991
+ // If we fail to find a reasonable config path, we'll just pretend everything is
992
+ // fine.
993
+ tracing::error!(
994
+ "Failed to find FunctionTrace config directory - proceeding without loading config"
995
+ );
996
+ return Ok(());
997
+ }
998
+ };
999
+
1000
+ let config = if let Ok(config) = std::fs::read(&config_path)
1001
+ .wrap_err("Failed to read config file")
1002
+ .and_then(|x| toml::from_slice::<Config>(&x).wrap_err("Failed to parse config file"))
1003
+ {
1004
+ // We successfully loaded an existing config file.
1005
+ config
1006
+ } else {
1007
+ // There's no config file, so assume this is the first time FunctionTrace has been
1008
+ // used and generate a new one.
1009
+ let logo = "\
1010
+ ███████╗██╗ ██╗███╗ ██╗ ██████╗████████╗██╗ ██████╗ ███╗ ██╗████████╗██████╗ █████╗ ██████╗███████╗\n\
1011
+ ██╔════╝██║ ██║████╗ ██║██╔════╝╚══██╔══╝██║██╔═══██╗████╗ ██║╚══██╔══╝██╔══██╗██╔══██╗██╔════╝██╔════╝\n\
1012
+ █████╗ ██║ ██║██╔██╗ ██║██║ ██║ ██║██║ ██║██╔██╗ ██║ ██║ ██████╔╝███████║██║ █████╗ \n\
1013
+ ██╔══╝ ██║ ██║██║╚██╗██║██║ ██║ ██║██║ ██║██║╚██╗██║ ██║ ██╔══██╗██╔══██║██║ ██╔══╝ \n\
1014
+ ██║ ╚██████╔╝██║ ╚████║╚██████╗ ██║ ██║╚██████╔╝██║ ╚████║ ██║ ██║ ██║██║ ██║╚██████╗███████╗\n\
1015
+ ╚═╝ ╚═════╝ ╚═╝ ╚═══╝ ╚═════╝ ╚═╝ ╚═╝ ╚═════╝ ╚═╝ ╚═══╝ ╚═╝ ╚═╝ ╚═╝╚═╝ ╚═╝ ╚═════╝╚══════╝\n";
1016
+
1017
+ // Render the logo like a flamechart using Python-esque colors.
1018
+ for (row, line) in logo.lines().enumerate() {
1019
+ for (x, c) in line.chars().enumerate() {
1020
+ // heights[x] represents the height of the flamechart at position x
1021
+ //
1022
+ // NOTE: These values were somewhat arbitrarily chosen, but the bottom row is
1023
+ // always covered.
1024
+ let heights = [
1025
+ 1, 1, 3, 3, 4, 4, 4, 4, 3, 2, 2, 2, 2, 2, 3, 4, 5, 5, 3, 3, 1, 2, 2, 3, 4,
1026
+ 4, 4, 4, 4, 4, 4, 3, 3, 3, 5, 5, 6, 6, 6, 5, 5, 3, 2, 2, 2, 3, 5, 5, 5, 5,
1027
+ 5, 2, 2, 4, 6, 6, 6, 3, 2, 3, 3, 5, 3, 3, 5, 5, 5, 5, 3, 3, 2, 2, 1, 1, 1,
1028
+ 1, 1, 2, 2, 2, 4, 4, 4, 3, 4, 3, 3, 3, 3, 2, 2, 5, 5, 5, 5, 5, 5, 4, 3, 2,
1029
+ 2, 2, 2, 3, 4, 5, 5, 0,
1030
+ ];
1031
+
1032
+ // We look at chart depth from top to bottom.
1033
+ let current_depth = 6 - row;
1034
+
1035
+ if current_depth > heights[x] {
1036
+ print!("{}", c.blue());
1037
+ } else {
1038
+ print!("{}", c.fg::<owo_colors::colors::xterm::FlushOrange>());
1039
+ }
1040
+ }
1041
+ println!();
1042
+ }
1043
+ println!();
1044
+ println!(
1045
+ "{}",
1046
+ "Welcome to FunctionTrace, a low overheard tracing profiler for Python".bold()
1047
+ );
1048
+ println!();
1049
+ println!(
1050
+ "Please see {} for instructions on setup and use.",
1051
+ "https://functiontrace.com".underline().cyan()
1052
+ );
1053
+ println!(
1054
+ "FunctionTrace will emit profile files ({}) that must be loaded by the",
1055
+ "functiontrace.*.tar.gz".dimmed()
1056
+ );
1057
+ println!(
1058
+ "Firefox Profiler ({}) for analysis.",
1059
+ "https://profiler.firefox.com/".underline().cyan()
1060
+ );
1061
+ println!();
1062
+
1063
+ let config = Config { license_key: None };
1064
+
1065
+ match std::fs::write(
1066
+ &config_path,
1067
+ toml::to_string_pretty(&config).wrap_err("Failed to serialize config file")?,
1068
+ ) {
1069
+ Ok(_) => {
1070
+ println!(
1071
+ "{}{}",
1072
+ format_args!("Default config saved to {}", config_path.display().bold())
1073
+ .dimmed(),
1074
+ ". This startup message will not be shown again.".dimmed()
1075
+ );
1076
+ println!();
1077
+ }
1078
+ Err(e) => {
1079
+ eprintln!(
1080
+ "Failed to create default FunctionTrace config file at {}: {e}",
1081
+ config_path.display().dimmed(),
1082
+ );
1083
+ }
1084
+ };
1085
+
1086
+ config
1087
+ };
1088
+
1089
+ // We need to write out the config if it's changed. Track the new and old configs so we
1090
+ // can tell.
1091
+ let mut new_config = config.clone();
1092
+
1093
+ ////////////////////////////////////////////////////////////////////////
1094
+ // Find their license key
1095
+ ////////////////////////////////////////////////////////////////////////
1096
+ // Demo keys have the same general shape as a UUID, but include the expiration time
1097
+ const DEMO_KEY_PREFIX: &str = "TRIALKEYX-";
1098
+ const DEMO_KEY_FORMAT: &str = "%Y-%m%d-%H%M-%S%z";
1099
+ const DEMO_KEY_SUFFIX: &str = "TRIAL";
1100
+
1101
+ let license_key = if let Some(license_key) = &config.license_key {
1102
+ license_key.clone()
1103
+ } else {
1104
+ println!("FunctionTrace is open-source and may freely be used noncommercially under");
1105
+ println!(
1106
+ "the Prosperity Public License 3.0 ({}).",
1107
+ "https://prosperitylicense.com/".underline().cyan()
1108
+ );
1109
+ println!();
1110
+
1111
+ /// Print a prompt, then read lines of input repeatedly until it successfully parses.
1112
+ fn read_input<T, F>(prompt: std::fmt::Arguments, parser: F) -> Result<T>
1113
+ where
1114
+ F: Fn(&String) -> Option<T>,
1115
+ {
1116
+ let mut input = String::new();
1117
+
1118
+ loop {
1119
+ print!("{}", prompt);
1120
+ std::io::stdout().flush()?;
1121
+
1122
+ input.clear();
1123
+ std::io::stdin().read_line(&mut input)?;
1124
+
1125
+ if let Some(result) = parser(&input) {
1126
+ return Ok(result);
1127
+ }
1128
+ Python::with_gil(|py| {
1129
+ // We need to check with the Python interpreter during loops, otherwise we
1130
+ // might not receive ctrl-c events.
1131
+ py.check_signals()
1132
+ })?;
1133
+ }
1134
+ }
1135
+
1136
+ let commercial = read_input(
1137
+ format_args!(
1138
+ "{} {}: ",
1139
+ "Are you using FunctionTrace commercially?".bold(),
1140
+ "[yes/no]".dimmed()
1141
+ ),
1142
+ |input| match input.chars().next() {
1143
+ Some('y') => Some(true),
1144
+ Some('n') => Some(false),
1145
+ _ => {
1146
+ eprintln!(
1147
+ "Failed to read {} or {} input.",
1148
+ "yes".dimmed(),
1149
+ "no".dimmed()
1150
+ );
1151
+ None
1152
+ }
1153
+ },
1154
+ )?;
1155
+
1156
+ let key = if !commercial {
1157
+ "NONCOMMERCIAL".to_string()
1158
+ } else {
1159
+ println!();
1160
+ println!(
1161
+ "The Prosperity Public License allows for a single 30 day trial period for commercial users."
1162
+ );
1163
+ let demo = read_input(
1164
+ format_args!(
1165
+ "{} {}: ",
1166
+ "Is this your company's first time using FunctionTrace?".bold(),
1167
+ "[yes/no]".dimmed()
1168
+ ),
1169
+ |input| match input.chars().next() {
1170
+ Some('y') => Some(true),
1171
+ Some('n') => Some(false),
1172
+ _ => {
1173
+ eprintln!(
1174
+ "Failed to read {} or {} input.",
1175
+ "yes".dimmed(),
1176
+ "no".dimmed()
1177
+ );
1178
+ None
1179
+ }
1180
+ },
1181
+ )?;
1182
+
1183
+ if demo {
1184
+ // They haven't used FunctionTrace yet, so we'll generate them a demo key
1185
+ // that's valid for 30 days.
1186
+ let expiration = Timestamp::now()
1187
+ .checked_add(Duration::from_secs(3600 * 24 * 30))
1188
+ .expect("30 days in the future is valid");
1189
+
1190
+ // Generate a key formatted like a UUID, but with the expiration time included
1191
+ format!(
1192
+ "{DEMO_KEY_PREFIX}{}{DEMO_KEY_SUFFIX}",
1193
+ expiration.strftime(DEMO_KEY_FORMAT)
1194
+ )
1195
+ } else {
1196
+ println!();
1197
+ println!(
1198
+ "Commercial users must purchase an alternative license via {}.",
1199
+ "https://functiontrace.com".underline().cyan()
1200
+ );
1201
+
1202
+ read_input(
1203
+ format_args!("{}: ", "Enter your FunctionTrace license key".bold()),
1204
+ |input| {
1205
+ if input.trim().chars().filter(|&c| c == '-').count() == 4 {
1206
+ // The key looks plausible. We'll validate it later.
1207
+ Some(input.trim().to_string())
1208
+ } else {
1209
+ eprintln!(
1210
+ "Invalid license format. FunctionTrace license keys look like {}.",
1211
+ "XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXXXXXX".dimmed()
1212
+ );
1213
+ None
1214
+ }
1215
+ },
1216
+ )?
1217
+ }
1218
+ };
1219
+
1220
+ new_config.license_key = Some(key.clone());
1221
+ key
1222
+ };
1223
+
1224
+ ////////////////////////////////////////////////////////////////////////
1225
+ // Validate the license key
1226
+ ////////////////////////////////////////////////////////////////////////
1227
+ let valid_license = if license_key == "NONCOMMERCIAL" {
1228
+ // The user has asserted that they're not using FunctionTrace for commercial
1229
+ // purposes. Nothing left to validate here.
1230
+ true
1231
+ } else if let Some(demo_key) = license_key
1232
+ .strip_prefix(DEMO_KEY_PREFIX)
1233
+ .and_then(|s| s.strip_suffix(DEMO_KEY_SUFFIX))
1234
+ {
1235
+ // The user has a demo key that's valid for 30 days from creation.
1236
+ let now = Timestamp::now();
1237
+
1238
+ if let Ok(valid_expiration) = Timestamp::strptime(DEMO_KEY_FORMAT, demo_key) {
1239
+ if now
1240
+ .until(valid_expiration)
1241
+ .map_or(false, |remaining| remaining.is_positive())
1242
+ {
1243
+ // The license is still valid
1244
+ eprintln!(
1245
+ "{}: You're using a trial license that expires on {}. Please visit {} to acquire a full license.",
1246
+ "FunctionTrace trial".yellow().bold(),
1247
+ valid_expiration.strftime("%Y-%m-%d").dimmed(),
1248
+ "https://functiontrace.com".underline().cyan(),
1249
+ );
1250
+ true
1251
+ } else {
1252
+ eprintln!(
1253
+ "{}: You're using a trial license that expired on {}. Please visit {} to acquire a full license.",
1254
+ "FunctionTrace trial expired".red().bold(),
1255
+ valid_expiration.strftime("%Y-%m-%d").dimmed(),
1256
+ "https://functiontrace.com".underline().cyan(),
1257
+ );
1258
+ false
1259
+ }
1260
+ } else {
1261
+ // This demo key doesn't parse, so it's presumably invalid.
1262
+ eprintln!(
1263
+ "{}: your company's trial has expired. Please visit {} to acquire a license.",
1264
+ "FunctionTrace trial expired".red().bold().dimmed(),
1265
+ "https://functiontrace.com".underline().cyan(),
1266
+ );
1267
+ false
1268
+ }
1269
+ } else {
1270
+ // The user has a commercial license key. Validate it against our server
1271
+ let client = reqwest::blocking::Client::new();
1272
+ if let Ok(license_check) = client
1273
+ .get("https://license.functiontrace.com")
1274
+ .header("functiontrace-key", &license_key)
1275
+ .timeout(
1276
+ // We set a short timeout, since the license server should respond within a few
1277
+ // ms and we don't want to block users if it can't be contacted.
1278
+ Duration::from_secs(2),
1279
+ )
1280
+ .send()
1281
+ {
1282
+ /// Represents useful information about a given license
1283
+ #[derive(Deserialize)]
1284
+ struct LicenseInfo {
1285
+ #[serde(rename = "Seats")]
1286
+ seats: u32,
1287
+ #[serde(rename = "Expires")]
1288
+ expires: Timestamp,
1289
+ }
1290
+
1291
+ match license_check.status() {
1292
+ reqwest::StatusCode::OK => {
1293
+ // This key is live!
1294
+ if let Ok(info) = license_check.json::<LicenseInfo>() {
1295
+ tracing::info!(
1296
+ seats = info.seats,
1297
+ expiration = info.expires.strftime("%F").to_string(),
1298
+ "Valid license loaded"
1299
+ );
1300
+ } else {
1301
+ tracing::info!(
1302
+ seats = "unknown",
1303
+ expiration = "unknown",
1304
+ "Valid license loaded"
1305
+ );
1306
+ }
1307
+ }
1308
+ reqwest::StatusCode::BAD_REQUEST => {
1309
+ // This license does not and has never existed.
1310
+ eprintln!(
1311
+ "{}: {} is not a valid FunctionTrace license.",
1312
+ "Invalid Functiontrace license".red().bold(),
1313
+ license_key.dimmed()
1314
+ );
1315
+
1316
+ // Mark this license key for deletion. We don't need to inform the user,
1317
+ // since it never could have worked.
1318
+ new_config.license_key = None;
1319
+ }
1320
+ reqwest::StatusCode::PAYMENT_REQUIRED => {
1321
+ // This license used to be valid but isn't anymore.
1322
+ if let Ok(info) = license_check.json::<LicenseInfo>() {
1323
+ eprintln!(
1324
+ "{}: {} expired on {}. Please visit {} to acquire a new license.",
1325
+ "FunctionTrace license expired".red().bold(),
1326
+ license_key.dimmed(),
1327
+ info.expires.strftime("%F").bold(),
1328
+ "https://functiontrace.com".underline().cyan(),
1329
+ );
1330
+ } else {
1331
+ eprintln!(
1332
+ "{}: {} has expired. Please visit {} to acquire a new license.",
1333
+ "FunctionTrace license expired".red().bold(),
1334
+ license_key.dimmed(),
1335
+ "https://functiontrace.com".underline().cyan(),
1336
+ );
1337
+ }
1338
+
1339
+ eprintln!(
1340
+ "Removing expired FunctionTrace license key from {}.",
1341
+ config_path.display().dimmed()
1342
+ );
1343
+ new_config.license_key = None;
1344
+ }
1345
+ reqwest::StatusCode::INTERNAL_SERVER_ERROR => {
1346
+ // The license server crashed?
1347
+ tracing::warn!(
1348
+ "License server returned an error - assuming license is valid"
1349
+ );
1350
+ }
1351
+ status => {
1352
+ tracing::warn!(
1353
+ code = status.as_u16(),
1354
+ "License server returned an unexpected code - assuming license is valid"
1355
+ );
1356
+ }
1357
+ };
1358
+
1359
+ // If there's still a license key stored in the config file, we must've validated
1360
+ // it.
1361
+ new_config.license_key.is_some()
1362
+ } else {
1363
+ // We didn't successfully talk to the license server. Assume everything is fine.
1364
+ true
1365
+ }
1366
+ };
1367
+
1368
+ ////////////////////////////////////////////////////////////////////////
1369
+ // Save the updated config
1370
+ ////////////////////////////////////////////////////////////////////////
1371
+ if new_config != config {
1372
+ // We updated the configuration file, so save it out.
1373
+ if let Err(e) = std::fs::write(
1374
+ &config_path,
1375
+ toml::to_string_pretty(&new_config).wrap_err("Failed to serialize config file")?,
1376
+ ) {
1377
+ eprintln!(
1378
+ "Failed to update FunctionTrace config file at {}: {e}",
1379
+ config_path.display().dimmed(),
1380
+ );
1381
+ };
1382
+ }
1383
+
1384
+ if !valid_license {
1385
+ // We've already printed information about what's happening, so it's safe to just kill
1386
+ // the process.
1387
+ std::process::exit(1);
1388
+ }
1389
+
1390
+ Ok(())
1391
+ }
891
1392
  }
892
1393
 
893
1394
  /// Allocation hooks and helpers for tracing Python memory allocator operations.
@@ -997,7 +1498,7 @@ mod hooks {
997
1498
  use color_eyre::eyre::{Result, WrapErr};
998
1499
  use functiontrace_server::function_trace;
999
1500
  use pyo3::ffi::PyObject;
1000
- use pyo3::prelude::*; // TODO: What's in here?
1501
+ use pyo3::prelude::*;
1001
1502
  use serde::Serialize;
1002
1503
  use std::borrow::Cow;
1003
1504
  use std::sync::OnceLock;
@@ -1148,7 +1649,7 @@ mod hooks {
1148
1649
  .get()
1149
1650
  .expect("Wrapper functions can't be called until hooked");
1150
1651
 
1151
- Python::with_gil(|py| -> Result<()> {
1652
+ let proxied = Python::with_gil(|py| -> Result<()> {
1152
1653
  if let Some(mut writer) = crate::_functiontrace_rs::thread_state() {
1153
1654
  let args = unsafe {
1154
1655
  // SAFETY: `args` represents a valid argument tuple
@@ -1167,10 +1668,15 @@ mod hooks {
1167
1668
 
1168
1669
  Ok(())
1169
1670
  })
1170
- .wrap_err_with(|| format!("FunctionTrace proxy logging for {} failed", hook.target))
1171
- // TODO: This shouldn't be catastrophic - we can log it and continue on
1172
- // instead once we introduce `tracing`.
1173
- .expect("Proxying should never fail");
1671
+ .wrap_err_with(|| format!("FunctionTrace proxy logging for {} failed", hook.target));
1672
+
1673
+ if let Err(e) = proxied {
1674
+ tracing::error!(
1675
+ error = e.to_string(),
1676
+ function = hook.target,
1677
+ "Failed to properly handle hooked function"
1678
+ );
1679
+ }
1174
1680
 
1175
1681
  unsafe {
1176
1682
  // SAFETY: This is a direct translation of the original call we're proxying
File without changes
File without changes
File without changes
File without changes
File without changes