common/grpc_client_manager.py

"""
Copyright 2020 The Magma Authors.

This source code is licensed under the BSD-style license found in the
LICENSE file in the root directory of this source tree.

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.
"""

import logging
import sys

import grpc
import psutil
from common.service_registry import ServiceRegistry


class GRPCClientManager:
    def __init__(
        self, service_name: str, service_stub,
        max_client_reuse: int = 60,
    ):
        self._client = None
        self._service_stub = service_stub
        self._service_name = service_name
        self._num_client_use = 0
        self._max_client_reuse = max_client_reuse

    def get_client(self):
        """
        get_client returns a grpc client of the specified service in the cloud.
        it will return a recycled client until the client fails or the number
        of recycling reaches the max_client_use.
        """
        if self._client is None or \
                self._num_client_use > self._max_client_reuse:
            chan = ServiceRegistry.get_rpc_channel(
                self._service_name,
                ServiceRegistry.CLOUD,
            )
            self._client = self._service_stub(chan)
            self._num_client_use = 0

        self._num_client_use += 1
        return self._client

    def on_grpc_fail(self, err_code):
        """
        Try to reuse the grpc client if possible. We are yet to fix a
        grpc behavior, where if DNS request blackholes then the DNS request
        is retried infinitely even after the channel is deleted. To prevent
        running out of fds, we try to reuse the channel during such failures
        as much as possible.
        """
        if err_code != grpc.StatusCode.DEADLINE_EXCEEDED:
            # Not related to the DNS issue
            self._reset_client()
        if self._num_client_use >= self._max_client_reuse:
            logging.info('Max client reuse reached. Cleaning up client')
            self._reset_client()

            # Sanity check if we are not leaking fds
            proc = psutil.Process()
            max_fds, _ = proc.rlimit(psutil.RLIMIT_NOFILE)
            open_fds = proc.num_fds()
            logging.info('Num open fds: %d', open_fds)
            if open_fds >= (max_fds * 0.8):
                logging.error("Reached 80% of allowed fds. Restarting process")
                sys.exit(1)

    def _reset_client(self):
        self._client = None
        self._num_client_use = 0