1 /*
2 * FollyRequestContextSwitch Monitor RequestContext switch events for any binary
3 * uses the class from [folly](http://bit.ly/2h6S1yx).
4 * For Linux, uses BCC, eBPF. Embedded C.
5 *
6 * Basic example of using USDT with BCC.
7 *
8 * USAGE: FollyRequestContextSwitch PATH_TO_BINARY
9 *
10 * Copyright (c) Facebook, Inc.
11 * Licensed under the Apache License, Version 2.0 (the "License")
12 */
13
14 #include <signal.h>
15 #include <functional>
16 #include <iostream>
17 #include <vector>
18
19 #include "BPF.h"
20
21 const std::string BPF_PROGRAM = R"(
22 #include <linux/sched.h>
23 #include <uapi/linux/ptrace.h>
24
25 struct event_t {
26 int pid;
27 char name[16];
28 uint64_t old_addr;
29 uint64_t new_addr;
30 };
31
32 BPF_PERF_OUTPUT(events);
33
34 int on_context_switch(struct pt_regs *ctx) {
35 struct event_t event = {};
36
37 event.pid = bpf_get_current_pid_tgid();
38 bpf_get_current_comm(&event.name, sizeof(event.name));
39
40 bpf_usdt_readarg(1, ctx, &event.old_addr);
41 bpf_usdt_readarg(2, ctx, &event.new_addr);
42
43 events.perf_submit(ctx, &event, sizeof(event));
44 return 0;
45 }
46 )";
47
48 // Define the same struct to use in user space.
49 struct event_t {
50 int pid;
51 char name[16];
52 uint64_t old_addr;
53 uint64_t new_addr;
54 };
55
handle_output(void * cb_cookie,void * data,int data_size)56 void handle_output(void* cb_cookie, void* data, int data_size) {
57 auto event = static_cast<event_t*>(data);
58 std::cout << "PID " << event->pid << " (" << event->name << ") ";
59 std::cout << "folly::RequestContext switch from " << event->old_addr << " to "
60 << event->new_addr << std::endl;
61 }
62
63 std::function<void(int)> shutdown_handler;
64
signal_handler(int s)65 void signal_handler(int s) { shutdown_handler(s); }
66
main(int argc,char ** argv)67 int main(int argc, char** argv) {
68 std::string binary;
69 pid_t pid = -1;
70 for (int i = 0; i < argc; i++) {
71 if (strncmp(argv[i], "--pid", 5) == 0) {
72 pid = std::stoi(argv[i + 1]);
73 i++;
74 continue;
75 }
76 if (strncmp(argv[i], "--binary", 8) == 0) {
77 binary = argv[i + 1];
78 i++;
79 continue;
80 }
81 }
82
83 if (pid <= 0 && binary.empty()) {
84 std::cout << "Must specify at least one of binary or PID:" << std::endl
85 << "FollyRequestContextSwitch [--pid PID] [--binary BINARY]"
86 << std::endl;
87 exit(1);
88 }
89
90 ebpf::USDT u(binary, pid, "folly", "request_context_switch_before",
91 "on_context_switch");
92
93 ebpf::BPF* bpf = new ebpf::BPF();
94
95 auto init_res = bpf->init(BPF_PROGRAM, {}, {u});
96 if (init_res.code() != 0) {
97 std::cerr << init_res.msg() << std::endl;
98 return 1;
99 }
100
101 auto attach_res = bpf->attach_usdt(u);
102 if (attach_res.code() != 0) {
103 std::cerr << attach_res.msg() << std::endl;
104 return 1;
105 } else {
106 std::cout << "Attached to USDT " << u;
107 }
108
109 auto open_res = bpf->open_perf_buffer("events", &handle_output);
110 if (open_res.code() != 0) {
111 std::cerr << open_res.msg() << std::endl;
112 return 1;
113 }
114
115 shutdown_handler = [&](int s) {
116 std::cerr << "Terminating..." << std::endl;
117 bpf->detach_usdt(u);
118 delete bpf;
119 exit(0);
120 };
121 signal(SIGINT, signal_handler);
122
123 std::cout << "Started tracing, hit Ctrl-C to terminate." << std::endl;
124 auto perf_buffer = bpf->get_perf_buffer("events");
125 if (perf_buffer)
126 while (true)
127 // 100ms timeout
128 perf_buffer->poll(100);
129
130 return 0;
131 }
132