From 2e0fccdbf909293b6933efed46ebb2cb60ab7e8e Mon Sep 17 00:00:00 2001
From: Brian Martin <bmartin@twitter.com>
Date: Mon, 24 Jun 2019 11:43:04 -0700
Subject: [PATCH] adds signal handler to examples

Adds signal handlers for SIGINT and SIGTEM to the examples so that
the `Drop` implementations are invoked following ctrl-c.
---
 Cargo.toml            |  7 ++++---
 examples/opensnoop.rs | 20 +++++++++++++++-----
 examples/softirqs.rs  | 16 +++++++++++++---
 examples/strlen.rs    | 17 ++++++++++++++---
 4 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index a469cbf4..6a96d2b9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,14 +12,15 @@ homepage = "https://github.com/rust-bpf/rust-bcc"
 edition = '2018'
 
 [dependencies]
-libc = "0.2.55"
 bcc-sys = "0.9.2"
-failure = "0.1.5"
 byteorder = "1.3.1"
+failure = "0.1.5"
+libc = "0.2.55"
 
 [dev-dependencies]
-lazy_static = "1.3.0"
 clap = "2.33.0"
+ctrlc = "3.1.3"
+lazy_static = "1.3.0"
 
 [features]
 static = ["bcc-sys/static"]
diff --git a/examples/opensnoop.rs b/examples/opensnoop.rs
index 7948f0b3..bf9f6a7e 100644
--- a/examples/opensnoop.rs
+++ b/examples/opensnoop.rs
@@ -3,12 +3,14 @@ extern crate byteorder;
 extern crate failure;
 extern crate libc;
 
-use std::ptr;
-
 use bcc::core::BPF;
 use bcc::perf::init_perf_map;
 use failure::Error;
 
+use core::sync::atomic::{AtomicBool, Ordering};
+use std::ptr;
+use std::sync::Arc;
+
 /*
  * Basic Rust clone of `opensnoop`, from the iovisor tools.
  * https://github.com/iovisor/bcc/blob/master/tools/opensnoop.py
@@ -31,7 +33,7 @@ struct data_t {
     fname: [u8; 255], // NAME_MAX
 }
 
-fn do_main() -> Result<(), Error> {
+fn do_main(runnable: Arc<AtomicBool>) -> Result<(), Error> {
     let code = include_str!("opensnoop.c");
     // compile the above BPF code!
     let mut module = BPF::new(code)?;
@@ -47,9 +49,10 @@ fn do_main() -> Result<(), Error> {
     // print a header
     println!("{:-7} {:-16} {}", "PID", "COMM", "FILENAME");
     // this `.poll()` loop is what makes our callback get called
-    loop {
+    while runnable.load(Ordering::SeqCst) {
         perf_map.poll(200);
     }
+    Ok(())
 }
 
 fn perf_data_callback() -> Box<FnMut(&[u8]) + Send> {
@@ -77,7 +80,14 @@ fn get_string(x: &[u8]) -> String {
 }
 
 fn main() {
-    match do_main() {
+    let runnable = Arc::new(AtomicBool::new(true));
+    let r = runnable.clone();
+    ctrlc::set_handler(move || {
+        r.store(false, Ordering::SeqCst);
+    })
+    .expect("Failed to set handler for SIGINT / SIGTERM");
+
+    match do_main(runnable) {
         Err(x) => {
             eprintln!("Error: {}", x);
             eprintln!("{}", x.backtrace());
diff --git a/examples/softirqs.rs b/examples/softirqs.rs
index 579f1619..3f0093f6 100644
--- a/examples/softirqs.rs
+++ b/examples/softirqs.rs
@@ -2,6 +2,8 @@ use bcc::core::BPF;
 use clap::{App, Arg};
 use failure::Error;
 
+use core::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::{fmt, mem, ptr, thread, time};
 
 // A simple tool for reporting on time spent in softirq handlers
@@ -71,7 +73,7 @@ impl fmt::Display for SoftIRQ {
     }
 }
 
-fn do_main() -> Result<(), Error> {
+fn do_main(runnable: Arc<AtomicBool>) -> Result<(), Error> {
     let matches = App::new("softirqs")
         .about("Reports time spent in IRQ Handlers")
         .arg(
@@ -112,7 +114,7 @@ fn do_main() -> Result<(), Error> {
     let table = module.table("dist");
     let mut window = 0;
 
-    loop {
+    while runnable.load(Ordering::SeqCst) {
         thread::sleep(time::Duration::new(interval as u64, 0));
         println!("======");
         for entry in table.iter() {
@@ -138,6 +140,7 @@ fn do_main() -> Result<(), Error> {
             }
         }
     }
+    Ok(())
 }
 
 fn parse_struct(x: &[u8]) -> irq_key_t {
@@ -145,7 +148,14 @@ fn parse_struct(x: &[u8]) -> irq_key_t {
 }
 
 fn main() {
-    match do_main() {
+    let runnable = Arc::new(AtomicBool::new(true));
+    let r = runnable.clone();
+    ctrlc::set_handler(move || {
+        r.store(false, Ordering::SeqCst);
+    })
+    .expect("Failed to set handler for SIGINT / SIGTERM");
+
+    match do_main(runnable) {
         Err(x) => {
             eprintln!("Error: {}", x);
             eprintln!("{}", x.backtrace());
diff --git a/examples/strlen.rs b/examples/strlen.rs
index ddda0bdd..b44be1dd 100644
--- a/examples/strlen.rs
+++ b/examples/strlen.rs
@@ -6,9 +6,12 @@ extern crate libc;
 use bcc::core::BPF;
 use byteorder::{NativeEndian, ReadBytesExt};
 use failure::Error;
+
+use core::sync::atomic::{AtomicBool, Ordering};
 use std::io::Cursor;
+use std::sync::Arc;
 
-fn do_main() -> Result<(), Error> {
+fn do_main(runnable: Arc<AtomicBool>) -> Result<(), Error> {
     let code = "
 #include <uapi/linux/ptrace.h>
 
@@ -39,7 +42,7 @@ int count(struct pt_regs *ctx) {
         -1, /* all PIDs */
     )?;
     let table = module.table("counts");
-    loop {
+    while runnable.load(Ordering::SeqCst) {
         std::thread::sleep(std::time::Duration::from_millis(1000));
         for e in &table {
             // key and value are each a Vec<u8> so we need to transform them into a string and
@@ -51,6 +54,7 @@ int count(struct pt_regs *ctx) {
             }
         }
     }
+    Ok(())
 }
 
 fn get_string(x: &[u8]) -> String {
@@ -61,7 +65,14 @@ fn get_string(x: &[u8]) -> String {
 }
 
 fn main() {
-    match do_main() {
+    let runnable = Arc::new(AtomicBool::new(true));
+    let r = runnable.clone();
+    ctrlc::set_handler(move || {
+        r.store(false, Ordering::SeqCst);
+    })
+    .expect("Failed to set handler for SIGINT / SIGTERM");
+
+    match do_main(runnable) {
         Err(x) => {
             eprintln!("Error: {}", x);
             eprintln!("{}", x.backtrace());