import tracing

* add argument to parser
* add out_path as fn parameter
* set necessary env variables for import tracing

12 files changed, 296 insertions, 2 deletions

diff --git a/library/python/import_tracing/constructor/__init__.py b/library/python/import_tracing/constructor/__init__.py
new file mode 100644
index 0000000000..47ae8527d4
--- /dev/null
+++ b/library/python/import_tracing/constructor/__init__.py
@@ -0,0 +1,10 @@
+import os
+
+
+def init():
+    if "Y_PYTHON_TRACE_FILE" in os.environ:
+        import atexit
+        import library.python.import_tracing.lib.regulator as regulator
+
+        regulator.enable(os.getenv("Y_PYTHON_TRACE_FILE"))
+        atexit.register(regulator.disable)
diff --git a/library/python/import_tracing/constructor/ya.make b/library/python/import_tracing/constructor/ya.make
new file mode 100644
index 0000000000..8ba6403897
--- /dev/null
+++ b/library/python/import_tracing/constructor/ya.make
@@ -0,0 +1,15 @@
+PY23_LIBRARY()
+
+STYLE_PYTHON()
+
+PY_CONSTRUCTOR(library.python.import_tracing.constructor)
+
+PY_SRCS(
+    __init__.py
+)
+
+PEERDIR(
+    library/python/import_tracing/lib
+)
+
+END()
diff --git a/library/python/import_tracing/lib/constants.py b/library/python/import_tracing/lib/constants.py
new file mode 100644
index 0000000000..8c445c2067
--- /dev/null
+++ b/library/python/import_tracing/lib/constants.py
@@ -0,0 +1 @@
+MCS_IN_SEC = 1e6
diff --git a/library/python/import_tracing/lib/converters/base.py b/library/python/import_tracing/lib/converters/base.py
new file mode 100644
index 0000000000..8195d6051c
--- /dev/null
+++ b/library/python/import_tracing/lib/converters/base.py
@@ -0,0 +1,3 @@
+class BaseTraceConverter:
+    def dump(self, events, filepath):
+        raise NotImplementedError()
diff --git a/library/python/import_tracing/lib/converters/chrometrace.py b/library/python/import_tracing/lib/converters/chrometrace.py
new file mode 100644
index 0000000000..931f72db58
--- /dev/null
+++ b/library/python/import_tracing/lib/converters/chrometrace.py
@@ -0,0 +1,33 @@
+import json
+import os
+
+import library.python.import_tracing.lib.converters.base as base_converter
+
+
+class ChromiumTraceConverter(base_converter.BaseTraceConverter):
+    @staticmethod
+    def _yield_in_chrome_trace_format(events, pid):
+        for event in events:
+            yield {
+                "cat": event.modname,
+                "name": event.filename,
+                "ph": "B",
+                "ts": event.start_time,
+                "pid": pid,
+                "tid": event.tid,
+                "args": {},
+            }
+
+            yield {
+                "cat": event.modname,
+                "name": event.filename,
+                "ph": "E",
+                "ts": event.end_time,
+                "pid": pid,
+                "tid": event.tid,
+            }
+
+    def dump(self, events, filepath):
+        pid = os.getpid()
+        with open(filepath, "w") as file:
+            file.write(json.dumps(tuple(self._yield_in_chrome_trace_format(events, pid))))
diff --git a/library/python/import_tracing/lib/converters/raw.py b/library/python/import_tracing/lib/converters/raw.py
new file mode 100644
index 0000000000..7c307a5984
--- /dev/null
+++ b/library/python/import_tracing/lib/converters/raw.py
@@ -0,0 +1,56 @@
+import library.python.import_tracing.lib.converters.base as base_converter
+import library.python.import_tracing.lib.constants as constants
+
+
+class RawTextTraceConverter(base_converter.BaseTraceConverter):
+    @staticmethod
+    def _get_columns_length(events):
+        max_filename = 0
+        max_cumtime = 0
+        max_end_time = 0
+
+        for event in events:
+            max_filename = max(max_filename, len(event.filename))
+            max_cumtime = max(max_cumtime, event.end_time - event.start_time)
+            max_end_time = max(max_end_time, event.end_time)
+
+        return len(str(max_cumtime)), max_filename, max_end_time
+
+    @staticmethod
+    def _get_sorted_events(events):
+        return sorted(events, key=lambda event: event.end_time - event.start_time, reverse=True)
+
+    @staticmethod
+    def _format_line(cumtime, filename, max_cumtime, max_filename):
+        return "{0:<{max_cumtime}}\t{1:<{max_filename}}\n".format(
+            cumtime,
+            filename,
+            max_cumtime=max_cumtime,
+            max_filename=max_filename,
+        )
+
+    def dump(self, events, filepath):
+        max_cumtime, max_filename, max_end_time = self._get_columns_length(events)
+        max_line_length = max_cumtime + max_filename
+
+        with open(filepath, "w") as file:
+            # total time taken
+            file.write("total time taken (seconds): {0:.4f}\n".format(max_end_time / constants.MCS_IN_SEC))
+            file.write("-" * max_line_length + "\n")
+
+            # header
+            file.write(self._format_line("cumtime", "filename", max_cumtime, max_filename))
+            file.write("-" * max_line_length + "\n")
+
+            # trace info
+            for event in self._get_sorted_events(events):
+                time_taken = format(((event.end_time - event.start_time) / constants.MCS_IN_SEC), ".6f")
+
+                file.write(
+                    self._format_line(
+                        time_taken,
+                        event.filename,
+                        max_cumtime,
+                        max_filename,
+                    )
+                )
diff --git a/library/python/import_tracing/lib/event.py b/library/python/import_tracing/lib/event.py
new file mode 100644
index 0000000000..94547e1398
--- /dev/null
+++ b/library/python/import_tracing/lib/event.py
@@ -0,0 +1,9 @@
+class Event:
+    __slots__ = ("modname", "filename", "tid", "start_time", "end_time")
+
+    def __init__(self, modname, filename, tid=None, start_time=None, end_time=None):
+        self.modname = modname
+        self.filename = filename
+        self.tid = tid
+        self.start_time = start_time
+        self.end_time = end_time
diff --git a/library/python/import_tracing/lib/import_tracer.py b/library/python/import_tracing/lib/import_tracer.py
new file mode 100644
index 0000000000..39dfc11bc9
--- /dev/null
+++ b/library/python/import_tracing/lib/import_tracer.py
@@ -0,0 +1,52 @@
+import threading
+import time
+import collections
+import library.python.import_tracing.lib.event as events
+import library.python.import_tracing.lib.constants as constants
+
+
+class ImportTracer:
+    def __init__(self):
+        self.events = collections.OrderedDict()
+        self.start_time = time.time()
+
+    def start_event(self, modname, filename, tid=None):
+        tid = tid if tid is not None else threading.current_thread().ident
+        time_from_start = self._get_current_time_from_start()
+
+        event_key = (modname, tid)
+        new_event = events.Event(
+            modname=modname,
+            filename=filename,
+            tid=tid,
+            start_time=time_from_start,
+            end_time=None,
+        )
+
+        self.events[event_key] = new_event
+
+    def finish_event(self, modname, filename, tid=None):
+        tid = tid if tid is not None else threading.current_thread().ident
+        event_key = (modname, tid)
+        event = self.events[event_key]
+
+        end_time = self._get_current_time_from_start()
+        event.end_time = end_time
+
+    def get_events(self, close_not_finished=False):
+        end_time = self._get_current_time_from_start()
+
+        for event in self.events.values():
+            if close_not_finished and event.end_time is None:
+                yield events.Event(
+                    modname=event.modname,
+                    filename=event.filename,
+                    tid=event.tid,
+                    start_time=event.start_time,
+                    end_time=end_time,
+                )
+            else:
+                yield event
+
+    def _get_current_time_from_start(self):
+        return (time.time() - self.start_time) * constants.MCS_IN_SEC
diff --git a/library/python/import_tracing/lib/regulator.py b/library/python/import_tracing/lib/regulator.py
new file mode 100644
index 0000000000..721b022ef7
--- /dev/null
+++ b/library/python/import_tracing/lib/regulator.py
@@ -0,0 +1,73 @@
+import collections
+import os
+
+_Instance = collections.namedtuple("_Instance", ("import_tracer", "converter", "filepath"))
+
+INSTANCE = None
+
+
+def _get_converter_instance():
+    import library.python.import_tracing.lib.converters.raw as text_converter
+    import library.python.import_tracing.lib.converters.chrometrace as chrome_converter
+
+    converter_mapping = {"text": text_converter.RawTextTraceConverter, "evlog": chrome_converter.ChromiumTraceConverter}
+
+    env_val = os.getenv("Y_PYTHON_TRACE_FORMAT")
+
+    converter = converter_mapping.get(env_val, text_converter.RawTextTraceConverter)
+
+    return converter()
+
+
+def _resolve_filepath(filemask):
+    import socket
+    import sys
+
+    pid = os.getpid()
+    hostname = socket.gethostname()
+    executable_filename = os.path.basename(sys.executable)
+
+    return filemask.replace("%p", str(pid)).replace("%h", hostname).replace("%e", executable_filename)
+
+
+def enable(filemask):
+    import library.python.import_tracing.lib.import_tracer as import_tracer
+    import __res
+
+    global INSTANCE
+
+    if INSTANCE is not None:
+        return INSTANCE
+
+    converter = _get_converter_instance()
+    import_tracer = import_tracer.ImportTracer()
+
+    def before_import_callback(modname, filename):
+        import_tracer.start_event(modname, filename)
+
+    def after_import_callback(modname, filename):
+        import_tracer.finish_event(modname, filename)
+
+    __res.importer.set_callbacks(before_import_callback, after_import_callback)
+
+    filepath = _resolve_filepath(filemask)
+
+    new_instance = _Instance(import_tracer, converter, filepath)
+    INSTANCE = new_instance
+
+    return new_instance
+
+
+def disable(close_not_finished=False):
+    global INSTANCE
+
+    if INSTANCE is None:
+        return
+
+    import_tracer = INSTANCE.import_tracer
+    converter = INSTANCE.converter
+    filepath = INSTANCE.filepath
+
+    converter.dump(import_tracer.get_events(close_not_finished), filepath)
+
+    INSTANCE = None
diff --git a/library/python/import_tracing/lib/ya.make b/library/python/import_tracing/lib/ya.make
new file mode 100644
index 0000000000..118d69fbfa
--- /dev/null
+++ b/library/python/import_tracing/lib/ya.make
@@ -0,0 +1,15 @@
+PY23_LIBRARY()
+
+STYLE_PYTHON()
+
+PY_SRCS(
+    event.py
+    import_tracer.py
+    constants.py
+    regulator.py
+    converters/base.py
+    converters/chrometrace.py
+    converters/raw.py
+)
+
+END()
diff --git a/library/python/testing/import_test/import_test.py b/library/python/testing/import_test/import_test.py
index 02331c18a7..7795c3c321 100644
--- a/library/python/testing/import_test/import_test.py
+++ b/library/python/testing/import_test/import_test.py
@@ -53,7 +53,7 @@ def check_imports(no_check=(), extra=(), skip_func=None, py_main=None):
     import_times = {}
 
     def norm(s):
