src/python/grpcio_tests/tests/fork/methods.py

# Copyright 2018 gRPC authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementations of fork support test methods."""

import enum
import json
import logging
import multiprocessing
import os
import queue
import subprocess
import sys
import tempfile
import threading
import time
import traceback

import grpc

from src.proto.grpc.testing import empty_pb2
from src.proto.grpc.testing import messages_pb2
from src.proto.grpc.testing import test_pb2_grpc

_LOGGER = logging.getLogger(__name__)
_RPC_TIMEOUT_S = 10
_CHILD_FINISH_TIMEOUT_S = 20
_GDB_TIMEOUT_S = 60


def _channel(args):
    target = "{}:{}".format(args["server_host"], args["server_port"])
    if args["use_tls"]:
        channel_credentials = grpc.ssl_channel_credentials()
        channel = grpc.secure_channel(target, channel_credentials)
    else:
        channel = grpc.insecure_channel(target)
    return channel


def _validate_payload_type_and_length(response, expected_type, expected_length):
    if response.payload.type is not expected_type:
        raise ValueError(
            "expected payload type %s, got %s"
            % (expected_type, type(response.payload.type))
        )
    elif len(response.payload.body) != expected_length:
        raise ValueError(
            "expected payload body size %d, got %d"
            % (expected_length, len(response.payload.body))
        )


def _async_unary(stub):
    size = 314159
    request = messages_pb2.SimpleRequest(
        response_type=messages_pb2.COMPRESSABLE,
        response_size=size,
        payload=messages_pb2.Payload(body=b"\x00" * 271828),
    )

    response_future = stub.UnaryCall.future(request, timeout=_RPC_TIMEOUT_S)
    response = response_future.result()
    _validate_payload_type_and_length(response, messages_pb2.COMPRESSABLE, size)


def _blocking_unary(stub):
    size = 314159
    request = messages_pb2.SimpleRequest(
        response_type=messages_pb2.COMPRESSABLE,
        response_size=size,
        payload=messages_pb2.Payload(body=b"\x00" * 271828),
    )
    response = stub.UnaryCall(request, timeout=_RPC_TIMEOUT_S)
    _validate_payload_type_and_length(response, messages_pb2.COMPRESSABLE, size)


class _Pipe(object):
    def __init__(self):
        self._condition = threading.Condition()
        self._values = []
        self._open = True

    def __iter__(self):
        return self

    def __next__(self):
        return self.next()

    def next(self):
        with self._condition:
            while not self._values and self._open:
                self._condition.wait()
            if self._values:
                return self._values.pop(0)
            else:
                raise StopIteration()

    def add(self, value):
        with self._condition:
            self._values.append(value)
            self._condition.notify()

    def close(self):
        with self._condition:
            self._open = False
            self._condition.notify()

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        self.close()


