waf_trace.py

import collections
import queue
import time

from waflib import Runner, Task
TaskEvent = collections.namedtuple("TaskEvent", "id start end task outputs".split())

def monkey_patch():
    prev_start = Runner.Parallel.start
    def start(self):
        self.task_trace = queue.Queue()
        prev_start(self)

        if self.dirty and not self.task_trace.empty():
            tasks = []
            while True:
                try:
                    tasks.append(self.task_trace.get(block=False))
                except queue.Empty:
                    break

            workers = []
            def next_free_worker_index_for(task):
                # This does not actually capture the correct worker but
                # reproduces the level of parallelism.
                for worker_index, busy_until in enumerate(workers):
                    if busy_until <= task.start:
                        workers[worker_index] = task.end
                        break
                else:
                    worker_index = len(workers)
                    workers.append(task.end)
                return worker_index

            entries = []
            for task in tasks:
                # TODO: Can we get the name of the target corresponding to this task?
                name = task.task

                worker_index = next_free_worker_index_for(task)
                # Trace Event Format, see https://docs.google.com/document/d/1CvAClvFfyA5R-PhYUmn5OOQtYMH4h6I0nSsKchNAySU/preview
                entries.append({
                    # Name is used as a label for the slice and to derive its color.
                    "name": name,
                    # Category can be used for filtering.
                    "cat": task.task,
                    # Event type "X" corresponds to "complete" events.
                    "ph": "X",
                    # Timestamp in us
                    "ts": int(task.start * 1e6),
                    "dur": int((task.end - task.start) * 1e6),
                    "pid": 0,
                    "tid": worker_index,
                    "args": {"outputs": task.outputs},
                })

            trace_node = self.bld.path.make_node("waf_trace.json")
            trace_node.write_json(entries, pretty=False)
    Runner.Parallel.start = start

    prev_process = Task.Task.process
    def process(self):
        start = time.time()
        prev_process(self)
        end = time.time()

        self.generator.bld.producer.task_trace.put(TaskEvent(
            id=id(self),
            start=start,
            end=end,
            task=self.__class__.__name__,
            outputs=[str(output) for output in self.outputs],
        ))
    Task.Task.process = process

monkey_patch()