-        return (s[:-9] if s.endswith('.__init__') else s)
+        return s[:-9] if s.endswith('.__init__') else s
 
     modules = sys.extra_modules | set(extra)
     modules = sorted(modules, key=norm)
@@ -87,7 +87,7 @@ def check_imports(no_check=(), extra=(), skip_func=None, py_main=None):
                 if module == '__main__':
                     importer.load_module('__main__', '__main__py')
                 elif module.endswith('.__init__'):
-                    __import__(module[:-len('.__init__')])
+                    __import__(module[: -len('.__init__')])
                 else:
                     __import__(module)
 
@@ -123,6 +123,29 @@ def main():
 
     skip_names = sys.argv[1:]
 
+    # SIGUSR2 is used by test_tool to teardown tests
+    if hasattr(signal, "SIGUSR2"):
+        # Dump python import tracing
+        import library.python.import_tracing.lib.regulator as regulator
+
+        # get the original handler to return control to it after dumping
+        signum = signal.SIGUSR2
+        orig_handler = signal.getsignal(signum)
+
+        if not hasattr(signal, 'raise_signal'):
+            # Only available for Python 3.8+
+            def raise_signal(signum):
+                os.kill(os.getpid(), signum)
+        else:
+            raise_signal = signal.raise_signal
+
+        def stop_tracing_handler(s, f):
+            regulator.disable(close_not_finished=True)
+            signal.signal(signal.SIGUSR2, orig_handler)
+            raise_signal(signum)
+
+        signal.signal(signal.SIGUSR2, stop_tracing_handler)
+
     try:
         import faulthandler
     except ImportError:
diff --git a/library/python/testing/import_test/ya.make b/library/python/testing/import_test/ya.make
index 3f95288f4d..504e3f3a57 100644
--- a/library/python/testing/import_test/ya.make
+++ b/library/python/testing/import_test/ya.make
@@ -2,4 +2,8 @@ PY23_LIBRARY()
 
 PY_SRCS(import_test.py)
 
+PEERDIR(
+    library/python/import_tracing/lib
+)
+
 END()