class _ChildProcess(object):
    def __init__(self, task, args=None):
        if args is None:
            args = ()
        self._exceptions = multiprocessing.Queue()
        self._stdout_path = tempfile.mkstemp()[1]
        self._stderr_path = tempfile.mkstemp()[1]
        self._child_pid = None
        self._rc = None
        self._args = args

        self._task = task

    def _child_main(self):
        import faulthandler

        faulthandler.enable(all_threads=True)

        try:
            self._task(*self._args)
        except grpc.RpcError as rpc_error:
            traceback.print_exc()
            self._exceptions.put("RpcError: %s" % rpc_error)
        except Exception as e:  # pylint: disable=broad-except
            traceback.print_exc()
            self._exceptions.put(e)
        sys.exit(0)

    def _orchestrate_child_gdb(self):
        cmd = [
            "gdb",
            "-ex",
            "set confirm off",
            "-ex",
            "attach {}".format(os.getpid()),
            "-ex",
            "set follow-fork-mode child",
            "-ex",
            "continue",
            "-ex",
            "bt",
        ]
        streams = tuple(tempfile.TemporaryFile() for _ in range(2))
        sys.stderr.write("Invoking gdb\n")
        sys.stderr.flush()
        process = subprocess.Popen(cmd, stdout=sys.stderr, stderr=sys.stderr)
        time.sleep(5)

    def start(self):
        # NOTE: Try uncommenting the following line if the child is segfaulting.
        # self._orchestrate_child_gdb()
        ret = os.fork()
        if ret == 0:
            self._child_main()
        else:
            self._child_pid = ret

    def wait(self, timeout):
        total = 0.0
        wait_interval = 1.0
        while total < timeout:
            ret, termination = os.waitpid(self._child_pid, os.WNOHANG)
            if ret == self._child_pid:
                self._rc = termination
                return True
            time.sleep(wait_interval)
            total += wait_interval
        else:
            return False

    def _print_backtraces(self):
        cmd = [
            "gdb",
            "-ex",
            "set confirm off",
            "-ex",
            "echo attaching",
            "-ex",
            "attach {}".format(self._child_pid),
            "-ex",
            "echo print_backtrace",
            "-ex",
            "thread apply all bt",
            "-ex",
            "echo printed_backtrace",
            "-ex",
            "quit",
        ]
        streams = tuple(tempfile.TemporaryFile() for _ in range(2))
        sys.stderr.write("Invoking gdb\n")
        sys.stderr.flush()
        process = subprocess.Popen(cmd, stdout=streams[0], stderr=streams[1])
        try:
            process.wait(timeout=_GDB_TIMEOUT_S)
        except subprocess.TimeoutExpired:
            sys.stderr.write("gdb stacktrace generation timed out.\n")
        finally:
            for stream_name, stream in zip(("STDOUT", "STDERR"), streams):
                stream.seek(0)
                sys.stderr.write(
                    "gdb {}:\n{}\n".format(
                        stream_name, stream.read().decode("ascii")
                    )
                )
                stream.close()
            sys.stderr.flush()

    def finish(self):
        terminated = self.wait(_CHILD_FINISH_TIMEOUT_S)
        sys.stderr.write("Exit code: {}\n".format(self._rc))
        if not terminated:
            self._print_backtraces()
            raise RuntimeError("Child process did not terminate")
        if self._rc != 0:
            raise ValueError("Child process failed with exitcode %d" % self._rc)
        try:
            exception = self._exceptions.get(block=False)
            raise ValueError(
                'Child process failed: "%s": "%s"'
                % (repr(exception), exception)
            )
        except queue.Empty:
            pass


def _async_unary_same_channel(channel):
    def child_target():
        try:
            _async_unary(stub)
            raise Exception(
                "Child should not be able to re-use channel after fork"
            )
        except ValueError as expected_value_error:
            pass

    stub = test_pb2_grpc.TestServiceStub(channel)
    _async_unary(stub)
    child_process = _ChildProcess(child_target)
    child_process.start()
    _async_unary(stub)
    child_process.finish()


def _async_unary_new_channel(channel, args):
    def child_target():
        with _channel(args) as child_channel:
            child_stub = test_pb2_grpc.TestServiceStub(child_channel)
            _async_unary(child_stub)
            child_channel.close()

    stub = test_pb2_grpc.TestServiceStub(channel)
    _async_unary(stub)
    child_process = _ChildProcess(child_target)
    child_process.start()
    _async_unary(stub)
    child_process.finish()


def _blocking_unary_same_channel(channel):
    def child_target():
        try:
            _blocking_unary(stub)
            raise Exception(
                "Child should not be able to re-use channel after fork"
            )
        except ValueError as expected_value_error:
            pass

    stub = test_pb2_grpc.TestServiceStub(channel)
    _blocking_unary(stub)
    child_process = _ChildProcess(child_target)
    child_process.start()
    child_process.finish()


def _blocking_unary_new_channel(channel, args):
    def child_target():
        with _channel(args) as child_channel:
            child_stub = test_pb2_grpc.TestServiceStub(child_channel)
            _blocking_unary(child_stub)

    stub = test_pb2_grpc.TestServiceStub(channel)
    _blocking_unary(stub)
    child_process = _ChildProcess(child_target)
    child_process.start()
    _blocking_unary(stub)
    child_process.finish()


# Verify that the fork channel registry can handle already closed channels
def _close_channel_before_fork(channel, args):
    def child_target():
        new_channel.close()
        with _channel(args) as child_channel:
            child_stub = test_pb2_grpc.TestServiceStub(child_channel)
            _blocking_unary(child_stub)

    stub = test_pb2_grpc.TestServiceStub(channel)
    _blocking_unary(stub)
    channel.close()

    with _channel(args) as new_channel:
        new_stub = test_pb2_grpc.TestServiceStub(new_channel)
        child_process = _ChildProcess(child_target)
        child_process.start()
        _blocking_unary(new_stub)
        child_process.finish()


