Python: Add telemetry for parser usage

Adds statistics on how many files were extracted using the old parser
and using the tree-sitter parser.

Because parsing is done in parallel across many workers, I opted not to
consolidate these statistics for the entire run. Instead, we emit the
statistics for each worker and then need to aggregate themselves after
the telemetry has been ingested. (In practice the number of workers is
~16 at most, so is unlikely to be an issue.)

In terms of implementation, I opted to simply extend the existing
`DiagnosticsWriter` object (instantiatied once per worker) with methods
for counting the number of parsed files, and then thread this object
through to `modules.py` where the magic happens.

Finally, this also required instantiating such an object in cases where
we call directly into the extractor for debugging purposes (e.g. dumping
the AST or CFG). Note that in these cases we do not actually print any
diagnostics, so it's harmless to create these objects.

Author: tausbn

python/extractor/cli-integration-test/writing-diagnostics/test_diagnostics_output.py

@@ -1,7 +1,31 @@
 import os
 import sys
+import glob
+import json
 sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..", "integration-tests"))
 import diagnostics_test_utils
 
 test_db = "db"
-diagnostics_test_utils.check_diagnostics(".", test_db, skip_attributes=True)
+diagnostics = []
+diagnostic_dir = os.path.join(test_db, "diagnostic", "extractors", "python")
+for path in glob.glob(os.path.join(diagnostic_dir, "*.jsonl")):
+    with open(path) as diagnostic_file:
+        diagnostics.extend(json.loads(line) for line in diagnostic_file)
+parser_statistics = [
+    diagnostic
+    for diagnostic in diagnostics
+    if diagnostic["source"]["id"] == "py/extractor/parser-statistics"
+]
+assert sum(
+    diagnostic["attributes"]["old_parser_file_count"]
+    + diagnostic["attributes"]["tree_sitter_parser_file_count"]
+    for diagnostic in parser_statistics
+) == 2
+diagnostics = [
+    diagnostic
+    for diagnostic in diagnostics
+    if diagnostic["source"]["id"] != "py/extractor/parser-statistics"
+]
+diagnostics_test_utils.check_diagnostics(
+    ".", test_db, actual=json.dumps(diagnostics), skip_attributes=True
+)

python/extractor/semmle/extractors/module_printer.py

@@ -6,9 +6,9 @@ class ModulePrinter(object):
 
     name = "module printer"
 
-    def __init__(self, options, trap_folder, src_archive, renamer, logger):
+    def __init__(self, options, trap_folder, src_archive, renamer, logger, diagnostics_writer):
         self.logger = logger
-        self.py_extractor = PythonExtractor(options, trap_folder, src_archive, logger)
+        self.py_extractor = PythonExtractor(options, trap_folder, src_archive, logger, diagnostics_writer)
 
     def process(self, unit):
         imports = ()

python/extractor/semmle/extractors/py_extractor.py

@@ -16,6 +16,7 @@ def __init__(self, options, trap_folder, src_archive, logger: Logger, diagnostic
         self.module_extractor = extractor.Extractor.from_options(options, trap_folder, src_archive, logger, diagnostics_writer)
         self.finder = finder.Finder.from_options_and_env(options, logger)
         self.importer = imports.importer_from_options(options, self.finder, logger)
+        self.diagnostics_writer = diagnostics_writer
 
     def _get_module_and_imports(self, unit):
         if not isinstance(unit, util.FileExtractable):
@@ -24,7 +25,7 @@ def _get_module_and_imports(self, unit):
         module = self.finder.from_extractable(unit)
         if module is None:
             return None, ()
-        py_module = module.load(self.logger)
+        py_module = module.load(self.logger, self.diagnostics_writer)
         if py_module is None:
             return None, ()
         imports = set(mod.get_extractable() for mod in self.importer.get_imports(module, py_module))

python/extractor/semmle/logging.py

@@ -367,6 +367,14 @@ def extractor_telemetry_message():
             .telemetry()
     )
 
