Merge pull request #9 from victormlg/cf-profile

Moved cf-profile.py to cfengine-cli

Author: olehermanse

src/cfengine_cli/commands.py

@@ -1,5 +1,8 @@
 import sys
 import os
+import re
+import json
+from cfengine_cli.profile import profile_cfengine, generate_callstack
 from cfengine_cli.dev import dispatch_dev_subcommand
 from cfengine_cli.lint import lint_cfbs_json, lint_json, lint_policy_file
 from cfengine_cli.shell import user_command
@@ -128,3 +131,19 @@ def run() -> int:
 
 def dev(subcommand, args) -> int:
     return dispatch_dev_subcommand(subcommand, args)
+
+
+def profile(args) -> int:
+    data = None
+    with open(args.profiling_input, "r") as f:
+        m = re.search(r"\[[.\s\S]*\]", f.read())
+        if m is not None:
+            data = json.loads(m.group(0))
+
+    if data is not None and any([args.bundles, args.functions, args.promises]):
+        profile_cfengine(data, args)
+
+    if args.flamegraph:
+        generate_callstack(data, args.flamegraph)
+
+    return 0

src/cfengine_cli/main.py

@@ -59,6 +59,20 @@ def _get_arg_parser():
         "run", help="Run the CFEngine agent, fetching, evaluating, and enforcing policy"
     )
 
+    profile_parser = subp.add_parser(
+        "profile", help="Parse CFEngine profiling output (cf-agent -Kp)"
+    )
+    profile_parser.add_argument(
+        "profiling_input", help="Path to the profiling input file"
+    )
+    profile_parser.add_argument("--top", type=int, default=10)
+    profile_parser.add_argument("--bundles", action="store_true")
+    profile_parser.add_argument("--promises", action="store_true")
+    profile_parser.add_argument("--functions", action="store_true")
+    profile_parser.add_argument(
+        "--flamegraph", type=str, help="Generate input file for ./flamegraph.pl"
+    )
+
     dev_parser = subp.add_parser(
         "dev", help="Utilities intended for developers / maintainers of CFEngine"
     )
@@ -101,6 +115,8 @@ def run_command_with_args(args) -> int:
         return commands.run()
     if args.command == "dev":
         return commands.dev(args.dev_command, args)
+    if args.command == "profile":
+        return commands.profile(args)
     raise UserError(f"Unknown command: '{args.command}'")
 
 

src/cfengine_cli/profile.py

@@ -0,0 +1,112 @@
+import os
+from collections import defaultdict
+
+
+def format_elapsed_time(elapsed_ns):
+    elapsed_ms = float(elapsed_ns) / 1e6
+
+    if elapsed_ms < 1000:
+        return "%.2f ms" % elapsed_ms
+    elif elapsed_ms < 60000:
+        elapsed_s = elapsed_ms / 1000.0
+        return "%.2fs" % elapsed_s
+    else:
+        elapsed_s = elapsed_ms / 1000.0
+        minutes = int(elapsed_s // 60)
+        seconds = int(elapsed_s % 60)
+        return "%dm%ds" % (minutes, seconds)
+
+
+def format_label(component, event_type, ns, name):
+    if component == "function":
+        return "%s %s" % (component, name)
+    elif event_type == "methods":
+        return "bundle invocation"
+    elif component == "promise":
+        return "%s %s" % (component, event_type)
+    return "%s %s %s:%s" % (component, event_type, ns, name)
+
+
+def format_columns(events, top):
+
+    labels = []
+
+    for event in events[:top]:
+        label = format_label(
+            event["component"], event["type"], event["namespace"], event["name"]
+        )
+        location = "%s:%s" % (event["source"], event["offset"]["line"])
+        time = format_elapsed_time(event["elapsed"])
+
+        labels.append((label, location, time))
+
+    return labels
+
+
+def get_max_column_lengths(lines, indent=4):
+
+    max_type, max_location, max_time = 0, 0, 0
+
+    for label, location, time_ms in lines:
+        max_type = max(max_type, len(label))
+        max_location = max(max_location, len(location))
+        max_time = max(max_time, len(time_ms))
+
+    return max_type + indent, max_location + indent, max_time + indent
+
+
+def profile_cfengine(events, args):
+
+    filter = defaultdict(list)
+
+    if args.bundles:
+        filter["component"].append("bundle")
+        filter["type"].append("methods")
+
+    if args.promises:
+        filter["type"] += list(
+            set(
+                event["type"]
+                for event in events
+                if event["component"] == "promise" and event["type"] != "methods"
+            )
+        )
+
+    if args.functions:
+        filter["component"].append("function")
+
+    # filter events
+    if filter is not None:
+        events = [
+            event
+            for field in filter.keys()
+            for event in events
+            if event[field] in filter[field]
+        ]
+
+    # sort events
+    events = sorted(events, key=lambda x: x["elapsed"], reverse=True)
+
+    lines = format_columns(events, args.top)
+    line_format = "%-{}s %-{}s %{}s".format(*get_max_column_lengths(lines))
+
+    # print top k filtered events
+    print(line_format % ("Type", "Location", "Time"))
+    for label, location, time_ms in lines:
+        print(line_format % (label, location, time_ms))
+
+
+def generate_callstack(data, stack_path):
+
+    with open(stack_path, "w") as f:
+        for event in data:
+            f.write("%s %d\n" % (event["callstack"], event["elapsed"]))
+
+    print(
+        "Successfully generated callstack at '{}'".format(os.path.abspath(stack_path))
+    )
+    print(
+        "Run './flamgraph {} > flamegraph.svg' to generate the flamegraph".format(
+            stack_path
+        )
+    )