python: allow lines configuration (#199)

Author: korniltsev

pyroscope_backends/pyroscope_pyspy/src/lib.rs

@@ -1,4 +1,4 @@
-use py_spy::{config::Config, sampler::Sampler};
+use py_spy::{config::Config, sampler::Sampler, Pid};
 use pyroscope::{
     backend::{
         Backend, BackendConfig, BackendImpl, BackendUninitialized, Report, Rule, Ruleset,
@@ -14,7 +14,6 @@ use std::{
     },
     thread::JoinHandle,
 };
-use py_spy::Pid;
 
 const LOG_TAG: &str = "Pyroscope::Pyspy";
 
@@ -47,6 +46,8 @@ pub struct PyspyConfig {
     gil_only: bool,
     /// Profile native C extensions
     native: bool,
+
+    line_no: py_spy::config::LineNo,
 }
 
 impl Default for PyspyConfig {
@@ -61,6 +62,7 @@ impl Default for PyspyConfig {
             oncpu: false,
             gil_only: false,
             native: false,
+            line_no: py_spy::config::LineNo::LastInstruction,
         }
     }
 }
@@ -160,6 +162,13 @@ impl PyspyConfig {
     pub fn native(self, native: bool) -> Self {
         PyspyConfig { native, ..self }
     }
+
+    pub fn line_no(self, line_no: LineNo) -> Self {
+        PyspyConfig {
+            line_no: line_no.into(),
+            ..self
+        }
+    }
 }
 
 /// Pyspy Backend
@@ -258,6 +267,7 @@ impl Backend for Pyspy {
             include_thread_ids: true,
             subprocesses: self.config.detect_subprocesses,
             gil_only: self.config.gil_only,
+            lineno: self.config.line_no.clone().into(),
             duration,
             ..Config::default()
         });
@@ -403,3 +413,22 @@ impl From<(py_spy::StackTrace, &BackendConfig)> for StackTraceWrapper {
         StackTraceWrapper(stacktrace)
     }
 }
+
+// This is because we do not depdend on pyspy in ffikit
+// We probably need to change that and remove the glue
+// Not today
+pub enum LineNo {
+    LastInstruction = 0,
+    First = 1,
+    NoLine = 2,
+}
+
+impl Into<py_spy::config::LineNo> for LineNo {
+    fn into(self) -> py_spy::config::LineNo {
+        match self {
+            LineNo::LastInstruction => py_spy::config::LineNo::LastInstruction,
+            LineNo::First => py_spy::config::LineNo::First,
+            LineNo::NoLine => py_spy::config::LineNo::NoLine,
+        }
+    }
+}

pyroscope_ffi/python/lib/Cargo.toml

@@ -1,4 +1,5 @@
 [package]
+# todo rename this, include python in the name
 name = "pyroscope_ffi"
 version = "0.1.0"
 authors = ["Abid Omar <[email protected]>"]

pyroscope_ffi/python/lib/cbindgen.toml

@@ -17,6 +17,11 @@ braces = "SameLine"
 tab_width = 2
 line_length = 80
 
+
 [parse]
 # Do not parse dependent crates
 parse_deps = false
+include = ["LineNoFFI"]
+
+[export]
+include = ["LineNoFFI"]

pyroscope_ffi/python/lib/include/pyroscope_ffi.h

@@ -12,6 +12,12 @@
 #include <stdint.h>
 #include <stdlib.h>
 
+typedef enum {
+  LastInstruction = 0,
+  First = 1,
+  NoLine = 2,
+} LineNo;
+
 bool initialize_logging(uint32_t logging_level);
 
 bool initialize_agent(const char *application_name,
@@ -28,7 +34,8 @@ bool initialize_agent(const char *application_name,
                       bool report_thread_name,
                       const char *tags,
                       const char *tenant_id,
-                      const char *http_headers_json);
+                      const char *http_headers_json,
+                      LineNo line_no);
 
 bool drop_agent(void);
 

pyroscope_ffi/python/lib/src/lib.rs

@@ -1,7 +1,7 @@
 use ffikit::Signal;
 use pyroscope::backend::Tag;
-use pyroscope::PyroscopeAgent;
 use pyroscope::pyroscope::ReportEncoding;
+use pyroscope::PyroscopeAgent;
 use pyroscope_pyspy::{pyspy_backend, PyspyConfig};
 use std::collections::hash_map::DefaultHasher;
 use std::ffi::CStr;