def _connectivity_watch(channel, args):
    parent_states = []
    parent_channel_ready_event = threading.Event()

    def child_target():
        child_channel_ready_event = threading.Event()

        def child_connectivity_callback(state):
            if state is grpc.ChannelConnectivity.READY:
                child_channel_ready_event.set()

        with _channel(args) as child_channel:
            child_stub = test_pb2_grpc.TestServiceStub(child_channel)
            child_channel.subscribe(child_connectivity_callback)
            _async_unary(child_stub)
            if not child_channel_ready_event.wait(timeout=_RPC_TIMEOUT_S):
                raise ValueError("Channel did not move to READY")
            if len(parent_states) > 1:
                raise ValueError(
                    "Received connectivity updates on parent callback",
                    parent_states,
                )
            child_channel.unsubscribe(child_connectivity_callback)

    def parent_connectivity_callback(state):
        parent_states.append(state)
        if state is grpc.ChannelConnectivity.READY:
            parent_channel_ready_event.set()

    channel.subscribe(parent_connectivity_callback)
    stub = test_pb2_grpc.TestServiceStub(channel)
    child_process = _ChildProcess(child_target)
    child_process.start()
    _async_unary(stub)
    if not parent_channel_ready_event.wait(timeout=_RPC_TIMEOUT_S):
        raise ValueError("Channel did not move to READY")
    channel.unsubscribe(parent_connectivity_callback)
    child_process.finish()


def _ping_pong_with_child_processes_after_first_response(
    channel, args, child_target, run_after_close=True
):
    request_response_sizes = (
        31415,
        9,
        2653,
        58979,
    )
    request_payload_sizes = (
        27182,
        8,
        1828,
        45904,
    )
    stub = test_pb2_grpc.TestServiceStub(channel)
    pipe = _Pipe()
    parent_bidi_call = stub.FullDuplexCall(pipe)
    child_processes = []
    first_message_received = False
    for response_size, payload_size in zip(
        request_response_sizes, request_payload_sizes
    ):
        request = messages_pb2.StreamingOutputCallRequest(
            response_type=messages_pb2.COMPRESSABLE,
            response_parameters=(
                messages_pb2.ResponseParameters(size=response_size),
            ),
            payload=messages_pb2.Payload(body=b"\x00" * payload_size),
        )
        pipe.add(request)
        if first_message_received:
            child_process = _ChildProcess(
                child_target, (parent_bidi_call, channel, args)
            )
            child_process.start()
            child_processes.append(child_process)
        response = next(parent_bidi_call)
        first_message_received = True
        child_process = _ChildProcess(
            child_target, (parent_bidi_call, channel, args)
        )
        child_process.start()
        child_processes.append(child_process)
        _validate_payload_type_and_length(
            response, messages_pb2.COMPRESSABLE, response_size
        )
    pipe.close()
    if run_after_close:
        child_process = _ChildProcess(
            child_target, (parent_bidi_call, channel, args)
        )
        child_process.start()
        child_processes.append(child_process)
    for child_process in child_processes:
        child_process.finish()


def _in_progress_bidi_continue_call(channel):
    def child_target(parent_bidi_call, parent_channel, args):
        stub = test_pb2_grpc.TestServiceStub(parent_channel)
        try:
            _async_unary(stub)
            raise Exception(
                "Child should not be able to re-use channel after fork"
            )
        except ValueError as expected_value_error:
            pass
        inherited_code = parent_bidi_call.code()
        inherited_details = parent_bidi_call.details()
        if inherited_code != grpc.StatusCode.CANCELLED:
            raise ValueError(
                "Expected inherited code CANCELLED, got %s" % inherited_code
            )
        if inherited_details != "Channel closed due to fork":
            raise ValueError(
                "Expected inherited details Channel closed due to fork, got %s"
                % inherited_details
            )

    # Don't run child_target after closing the parent call, as the call may have
    # received a status from the  server before fork occurs.
    _ping_pong_with_child_processes_after_first_response(
        channel, None, child_target, run_after_close=False
    )


def _in_progress_bidi_same_channel_async_call(channel):
    def child_target(parent_bidi_call, parent_channel, args):
        stub = test_pb2_grpc.TestServiceStub(parent_channel)
        try:
            _async_unary(stub)
            raise Exception(
                "Child should not be able to re-use channel after fork"
            )
        except ValueError as expected_value_error:
            pass

    _ping_pong_with_child_processes_after_first_response(
        channel, None, child_target
    )


