src/llvm-project/clang-tools-extra/clangd/test/remote-index/pipeline_helper.py

#!/usr/bin/env python
#
# ===- pipeline_helper.py - Remote Index pipeline Helper *- python -------*--===#
#
# Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
#
# ===------------------------------------------------------------------------===#

import argparse
import os
import subprocess
from socket import socket
import sys
import time
import threading


def kill_process_after_delay(server_process):
    time.sleep(10)
    if server_process.poll() is None:
        server_process.kill()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-file-name", required=True)
    parser.add_argument("--project-root", required=True)
    parser.add_argument("--index-file", required=True)
    parser.add_argument("--server-arg", action="append", default=[])
    parser.add_argument(
        "--server-log", nargs="?", type=argparse.FileType("wb"), default=os.devnull
    )

    args = parser.parse_args()

    # Grab an available port.
    with socket() as s:
        s.bind(("localhost", 0))
        server_address = "localhost:" + str(s.getsockname()[1])

    print("Initializing clangd-index-server...", file=sys.stderr)
    index_server_process = subprocess.Popen(
        [
            "clangd-index-server",
            "--server-address=" + server_address,
            args.index_file,
            args.project_root,
        ]
        + args.server_arg,
        stderr=subprocess.PIPE,
    )

    # This will kill index_server_process if it hangs without printing init
    # message.
    shutdown_thread = threading.Thread(
        target=kill_process_after_delay, args=(index_server_process,)
    )
    shutdown_thread.daemon = True
    shutdown_thread.start()

    # Wait for the server to warm-up.
    found_init_message = False
    while index_server_process.poll() is None:
        line = index_server_process.stderr.readline()
        args.server_log.write(line)
        args.server_log.flush()
        if b"Server listening" in line:
            print("Server initialization complete.", file=sys.stderr)
            found_init_message = True
            break

    if not found_init_message:
        print("Server initialization failed. Shutting down.", file=sys.stderr)
        sys.exit(1)

    print("Running clangd-index-server-monitor...", file=sys.stderr)
    index_server_monitor_process = subprocess.Popen(
        [
            "clangd-index-server-monitor",
            server_address,
        ],
        stderr=subprocess.PIPE,
    )

    index_server_monitor_process.wait()

    in_file = open(args.input_file_name)

    print("Staring clangd...", file=sys.stderr)
    clangd_process = subprocess.Popen(
        [
            "clangd",
            "--remote-index-address=" + server_address,
            "--project-root=" + args.project_root,
            "--lit-test",
            "--sync",
        ],
        stdin=in_file,
    )
    clangd_process.wait()
    print(
        "Clangd executed successfully, shutting down child processes.", file=sys.stderr
    )
    index_server_process.kill()
    for line in index_server_process.stderr:
        args.server_log.write(line)
        args.server_log.flush()

    for line in index_server_monitor_process.stderr:
        args.server_log.write(line)
        args.server_log.flush()


if __name__ == "__main__":
    main()