"""gRPC channel pool."""
from __future__ import annotations
import ipaddress
import re
import sys
from threading import Lock
from types import TracebackType
from typing import TYPE_CHECKING, Literal
from urllib.parse import urlparse
import grpc
from ni_grpc_extensions.loggers import ClientLogger
if TYPE_CHECKING:
if sys.version_info >= (3, 11):
from typing import Self
else:
from typing_extensions import Self
[docs]
class GrpcChannelPool:
"""Class that manages :any:`grpc.Channel` lifetimes."""
def __init__(self) -> None:
"""Initialize the :class:`GrpcChannelPool` object."""
self._lock: Lock = Lock()
self._channel_cache: dict[str, grpc.Channel] = {}
[docs]
def __enter__(self: Self) -> Self:
"""Enter the runtime context of the GrpcChannelPool."""
return self
[docs]
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_val: BaseException | None,
traceback: TracebackType | None,
) -> Literal[False]:
"""Exit the runtime context of the GrpcChannelPool."""
self.close()
return False
[docs]
def get_channel(self, target: str) -> grpc.Channel:
"""Return a gRPC channel.
Args:
target (str): The server address
"""
new_channel = None
with self._lock:
if target not in self._channel_cache:
self._lock.release()
new_channel = self._create_channel(target)
self._lock.acquire()
if target not in self._channel_cache:
self._channel_cache[target] = new_channel
new_channel = None
channel = self._channel_cache[target]
# Close new_channel if it was not stored in _channel_cache.
if new_channel is not None:
new_channel.close()
return channel
[docs]
def close(self) -> None:
"""Close channels opened by get_channel()."""
with self._lock:
for channel in self._channel_cache.values():
channel.close()
self._channel_cache.clear()
def _create_channel(self, target: str) -> grpc.Channel:
options = [
("grpc.max_receive_message_length", -1),
("grpc.max_send_message_length", -1),
]
if self._is_local(target):
options.append(("grpc.enable_http_proxy", 0))
channel = grpc.insecure_channel(target, options)
if ClientLogger.is_enabled():
channel = grpc.intercept_channel(channel, ClientLogger())
return channel
def _is_local(self, target: str) -> bool:
hostname = ""
# First, check if the target string is in URL format
parse_result = urlparse(target)
if parse_result.scheme and parse_result.hostname and parse_result.port:
hostname = parse_result.hostname
else:
# Next, check for target string in <host_name>:<port> format
match = re.match(r"^(.*):(\d+)$", target)
if match:
hostname = match.group(1)
if not hostname:
return False
if hostname == "localhost" or hostname == "LOCALHOST":
return True
# IPv6 addresses don't support parsing with leading/trailing brackets
# so we need to remove them.
match = re.match(r"^\[(.*)\]$", hostname)
if match:
hostname = match.group(1)
try:
address = ipaddress.ip_address(hostname)
return address.is_loopback
except ValueError:
return False