Add Ability to Activate per Sample Timestamps

lightswitch-proto/src/profile.rs

@@ -478,6 +478,43 @@ mod tests {
         pprof.build();
     }
 
+    #[test]
+    fn collected_at_label_round_trips_through_pprof() {
+        let mut pprof = PprofBuilder::new(SystemTime::now(), Duration::from_secs(1), 19);
+        let mapping_id = pprof.add_mapping(0x1000, 0x2000, 0, "test.so", "build-id-abc");
+        let location_id = pprof.add_location(0x1500, mapping_id, vec![]);
+
+        let ms_timestamp: i64 = 1_748_865_070_123;
+        let label = pprof.new_label(
+            "collected_at",
+            LabelStringOrNumber::Number(ms_timestamp, "milliseconds".into()),
+        );
+        pprof.add_sample(vec![location_id], 1, &[label]);
+
+        let profile = pprof.build();
+
+        // Verify in memory
+        assert_eq!(profile.sample.len(), 1);
+        let sample = &profile.sample[0];
+        assert_eq!(sample.label.len(), 1);
+        let lbl = &sample.label[0];
+        assert_eq!(lbl.num, ms_timestamp);
+        assert_ne!(lbl.num_unit, 0, "num_unit must be set for numeric labels");
+        let key_str = &profile.string_table[lbl.key as usize];
+        assert_eq!(key_str, "collected_at");
+        let unit_str = &profile.string_table[lbl.num_unit as usize];
+        assert_eq!(unit_str, "milliseconds");
+
+        // Verify after encode/decode round-trip
+        let mut buf = bytes::BytesMut::new();
+        profile.encode(&mut buf).expect("encode failed");
+        let decoded = pprof::Profile::decode(buf.freeze()).expect("decode failed");
+        let decoded_lbl = &decoded.sample[0].label[0];
+        assert_eq!(decoded_lbl.num, ms_timestamp);
+        let decoded_key = &decoded.string_table[decoded_lbl.key as usize];
+        assert_eq!(decoded_key, "collected_at");
+    }
+
     #[test]
     fn test_encode_decode() {
         let mut pprof = PprofBuilder::new(SystemTime::now(), Duration::from_secs(5), 27);

src/aggregator.rs

@@ -167,6 +167,49 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_aggregate_same_stack_different_ms_bucket_stays_separate() {
+        let stack = RawSample {
+            pid: 1,
+            tid: 2,
+            collected_at: 1_000_000_000, // 1000 ms
+            ustack: vec![0xdead],
+            kstack: vec![0xbeef],
+        };
+        // 2 ms later = different ms bucket: must NOT be merged
+        let later = RawSample {
+            collected_at: 1_002_000_000,
+            ..stack.clone()
+        };
+        let aggregator = Aggregator::default();
+        let result = aggregator.aggregate(vec![stack, later]);
+        assert_eq!(
+            result.len(),
+            2,
+            "samples in different ms buckets must stay separate"
+        );
+    }
+
+    #[test]
+    fn test_aggregate_same_stack_within_ms_bucket_merges() {
+        let stack = RawSample {
+            pid: 1,
+            tid: 2,
+            collected_at: 1_000_000_000, // 1000 ms
+            ustack: vec![0xdead],
+            kstack: vec![0xbeef],
+        };
+        // 999 µs later = same ms bucket: must be merged (count becomes 2)
+        let almost_same = RawSample {
+            collected_at: 1_000_999_000,
+            ..stack.clone()
+        };
+        let aggregator = Aggregator::default();
+        let result = aggregator.aggregate(vec![stack, almost_same]);
+        assert_eq!(result.len(), 1, "samples in the same ms bucket must merge");
+        assert_eq!(result[0].count, 2);
+    }
+
     #[test]
     fn test_aggregate_same_stack_traces_different_pid_tid() {
         let ustack = vec![0xffff, 0xdeadbeef];

src/bpf_objects.rs

@@ -262,7 +262,6 @@ impl Bpf {
             .lightswitch_config
             .use_task_pt_regs_helper
             .write(profiler_config.use_task_pt_regs_helper);
-
         // Disable BTF helpers if a BTF custom path is selected since
         // it will not load in machines that don't have one. TODO add override?
         rodata

src/cli/args.rs

@@ -185,6 +185,11 @@ pub(crate) struct CliArgs {
         help = "Read metadata for all threads when a new process is detected. This might be slow in older kernels."
     )]
     pub(crate) preload_thread_metadata: bool,
+    #[arg(
+        long,
+        help = "Attach a per-sample collected_at timestamp (ms) label to pprof output"
+    )]
+    pub(crate) enable_ts_per_sample: bool,
     #[arg(long, help = "Launch live flamegraph TUI")]
     pub(crate) live: bool,
     #[arg(

src/cli/main.rs

