Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
179 lines (140 sloc) 5.9 KB
#!/usr/bin/python
import argparse
import re
import sys
import time
from bcc import BPF, USDT
# minimum amount of total allocations for an object to be included in the flamegraph
MIN_ALLOCATION_SHARE_PERCENTAGE=1
def next_power_of_two(sample_interval):
if sample_interval == 1:
return sample_interval
bit_test = sample_interval
set_bit_count = 0
while bit_test != 0:
if (bit_test & 1) == 1:
set_bit_count += 1
bit_test /= 2
if set_bit_count == 1:
return sample_interval
next_power_of_two = 1
while next_power_of_two < sample_interval:
next_power_of_two *= 2
return next_power_of_two
def apply_regex(stack_counts, regex_list, should_include):
if regex_list is None:
return stack_counts
filtered_stack_counts=dict()
if should_include is False:
filtered_stack_counts=dict(stack_counts)
for pattern in regex_list:
regex = re.compile(pattern)
for k, v in stack_counts.iteritems():
if should_include and regex.search(k) is not None:
filtered_stack_counts[k] = v
elif should_include is False and regex.search(k) is not None:
if k in filtered_stack_counts:
del filtered_stack_counts[k]
return filtered_stack_counts
def apply_exclusion_regex(stack_counts, regex_list):
return apply_regex(stack_counts, regex_list, False)
def apply_inclusion_regex(stack_counts, regex_list):
return apply_regex(stack_counts, regex_list, True)
def remove_objects_with_small_allocation_count(stack_counts):
alloc_count_by_class = dict()
total_allocations = 0
large_allocations = dict()
for k, v in stack_counts.iteritems():
class_name = k.split(";")[0]
if class_name not in alloc_count_by_class:
alloc_count_by_class[class_name] = v
else:
alloc_count_by_class[class_name] += v
total_allocations += v
for k, v in stack_counts.iteritems():
class_name = k.split(";")[0]
percentage_of_allocations = (alloc_count_by_class[class_name] * 100) / float(total_allocations)
if percentage_of_allocations > MIN_ALLOCATION_SHARE_PERCENTAGE:
large_allocations[k] = v
return large_allocations
def remove_non_ascii(text):
return ''.join(i for i in text if ord(i)<128)
prog="""
#include <linux/types.h>
#include <uapi/linux/ptrace.h>
#include <uapi/linux/bpf_perf_event.h>
#include <linux/sched.h>
struct key_t {
u32 pid;
int user_stack_id;
char name[256];
};
BPF_HASH(counts, struct key_t);
BPF_HASH(start, u32);
BPF_HASH(sample_count, u64);
BPF_HASH(tids, long);
BPF_STACK_TRACE(stack_traces, 10240);
int trace_alloc(struct pt_regs *ctx, long tid, char *name, int nameLength, int wordSize) {
u32 pid = bpf_get_current_pid_tgid() >> 32;
u64 zero = 0, *val, *tc, *current_count;
u64 one = 1;
current_count = sample_count.lookup_or_init(&zero, &zero);
(*current_count)++;
if ((((*current_count)) & %d) != zero) {
return 0;
}
struct key_t key = {.pid = pid};
bpf_probe_read(&key.name, sizeof(key.name), (void *)PT_REGS_PARM3(ctx));
tc = tids.lookup_or_init(&tid, &zero);
(*tc)++;
key.user_stack_id = stack_traces.get_stackid(ctx, 0 | BPF_F_REUSE_STACKID | BPF_F_USER_STACK);
val = counts.lookup_or_init(&key, &zero);
(*val)++;
return 0;
};
"""
def get_arg_parser():
parser = argparse.ArgumentParser(description = 'Generate heap allocation flamegraphs')
parser.add_argument('-j', type=str, dest='lib_jvm_path', help='Path to libjvm.so', default='/usr/lib/jvm/java-8-openjdk-amd64/jre/lib/amd64/server/libjvm.so')
parser.add_argument('-p', type=int, dest='pid', help='PID of the target process', required=True)
parser.add_argument('-i', type=str, dest='include_regex', help='Regex for stacks to include', nargs='*')
parser.add_argument('-e', type=str, dest='exclude_regex', help='Regex for stacks to exclude', nargs='*')
parser.add_argument('-s', type=int, dest='sampling_frequency', help='Sample every N allocations', default=1)
parser.add_argument('-d', type=int, dest='duration_seconds', help='Recording duration in seconds', default=10)
return parser
args = get_arg_parser().parse_args()
print(prog % (next_power_of_two(args.sampling_frequency) - 1))
pid = args.pid
usdt = USDT(path=args.lib_jvm_path, pid=args.pid)
usdt.enable_probe(probe="object__alloc", fn_name="trace_alloc")
bpf = BPF(text=prog % (next_power_of_two(args.sampling_frequency) - 1), usdt_contexts=[usdt])
time.sleep(float(str(args.duration_seconds)))
if len(bpf["tids"]) == 0:
print("No data found - are DTrace probes enabled in running process?")
stack_traces = bpf["stack_traces"]
all_stacks=[]
stack_counts=dict()
for k, v in bpf["counts"].iteritems():
stack=[]
for addr in stack_traces.walk(k.user_stack_id):
pid = k.pid
symbol=bpf.sym(addr, pid)
if symbol == "[unknown]":
stack.append(("0x%-16x" % addr).strip())
else:
stack_trace_entry = symbol.strip().replace(';',':')
if stack_trace_entry[0] == 'L' and stack_trace_entry.find(':::'):
stack_trace_entry = stack_trace_entry[1:]
stack.append(stack_trace_entry.encode('utf-8', errors='replace'))
stack.reverse()
try:
key_name = remove_non_ascii(k.name).strip().encode('utf-8', errors='replace')
stack_counts[key_name + ";" + u";".join(stack).encode('utf-8', errors='replace')] = int(v.value)
except UnicodeDecodeError as e:
err_msg = str(e)
print("Failed to decode stack: " + k.name.strip() + ";" + str(stack) + ": " + err_msg)
stack_counts = remove_objects_with_small_allocation_count(stack_counts)
stack_counts = apply_inclusion_regex(stack_counts, args.include_regex)
stack_counts = apply_exclusion_regex(stack_counts, args.exclude_regex)
for k, v in stack_counts.iteritems():
print(k + " " + str(v))