@@ -15,22 +15,22 @@ pub extern "C" fn initialize_logging(logging_level: u32) -> bool {
     // Force rustc to display the log messages in the console.
     match logging_level {
         50 => {
-            std::env::set_var("RUST_LOG", "off");
+            unsafe { std::env::set_var("RUST_LOG", "off") };
         }
         40 => {
-            std::env::set_var("RUST_LOG", "error");
+            unsafe { std::env::set_var("RUST_LOG", "error") };
         }
         30 => {
-            std::env::set_var("RUST_LOG", "warn");
+            unsafe { std::env::set_var("RUST_LOG", "warn") };
         }
         20 => {
-            std::env::set_var("RUST_LOG", "info");
+            unsafe { std::env::set_var("RUST_LOG", "info") };
         }
         10 => {
-            std::env::set_var("RUST_LOG", "debug");
+            unsafe { std::env::set_var("RUST_LOG", "debug") };
         }
         _ => {
-            std::env::set_var("RUST_LOG", "debug");
+            unsafe { std::env::set_var("RUST_LOG", "debug") };
         }
     }
 
@@ -41,21 +41,11 @@ pub extern "C" fn initialize_logging(logging_level: u32) -> bool {
 
 #[no_mangle]
 pub extern "C" fn initialize_agent(
-    application_name: *const c_char,
-    server_address: *const c_char,
-    auth_token: *const c_char,
-    basic_auth_username: *const c_char,
-    basic_auth_password: *const c_char,
-    sample_rate: u32,
-    detect_subprocesses: bool,
-    oncpu: bool,
-    gil_only: bool,
-    report_pid: bool,
-    report_thread_id: bool,
-    report_thread_name: bool,
-    tags: *const c_char,
-    tenant_id: *const c_char,
-    http_headers_json: *const c_char,
+    application_name: *const c_char, server_address: *const c_char, auth_token: *const c_char,
+    basic_auth_username: *const c_char, basic_auth_password: *const c_char, sample_rate: u32,
+    detect_subprocesses: bool, oncpu: bool, gil_only: bool, report_pid: bool,
+    report_thread_id: bool, report_thread_name: bool, tags: *const c_char,
+    tenant_id: *const c_char, http_headers_json: *const c_char, line_no: LineNo
 ) -> bool {
     // Initialize FFIKit
     let recv = ffikit::initialize_ffi().unwrap();
@@ -117,6 +107,7 @@ pub extern "C" fn initialize_agent(
         .detect_subprocesses(detect_subprocesses)
         .oncpu(oncpu)
         .native(false)
+        .line_no(line_no.into())
         .gil_only(gil_only);
 
     // Report the PID.
@@ -162,17 +153,15 @@ pub extern "C" fn initialize_agent(
         Ok(http_headers) => {
             agent_builder = agent_builder.http_headers(http_headers);
         }
-        Err(e) => {
-            match e {
-                pyroscope::PyroscopeError::Json(e) => {
-                    log::error!(target: LOG_TAG, "parse_http_headers_json error {}", e);
-                }
-                pyroscope::PyroscopeError::AdHoc(e) => {
-                    log::error!(target: LOG_TAG, "parse_http_headers_json {}", e);
-                }
-                _ => {}
+        Err(e) => match e {
+            pyroscope::PyroscopeError::Json(e) => {
+                log::error!(target: LOG_TAG, "parse_http_headers_json error {}", e);
             }
-        }
+            pyroscope::PyroscopeError::AdHoc(e) => {
+                log::error!(target: LOG_TAG, "parse_http_headers_json {}", e);
+            }
+            _ => {}
+        },
     }
 
     // Build the agent.
@@ -291,3 +280,21 @@ fn string_to_tags<'a>(tags: &'a str) -> Vec<(&'a str, &'a str)> {
 
     tags_vec
 }
+
+#[repr(C)]
+#[derive(Debug)]
+pub enum LineNo {
+    LastInstruction = 0,
+    First = 1,
+    NoLine = 2,
+}
+
+impl Into<pyroscope_pyspy::LineNo> for LineNo {
+    fn into(self) -> pyroscope_pyspy::LineNo {
+        match self {
+            LineNo::LastInstruction => pyroscope_pyspy::LineNo::LastInstruction,
+            LineNo::First => pyroscope_pyspy::LineNo::First,
+            LineNo::NoLine => pyroscope_pyspy::LineNo::NoLine,
+        }
+    }
+}

pyroscope_ffi/python/pyroscope/__init__.py

@@ -2,10 +2,15 @@
 import warnings
 import logging
 import json
-from collections import namedtuple
-from pyroscope._native import ffi, lib
-from contextlib import contextmanager 
+from enum import Enum
 
+from pyroscope._native import lib
+from contextlib import contextmanager
+
+class LineNo(Enum):
+    LastInstruction = 0
+    First = 1
+    NoLine = 2
 
 def configure(
         app_name=None,
@@ -26,6 +31,7 @@ def configure(
         tags=None,
         tenant_id="",
         http_headers=None,
+        line_no=LineNo.LastInstruction,
 ):
 
     if app_name is not None:
@@ -56,7 +62,8 @@ def configure(
         tags_to_string(tags).encode("UTF-8"),
         (tenant_id or "").encode("UTF-8"),
         http_headers_to_json(http_headers).encode("UTF-8"),
-)
+        line_no.value
+    )
 
 def shutdown():
     drop = lib.drop_agent()