@@ -272,6 +272,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                 .available_bpf_features
                 .has_non_prealloc_hash_maps_in_tracing,
         userspace_pid_ns_level: system_info.available_bpf_features.userspace_pid_ns_level,
+        ts_per_sample: args.enable_ts_per_sample,
         ..Default::default()
     };
 
@@ -350,6 +351,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                 ProfilerConfig::default().session_duration,
                 args.sample_freq,
                 metadata_provider.clone(),
+                args.enable_ts_per_sample,
             )),
             ProfileSender::Pyroscope => Box::new(PyroscopeCollector::new(
                 args.symbolizer == Symbolizer::Local,
@@ -360,6 +362,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                 args.sample_freq,
                 metadata_provider.clone(),
                 args.node_name.clone().unwrap_or_default(),
+                args.enable_ts_per_sample,
             )),
             ProfileSender::Firefox => Box::new(FirefoxProfilerCollector::new(
                 "firefox-profiler.json",
@@ -431,6 +434,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                 &metadata_provider,
                 profile_duration,
                 args.sample_freq,
+                args.enable_ts_per_sample,
             );
             pprof_profile.encode(&mut buffer).unwrap();
             let profile_name = args.profile_name.unwrap_or_else(|| "profile.pb".into());
@@ -510,7 +514,7 @@ mod tests {
         cmd.arg("--help");
         cmd.assert().success();
         let actual = String::from_utf8(cmd.unwrap().stdout).unwrap();
-        insta::assert_yaml_snapshot!(actual, @r#""Usage: lightswitch [OPTIONS] [COMMAND]\n\nCommands:\n  object-info  \n  show-unwind  \n  system-info  \n  server       \n  help         Print this message or the help of the given subcommand(s)\n\nOptions:\n      --pids <PIDS>\n          Specific PIDs to profile\n\n  -D, --duration <DURATION>\n          How long this agent will run in seconds\n          \n          [default: 18446744073709551615]\n\n      --libbpf-debug\n          Enable libbpf logs. This includes the BPF verifier output\n\n      --bpf-logging\n          Enable BPF programs logging\n\n      --btf-custom-path <BTF_CUSTOM_PATH>\n          Alternative BTF file\n\n      --logging <LOGGING>\n          Set lightswitch's logging level\n          \n          [default: info]\n          [possible values: trace, debug, info, warn, error]\n\n      --sample-freq <SAMPLE_FREQ_IN_HZ>\n          Per-CPU Sampling Frequency in Hz\n          \n          [default: 19]\n\n      --profile-format <PROFILE_FORMAT>\n          Output file for Flame Graph in SVG format\n          \n          [default: flame-graph]\n          [possible values: none, flame-graph, pprof]\n\n      --flamegraph-aggregation <FLAMEGRAPH_AGGREGATION>\n          What information to show in the flamegraph. Won't do anything for other profile formats\n          \n          [default: function]\n          [possible values: function, all]\n\n      --profile-path <PROFILE_PATH>\n          Path for the generated profile\n\n      --profile-name <PROFILE_NAME>\n          Name for the generated profile\n\n      --sender <SENDER>\n          Where to write the profile\n\n          Possible values:\n          - none:       Discard the profile. Used for kernel tests\n          - local-disk\n          - remote\n          - pyroscope\n          - firefox\n          \n          [default: local-disk]\n\n      --server-url <SERVER_URL>\n          \n\n      --token <TOKEN>\n          \n\n      --pyroscope-app-name <PYROSCOPE_APP_NAME>\n          Application name for Pyroscope\n          \n          [default: lightswitch]\n\n      --pyroscope-tenant-id <PYROSCOPE_TENANT_ID>\n          Tenant ID for multi-tenant Pyroscope (sets `X-Scope-OrgID` header)\n\n      --perf-buffer-bytes <PERF_BUFFER_BYTES>\n          Size of each profiler perf buffer, in bytes (must be a power of 2)\n          \n          [default: 524288]\n\n      --mapsize-info\n          Print eBPF map sizes after creation\n\n      --mapsize-rate-limits <MAPSIZE_RATE_LIMITS>\n          max number of rate limit entries\n          \n          [default: 5000]\n\n      --exclude-self\n          Do not profile the profiler (myself)\n\n      --symbolizer <SYMBOLIZER>\n          [default: local]\n          [possible values: local, none]\n\n      --debug-info-backend <DEBUG_INFO_BACKEND>\n          [default: none]\n          [possible values: none, copy, remote]\n\n      --max-native-unwind-info-size-mb <MAX_NATIVE_UNWIND_INFO_SIZE_MB>\n          approximate max size in megabytes used for the BPF maps that hold unwind information\n          \n          [default: 2147483647]\n\n      --enable-deadlock-detector\n          enable parking_lot's deadlock detector\n\n      --cache-dir-base <CACHE_DIR_BASE>\n          [default: /tmp]\n\n      --killswitch-path-override <KILLSWITCH_PATH_OVERRIDE>\n          Override the default path to the killswitch file (/tmp/lightswitch/killswitch) which prevents the profiler from starting\n\n      --unsafe-start\n          Force the profiler to start even if the system killswitch is enabled\n\n      --force-perf-buffer\n          force perf buffers even if ring buffers can be used\n\n      --no-prealloc-bpf-hash-maps\n          Do not preallocate BPF hash maps\n\n      --preload-thread-metadata\n          Read metadata for all threads when a new process is detected. This might be slow in older kernels.\n\n      --live\n          Launch live flamegraph TUI\n\n      --kubernetes\n          Enable Kubernetes pod-level profiling\n\n      --node-name <NODE_NAME>\n          Node name for Kubernetes metadata (required with --kubernetes)\n\n  -h, --help\n          Print help (see a summary with '-h')\n\n  -V, --version\n          Print version\n""#);
+        insta::assert_yaml_snapshot!(actual, @r#""Usage: lightswitch [OPTIONS] [COMMAND]\n\nCommands:\n  object-info  \n  show-unwind  \n  system-info  \n  server       \n  help         Print this message or the help of the given subcommand(s)\n\nOptions:\n      --pids <PIDS>\n          Specific PIDs to profile\n\n  -D, --duration <DURATION>\n          How long this agent will run in seconds\n          \n          [default: 18446744073709551615]\n\n      --libbpf-debug\n          Enable libbpf logs. This includes the BPF verifier output\n\n      --bpf-logging\n          Enable BPF programs logging\n\n      --btf-custom-path <BTF_CUSTOM_PATH>\n          Alternative BTF file\n\n      --logging <LOGGING>\n          Set lightswitch's logging level\n          \n          [default: info]\n          [possible values: trace, debug, info, warn, error]\n\n      --sample-freq <SAMPLE_FREQ_IN_HZ>\n          Per-CPU Sampling Frequency in Hz\n          \n          [default: 19]\n\n      --profile-format <PROFILE_FORMAT>\n          Output file for Flame Graph in SVG format\n          \n          [default: flame-graph]\n          [possible values: none, flame-graph, pprof]\n\n      --flamegraph-aggregation <FLAMEGRAPH_AGGREGATION>\n          What information to show in the flamegraph. Won't do anything for other profile formats\n          \n          [default: function]\n          [possible values: function, all]\n\n      --profile-path <PROFILE_PATH>\n          Path for the generated profile\n\n      --profile-name <PROFILE_NAME>\n          Name for the generated profile\n\n      --sender <SENDER>\n          Where to write the profile\n\n          Possible values:\n          - none:       Discard the profile. Used for kernel tests\n          - local-disk\n          - remote\n          - pyroscope\n          - firefox\n          \n          [default: local-disk]\n\n      --server-url <SERVER_URL>\n          \n\n      --token <TOKEN>\n          \n\n      --pyroscope-app-name <PYROSCOPE_APP_NAME>\n          Application name for Pyroscope\n          \n          [default: lightswitch]\n\n      --pyroscope-tenant-id <PYROSCOPE_TENANT_ID>\n          Tenant ID for multi-tenant Pyroscope (sets `X-Scope-OrgID` header)\n\n      --perf-buffer-bytes <PERF_BUFFER_BYTES>\n          Size of each profiler perf buffer, in bytes (must be a power of 2)\n          \n          [default: 524288]\n\n      --mapsize-info\n          Print eBPF map sizes after creation\n\n      --mapsize-rate-limits <MAPSIZE_RATE_LIMITS>\n          max number of rate limit entries\n          \n          [default: 5000]\n\n      --exclude-self\n          Do not profile the profiler (myself)\n\n      --symbolizer <SYMBOLIZER>\n          [default: local]\n          [possible values: local, none]\n\n      --debug-info-backend <DEBUG_INFO_BACKEND>\n          [default: none]\n          [possible values: none, copy, remote]\n\n      --max-native-unwind-info-size-mb <MAX_NATIVE_UNWIND_INFO_SIZE_MB>\n          approximate max size in megabytes used for the BPF maps that hold unwind information\n          \n          [default: 2147483647]\n\n      --enable-deadlock-detector\n          enable parking_lot's deadlock detector\n\n      --cache-dir-base <CACHE_DIR_BASE>\n          [default: /bb/data/tmp]\n\n      --killswitch-path-override <KILLSWITCH_PATH_OVERRIDE>\n          Override the default path to the killswitch file (/tmp/lightswitch/killswitch) which prevents the profiler from starting\n\n      --unsafe-start\n          Force the profiler to start even if the system killswitch is enabled\n\n      --force-perf-buffer\n          force perf buffers even if ring buffers can be used\n\n      --no-prealloc-bpf-hash-maps\n          Do not preallocate BPF hash maps\n\n      --preload-thread-metadata\n          Read metadata for all threads when a new process is detected. This might be slow in older kernels.\n\n      --enable-ts-per-sample\n          Attach a per-sample collected_at timestamp (ms) label to pprof output\n\n      --live\n          Launch live flamegraph TUI\n\n      --kubernetes\n          Enable Kubernetes pod-level profiling\n\n      --node-name <NODE_NAME>\n          Node name for Kubernetes metadata (required with --kubernetes)\n\n  -h, --help\n          Print help (see a summary with '-h')\n\n  -V, --version\n          Print version\n""#);
     }
 
     #[rstest]

src/collector.rs

@@ -89,6 +89,7 @@ pub struct StreamingCollector {
     procs: HashMap<i32, ProcessInfo>,
     objs: HashMap<ExecutableId, ObjectFileInfo>,
     metadata_provider: ThreadSafeGlobalMetadataProvider,
+    ts_per_sample: bool,
 }
 
 impl StreamingCollector {
@@ -99,6 +100,7 @@ impl StreamingCollector {
         profile_duration: Duration,
         profile_frequency_hz: u64,
         metadata_provider: ThreadSafeGlobalMetadataProvider,
+        ts_per_sample: bool,
     ) -> Self {
         let client_builder = reqwest::blocking::Client::builder().timeout(Duration::from_secs(30));
         let client = client_builder
@@ -113,6 +115,7 @@ impl StreamingCollector {
             profile_duration,
             profile_frequency_hz,
             metadata_provider,
+            ts_per_sample,
             ..Default::default()
         }
     }
@@ -140,6 +143,7 @@ impl Collector for StreamingCollector {
             &self.metadata_provider,
             self.profile_duration,
             self.profile_frequency_hz,
+            self.ts_per_sample,
         );
 
         let mut request = self
@@ -182,6 +186,7 @@ pub struct PyroscopeCollector {
     objs: HashMap<ExecutableId, ObjectFileInfo>,
     metadata_provider: ThreadSafeGlobalMetadataProvider,
     node_name: String,
+    ts_per_sample: bool,
 }
 
 impl PyroscopeCollector {
@@ -195,6 +200,7 @@ impl PyroscopeCollector {
         profile_frequency_hz: u64,
         metadata_provider: ThreadSafeGlobalMetadataProvider,
         node_name: String,
+        ts_per_sample: bool,
     ) -> Self {
         let push_url = format!("{}/push.v1.PusherService/Push", server_url);
         let client_builder = reqwest::blocking::Client::builder().timeout(Duration::from_secs(30));
@@ -209,6 +215,7 @@ impl PyroscopeCollector {
             profile_frequency_hz,
             metadata_provider,
             node_name,
+            ts_per_sample,
             ..Default::default()
         }
     }
@@ -315,6 +322,7 @@ impl PyroscopeCollector {
             &self.metadata_provider,
             self.profile_duration,
             self.profile_frequency_hz,
+            self.ts_per_sample,
         );
 
         let mut buffer = Vec::new();

src/profile/convert.rs

@@ -45,6 +45,7 @@ pub fn to_pprof(
     metadata_provider: &ThreadSafeGlobalMetadataProvider,
     profile_duration: Duration,
     profile_frequency_hz: u64,
+    ts_per_sample: bool,
 ) -> pprof::Profile {
     // Not exactly when the profile session really started but works for now.
     let profile_start = SystemTime::now();
@@ -199,7 +200,7 @@ pub fn to_pprof(
             // todo: maybe append an error frame for debugging?
             continue;
         };
-        let labels = task_to_labels.entry(task_key).or_insert_with(|| {
+        let task_labels = task_to_labels.entry(task_key).or_insert_with(|| {
             let metadata = metadata_provider.lock().unwrap().get_metadata(task_key);
             metadata
                 .into_iter()
@@ -212,7 +213,15 @@ pub fn to_pprof(
                 })
                 .collect()
         });
-        pprof.add_sample(location_ids, sample.count as i64, labels);
+        let mut sample_labels = task_labels.clone();
+        if ts_per_sample {
+            let timestamp_label = pprof.new_label(
+                "collected_at",
+                LabelStringOrNumber::Number(sample.collected_at as i64, "milliseconds".into()),
+            );
+            sample_labels.push(timestamp_label);
+        }
+        pprof.add_sample(location_ids, sample.count as i64, &sample_labels);
     }
 
     pprof.build()
@@ -324,6 +333,7 @@ pub fn symbolize_profile(
             pid: sample.pid,
             tid: sample.tid,
             count: sample.count,
+            collected_at: sample.collected_at,
             ustack: symbolize_user_stack(
                 &addresses_per_sample,
                 procs,