def _in_progress_bidi_same_channel_blocking_call(channel):
    def child_target(parent_bidi_call, parent_channel, args):
        stub = test_pb2_grpc.TestServiceStub(parent_channel)
        try:
            _blocking_unary(stub)
            raise Exception(
                "Child should not be able to re-use channel after fork"
            )
        except ValueError as expected_value_error:
            pass

    _ping_pong_with_child_processes_after_first_response(
        channel, None, child_target
    )


def _in_progress_bidi_new_channel_async_call(channel, args):
    def child_target(parent_bidi_call, parent_channel, args):
        with _channel(args) as channel:
            stub = test_pb2_grpc.TestServiceStub(channel)
            _async_unary(stub)

    _ping_pong_with_child_processes_after_first_response(
        channel, args, child_target
    )


def _in_progress_bidi_new_channel_blocking_call(channel, args):
    def child_target(parent_bidi_call, parent_channel, args):
        with _channel(args) as channel:
            stub = test_pb2_grpc.TestServiceStub(channel)
            _blocking_unary(stub)

    _ping_pong_with_child_processes_after_first_response(
        channel, args, child_target
    )


@enum.unique
class TestCase(enum.Enum):
    CONNECTIVITY_WATCH = "connectivity_watch"
    CLOSE_CHANNEL_BEFORE_FORK = "close_channel_before_fork"
    ASYNC_UNARY_SAME_CHANNEL = "async_unary_same_channel"
    ASYNC_UNARY_NEW_CHANNEL = "async_unary_new_channel"
    BLOCKING_UNARY_SAME_CHANNEL = "blocking_unary_same_channel"
    BLOCKING_UNARY_NEW_CHANNEL = "blocking_unary_new_channel"
    IN_PROGRESS_BIDI_CONTINUE_CALL = "in_progress_bidi_continue_call"
    IN_PROGRESS_BIDI_SAME_CHANNEL_ASYNC_CALL = (
        "in_progress_bidi_same_channel_async_call"
    )
    IN_PROGRESS_BIDI_SAME_CHANNEL_BLOCKING_CALL = (
        "in_progress_bidi_same_channel_blocking_call"
    )
    IN_PROGRESS_BIDI_NEW_CHANNEL_ASYNC_CALL = (
        "in_progress_bidi_new_channel_async_call"
    )
    IN_PROGRESS_BIDI_NEW_CHANNEL_BLOCKING_CALL = (
        "in_progress_bidi_new_channel_blocking_call"
    )

    def run_test(self, args):
        _LOGGER.info("Running %s", self)
        channel = _channel(args)
        if self is TestCase.ASYNC_UNARY_SAME_CHANNEL:
            _async_unary_same_channel(channel)
        elif self is TestCase.ASYNC_UNARY_NEW_CHANNEL:
            _async_unary_new_channel(channel, args)
        elif self is TestCase.BLOCKING_UNARY_SAME_CHANNEL:
            _blocking_unary_same_channel(channel)
        elif self is TestCase.BLOCKING_UNARY_NEW_CHANNEL:
            _blocking_unary_new_channel(channel, args)
        elif self is TestCase.CLOSE_CHANNEL_BEFORE_FORK:
            _close_channel_before_fork(channel, args)
        elif self is TestCase.CONNECTIVITY_WATCH:
            _connectivity_watch(channel, args)
        elif self is TestCase.IN_PROGRESS_BIDI_CONTINUE_CALL:
            _in_progress_bidi_continue_call(channel)
        elif self is TestCase.IN_PROGRESS_BIDI_SAME_CHANNEL_ASYNC_CALL:
            _in_progress_bidi_same_channel_async_call(channel)
        elif self is TestCase.IN_PROGRESS_BIDI_SAME_CHANNEL_BLOCKING_CALL:
            _in_progress_bidi_same_channel_blocking_call(channel)
        elif self is TestCase.IN_PROGRESS_BIDI_NEW_CHANNEL_ASYNC_CALL:
            _in_progress_bidi_new_channel_async_call(channel, args)
        elif self is TestCase.IN_PROGRESS_BIDI_NEW_CHANNEL_BLOCKING_CALL:
            _in_progress_bidi_new_channel_blocking_call(channel, args)
        else:
            raise NotImplementedError(
                'Test case "%s" not implemented!' % self.name
            )
        channel.close()


# Useful if needing to find a block of code from an address in an SO.
def dump_object_map():
    with open("/proc/self/maps", "r") as f:
        sys.stderr.write("=============== /proc/self/maps ===============\n")
        sys.stderr.write(f.read())
        sys.stderr.write("\n")
        sys.stderr.flush()