+def parser_statistics_telemetry_message(old_parser_file_count, tree_sitter_parser_file_count):
+    return (DiagnosticMessage(Source("py/extractor/parser-statistics", "Python parser statistics"), Severity.NOTE)
+            .markdown("Internal parser telemetry for the Python extractor.\n\nNo action needed.")
+            .attribute("old_parser_file_count", old_parser_file_count)
+            .attribute("tree_sitter_parser_file_count", tree_sitter_parser_file_count)
+            .telemetry()
+    )
+
 def get_stack_trace_lines():
     """Creates a stack trace for inclusion into the `attributes` part of a diagnostic message.
     Limits the size of the stack trace to 5000 characters, so as to not make the SARIF file overly big.

python/extractor/semmle/python/finder.py

@@ -65,8 +65,8 @@ def all_sub_modules(self):
     def get_extractable(self):
         return FileExtractable(self.path)
 
-    def load(self, logger=None):
-        return PythonSourceModule(self.name, self.path, logger=logger)
+    def load(self, logger, diagnostics_writer):
+        return PythonSourceModule(self.name, self.path, logger=logger, diagnostics_writer=diagnostics_writer)
 
     def __str__(self):
         return "Python module at %s" % self.path

python/extractor/semmle/python/modules.py

@@ -18,7 +18,7 @@ class PythonSourceModule(object):
 
     kind = None
 
-    def __init__(self, name, path, logger, bytes_source = None):
+    def __init__(self, name, path, logger, diagnostics_writer, bytes_source = None):
         assert isinstance(path, str), path
         self.name = name # May be None
         self.path = path
@@ -34,6 +34,7 @@ def __init__(self, name, path, logger, bytes_source = None):
         self._line_types = None
         self._comments = None
         self._tokens = None
+        self.diagnostics_writer = diagnostics_writer
         self.logger = logger
         with timers["decode"]:
             self.encoding, self.bytes_source  = semmle.python.parser.tokenizer.encoding_from_source(bytes_source)
@@ -113,6 +114,7 @@ def old_py_ast(self):
                 self.logger.debug("Trying old parser on %s", self.path)
                 self._py_ast = semmle.python.parser.parse(self.tokens, self.logger)
                 self.logger.debug("Old parser successful on %s", self.path)
+                self.diagnostics_writer.record_old_parser()
         else:
             self.logger.debug("Found (during old_py_ast) parse tree for %s in cache", self.path)
         return self._py_ast
@@ -147,6 +149,7 @@ def py_ast(self):
                             self.logger.debug("Trying tsg-python on %s", self.path)
                             self._py_ast = semmle.python.parser.tsg_parser.parse(self.path, self.logger)
                             self.logger.debug("tsg-python successful on %s", self.path)
+                            self.diagnostics_writer.record_tree_sitter_parser()
                         else:
                             self.logger.debug("Found (during py_ast) parse tree for %s in cache", self.path)
                     return self._py_ast

python/extractor/semmle/python/parser/dump_ast.py

@@ -119,7 +119,8 @@ def reset_error_count(self):
         self.error_count = 0
 
 def old_parser(inputfile, logger):
-    mod = PythonSourceModule(None, inputfile, logger)
+    from semmle.worker import DiagnosticsWriter
+    mod = PythonSourceModule(None, inputfile, logger, DiagnosticsWriter(0))
     logger.close()
     return mod.old_py_ast
 

python/extractor/semmle/python/passes/flow.py

@@ -1916,7 +1916,8 @@ def write_ssa_phi(out, phi, arg):
         import semmle.python.parser.tsg_parser
         parsed_ast = semmle.python.parser.tsg_parser.parse(inputfile, FakeLogger())
     else:
-        module = modules.PythonSourceModule("__main__", inputfile, FakeLogger())
+        from semmle.worker import DiagnosticsWriter
+        module = modules.PythonSourceModule("__main__", inputfile, FakeLogger(), DiagnosticsWriter(0))
         parsed_ast = module.ast
     FlowPass(options.split, options.prune, options.unroll).extract(parsed_ast, writer)
     writer.close()

python/extractor/semmle/worker.py

@@ -11,6 +11,7 @@
 from semmle.profiling import get_profiler
 from semmle.path_rename import renamer_from_options_and_env
 from semmle.logging import WARN, recursion_error_message, internal_error_message, extractor_telemetry_message, Logger
+from semmle.logging import parser_statistics_telemetry_message
 from semmle.util import FileExtractable, FolderExtractable
 
 class ExtractorFailure(Exception):
@@ -245,9 +246,29 @@ def _write_extractor_telemetry(diagnostics_writer, logger: Logger):
     except OSError as ex:
         logger.warning("Failed to write extractor telemetry: %s", ex)
 
+def _write_parser_statistics_telemetry(diagnostics_writer, logger: Logger):
+    counts = diagnostics_writer.parser_statistics()
+    if counts == (0, 0):
+        return
+    try:
+        diagnostics_writer.write(parser_statistics_telemetry_message(*counts))
+    except OSError as ex:
+        logger.warning("Failed to write parser statistics telemetry: %s", ex)
+
 class DiagnosticsWriter(object):
     def __init__(self, proc_id):
         self.proc_id = proc_id
+        self.old_parser_file_count = 0
+        self.tree_sitter_parser_file_count = 0
+
+    def record_old_parser(self):
+        self.old_parser_file_count += 1
+
+    def record_tree_sitter_parser(self):
+        self.tree_sitter_parser_file_count += 1
+
+    def parser_statistics(self):
+        return self.old_parser_file_count, self.tree_sitter_parser_file_count
 
     def write(self, message):
         dir = os.environ.get("CODEQL_EXTRACTOR_PYTHON_DIAGNOSTIC_DIR")
@@ -286,7 +307,7 @@ def _extract_loop(proc_id, queue, trap_dir, archive, options, reply_queue, logge
         _write_extractor_telemetry(diagnostics_writer, logger)
     try:
         if options.trace_only:
-            extractor = ModulePrinter(options, trap_dir, archive, renamer, logger)
+            extractor = ModulePrinter(options, trap_dir, archive, renamer, logger, diagnostics_writer)
         else:
             extractor = SuperExtractor(options, trap_dir, archive, renamer, logger, diagnostics_writer)
         profiler = get_profiler(options, id, logger)
@@ -299,6 +320,7 @@ def _extract_loop(proc_id, queue, trap_dir, archive, options, reply_queue, logge
                     if write_global_data:
                         extractor.write_global_data()
                     extractor.close()
+                    _write_parser_statistics_telemetry(diagnostics_writer, logger)
                     return
                 try:
                     start = time.time()
@@ -352,4 +374,5 @@ def _extract_loop(proc_id, queue, trap_dir, archive, options, reply_queue, logge
     except _Empty:
         #Cleared queue enough to avoid deadlock.
         pass
+    _write_parser_statistics_telemetry(diagnostics_writer, logger)
     sys.exit(2)

python/extractor/tests/test_diagnostics.py

@@ -3,6 +3,7 @@
 from semmle import logging
 from semmle import util
 from semmle import worker
+from semmle.python.modules import PythonSourceModule
 
 
 def test_extractor_telemetry_message(mocker):
@@ -32,13 +33,44 @@ def test_extractor_telemetry_message(mocker):
     }
 
 
+def test_parser_statistics_telemetry_message():
+    message = logging.parser_statistics_telemetry_message(
+        old_parser_file_count=12, tree_sitter_parser_file_count=3
+    ).to_dict()
+    message.pop("timestamp")
+
+    assert message == {
+        "source": {
+            "id": "py/extractor/parser-statistics",
+            "name": "Python parser statistics",
+            "extractorName": "python",
+        },
+        "severity": "note",
+        "markdownMessage": "Internal parser telemetry for the Python extractor.\n\nNo action needed.",
+        "visibility": {
+            "statusPage": False,
+            "cliSummaryTable": False,
+            "telemetry": True,
+        },
+        "attributes": {
+            "old_parser_file_count": 12,
+            "tree_sitter_parser_file_count": 3,
+        },
+    }
+
+
 def test_write_extractor_telemetry(mocker):
     diagnostics_writer = mocker.Mock()
     logger = mocker.Mock()
 
     worker._write_extractor_telemetry(diagnostics_writer, logger)
 
     diagnostics_writer.write.assert_called_once()
+    assert diagnostics_writer.write.call_args.args[0].to_dict()["attributes"] == {
+        "python_analysis_version": util.get_analysis_version(),
+        "python_runtime_version": platform.python_version(),
+        "extractor_version": util.VERSION,
+    }
     logger.warning.assert_not_called()
 
 
@@ -52,3 +84,76 @@ def test_write_extractor_telemetry_handles_io_error(mocker):
     logger.warning.assert_called_once_with(
         "Failed to write extractor telemetry: %s", diagnostics_writer.write.side_effect
     )
+
+
+def test_write_parser_statistics_telemetry(mocker):
+    diagnostics_writer = mocker.Mock()
+    diagnostics_writer.parser_statistics.return_value = (1, 1)
+    logger = mocker.Mock()
+
+    worker._write_parser_statistics_telemetry(diagnostics_writer, logger)
+
+    diagnostics_writer.write.assert_called_once()
+    assert diagnostics_writer.write.call_args.args[0].to_dict()["attributes"] == {
+        "old_parser_file_count": 1,
+        "tree_sitter_parser_file_count": 1,
+    }
+    logger.warning.assert_not_called()
+
+
+def test_does_not_write_empty_parser_statistics_telemetry(mocker):
+    diagnostics_writer = mocker.Mock()
+    diagnostics_writer.parser_statistics.return_value = (0, 0)
+    logger = mocker.Mock()
+
+    worker._write_parser_statistics_telemetry(diagnostics_writer, logger)
+
+    diagnostics_writer.write.assert_not_called()
+    logger.warning.assert_not_called()
+
+
+def test_records_old_parser_usage_once(mocker, monkeypatch):
+    monkeypatch.delenv("CODEQL_PYTHON_DISABLE_OLD_PARSER", raising=False)
+    monkeypatch.delenv("CODEQL_PYTHON_DISABLE_TSG_PARSER", raising=False)
+    old_ast = object()
+    mocker.patch("semmle.python.parser.parse", return_value=old_ast)
+    diagnostics_writer = worker.DiagnosticsWriter(1)
+    module = PythonSourceModule(
+        None,
+        "test.py",
+        mocker.Mock(),
+        diagnostics_writer,
+        bytes_source=b"x = 1\n",
+    )
+
+    parsed_ast = module.py_ast
+    # Access the cached AST again to verify that it is not counted twice.
+    _ = module.py_ast
+
+    assert parsed_ast is old_ast
+    assert diagnostics_writer.parser_statistics() == (1, 0)
+
+
+def test_records_tree_sitter_parser_usage_once(mocker, monkeypatch):
+    monkeypatch.delenv("CODEQL_PYTHON_DISABLE_OLD_PARSER", raising=False)
+    monkeypatch.delenv("CODEQL_PYTHON_DISABLE_TSG_PARSER", raising=False)
+    tree_sitter_ast = object()
+    mocker.patch("semmle.python.parser.parse", side_effect=SyntaxError("old parser failed"))
+    mocker.patch(
+        "semmle.python.parser.tsg_parser.parse", return_value=tree_sitter_ast
+    )
+    diagnostics_writer = worker.DiagnosticsWriter(1)
+    module = PythonSourceModule(
+        None,
+        "test.py",
+        mocker.Mock(),
+        diagnostics_writer,
+        bytes_source=b"x = 1\n",
+    )
+
+    parsed_ast = module.py_ast
+    # Access the cached AST again to verify that it is not counted twice.
+    _ = module.py_ast
+
+    assert parsed_ast is tree_sitter_ast
+    assert diagnostics_writer.parser_statistics() == (0, 1)