homeserver/plugins/callback/trace.py

169 lines
5.3 KiB
Python
Raw Normal View History

# Copyright 2021 Google LLC
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# version 3 as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
DOCUMENTATION = '''
name: trace
type: aggregate
short_description: write playbook output to Chrome's Trace Event Format file
description:
- This callback writes playbook output to Chrome Trace Event Format file.
author: Mark Hansen (@mhansen)
options:
trace_output_dir:
name: output dir
default: ./trace
description: Directory to write files to.
env:
- name: TRACE_OUTPUT_DIR
hide_task_arguments:
name: Hide the arguments for a task
default: False
description: Hide the arguments for a task
env:
- name: HIDE_TASK_ARGUMENTS
requirements:
- enable in configuration
'''
import atexit
import json
import os
import time
from dataclasses import dataclass
from datetime import datetime
from typing import Dict, TextIO
from ansible.plugins.callback import CallbackBase
class CallbackModule(CallbackBase):
"""
This callback traces execution time of tasks to Trace Event Format.
This plugin makes use of the following environment variables:
TRACE_OUTPUT_DIR (optional): Directory to write JSON files to.
Default: ./trace
TRACE_HIDE_TASK_ARGUMENTS (optional): Hide the arguments for a task
Default: False
"""
CALLBACK_VERSION = 2.0
CALLBACK_TYPE = 'aggregate'
CALLBACK_NAME = 'trace'
CALLBACK_NEEDS_ENABLED = True
def __init__(self):
super(CallbackModule, self).__init__()
self._output_dir: str = os.getenv('TRACE_OUTPUT_DIR', os.path.join(os.path.expanduser('.'), 'trace'))
self._hide_task_arguments: str = os.getenv('TRACE_HIDE_TASK_ARGUMENTS', 'False').lower()
self._hosts: Dict[Host] = {}
self._next_pid: int = 1
self._first: bool = True
self._start_date: str = datetime.now().isoformat()
self._output_file: str = 'trace-%s.json' % self._start_date
if not os.path.exists(self._output_dir):
os.makedirs(self._output_dir)
output_file = os.path.join(self._output_dir, self._output_file)
self._f: TextIO = open(output_file, 'w')
self._f.write("[\n")
atexit.register(self._end)
def _write_event(self, e: Dict):
if not self._first:
self._f.write(",\n")
self._first = False
json.dump(e, self._f, sort_keys=True, indent=2) # sort for reproducibility
self._f.flush()
def v2_runner_on_start(self, host, task):
uuid = task._uuid
name = task.get_name().strip()
args = None
if not task.no_log and self._hide_task_arguments == 'false':
args = task.args
host_uuid = host._uuid
if host_uuid not in self._hosts:
pid = self._next_pid
self._hosts[host_uuid] = Host(pid=pid, name=host.name)
self._next_pid += 1
# https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#bookmark=id.iycbnb4z7i9g
self._write_event({
"name": "process_name",
"pid": pid,
"cat": "process",
"ph": "M",
"args": {
"name": host.name,
},
})
# See "Duration Events" in:
# https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.nso4gcezn7n1
self._write_event({
"name": name,
"cat": "runner",
"ph": "B", # Begin
"ts": time.time_ns() / 1000,
"pid": self._hosts[host_uuid].pid,
"id": abs(hash(uuid)),
"args": {
"args": args,
"task": name,
"path": task.get_path(),
"host": host.name,
},
})
def _end_span(self, result, status: str):
task = result._task
uuid = task._uuid
# See "Duration Events" in:
# https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview#heading=h.nso4gcezn7n1
self._write_event({
"name": task.get_name().strip(),
"cat": "runner",
"id": abs(hash(uuid)),
"ph": "E", # End
"ts": time.time_ns() / 1000,
"pid": self._hosts[result._host._uuid].pid,
"args": {
"status": status,
},
})
def v2_runner_on_ok(self, result):
self._end_span(result, status="ok")
def v2_runner_on_unreachable(self, result):
self._end_span(result, 'unreachable')
def v2_runner_on_failed(self, result, ignore_errors=False):
self._end_span(result, status='failed')
def v2_runner_on_skipped(self, result):
self._end_span(result, status='skipped')
def _end(self):
self._f.write("\n]")
self._f.close()
@dataclass
class Host:
name: str
pid: int