From a116c55199dfb64f180690bb6eb3c219ca677ca7 Mon Sep 17 00:00:00 2001 From: Ivana Kellyerova Date: Wed, 20 Mar 2024 10:56:12 +0100 Subject: [PATCH] feat: Add optional `keep_alive` (#2842) --- sentry_sdk/consts.py | 1 + sentry_sdk/transport.py | 35 +++++++++++++++++++++-- tests/test_transport.py | 62 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 95 insertions(+), 3 deletions(-) diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py index 83076c762f..6af08b4a40 100644 --- a/sentry_sdk/consts.py +++ b/sentry_sdk/consts.py @@ -264,6 +264,7 @@ def __init__( ignore_errors=[], # type: Sequence[Union[type, str]] # noqa: B006 max_request_body_size="medium", # type: str socket_options=None, # type: Optional[List[Tuple[int, int, int | bytes]]] + keep_alive=False, # type: bool before_send=None, # type: Optional[EventProcessor] before_breadcrumb=None, # type: Optional[BreadcrumbProcessor] debug=None, # type: Optional[bool] diff --git a/sentry_sdk/transport.py b/sentry_sdk/transport.py index b924ae502a..9ea9cd0c98 100644 --- a/sentry_sdk/transport.py +++ b/sentry_sdk/transport.py @@ -2,6 +2,7 @@ import io import gzip +import socket import time from datetime import timedelta from collections import defaultdict @@ -21,6 +22,7 @@ from typing import Callable from typing import Dict from typing import Iterable + from typing import List from typing import Optional from typing import Tuple from typing import Type @@ -40,6 +42,21 @@ from urllib import getproxies # type: ignore +KEEP_ALIVE_SOCKET_OPTIONS = [] +for option in [ + (socket.SOL_SOCKET, lambda: getattr(socket, "SO_KEEPALIVE"), 1), # noqa: B009 + (socket.SOL_TCP, lambda: getattr(socket, "TCP_KEEPIDLE"), 45), # noqa: B009 + (socket.SOL_TCP, lambda: getattr(socket, "TCP_KEEPINTVL"), 10), # noqa: B009 + (socket.SOL_TCP, lambda: getattr(socket, "TCP_KEEPCNT"), 6), # noqa: B009 +]: + try: + KEEP_ALIVE_SOCKET_OPTIONS.append((option[0], option[1](), option[2])) + except AttributeError: + # a specific option might not be available on specific systems, + # e.g. TCP_KEEPIDLE doesn't exist on macOS + pass + + class Transport(object): """Baseclass for all transports. @@ -446,8 +463,22 @@ def _get_pool_options(self, ca_certs): "ca_certs": ca_certs or certifi.where(), } - if self.options["socket_options"]: - options["socket_options"] = self.options["socket_options"] + socket_options = None # type: Optional[List[Tuple[int, int, int | bytes]]] + + if self.options["socket_options"] is not None: + socket_options = self.options["socket_options"] + + if self.options["keep_alive"]: + if socket_options is None: + socket_options = [] + + used_options = {(o[0], o[1]) for o in socket_options} + for default_option in KEEP_ALIVE_SOCKET_OPTIONS: + if (default_option[0], default_option[1]) not in used_options: + socket_options.append(default_option) + + if socket_options is not None: + options["socket_options"] = socket_options return options diff --git a/tests/test_transport.py b/tests/test_transport.py index aa471b9081..c1f70b0108 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -13,7 +13,7 @@ from sentry_sdk import Hub, Client, add_breadcrumb, capture_message, Scope from sentry_sdk._compat import datetime_utcnow -from sentry_sdk.transport import _parse_rate_limits +from sentry_sdk.transport import KEEP_ALIVE_SOCKET_OPTIONS, _parse_rate_limits from sentry_sdk.envelope import Envelope, parse_json from sentry_sdk.integrations.logging import LoggingIntegration @@ -167,6 +167,66 @@ def test_socket_options(make_client): assert options["socket_options"] == socket_options +def test_keep_alive_true(make_client): + client = make_client(keep_alive=True) + + options = client.transport._get_pool_options([]) + assert options["socket_options"] == KEEP_ALIVE_SOCKET_OPTIONS + + +def test_keep_alive_off_by_default(make_client): + client = make_client() + options = client.transport._get_pool_options([]) + assert "socket_options" not in options + + +def test_socket_options_override_keep_alive(make_client): + socket_options = [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + (socket.SOL_TCP, socket.TCP_KEEPINTVL, 10), + (socket.SOL_TCP, socket.TCP_KEEPCNT, 6), + ] + + client = make_client(socket_options=socket_options, keep_alive=False) + + options = client.transport._get_pool_options([]) + assert options["socket_options"] == socket_options + + +def test_socket_options_merge_with_keep_alive(make_client): + socket_options = [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 42), + (socket.SOL_TCP, socket.TCP_KEEPINTVL, 42), + ] + + client = make_client(socket_options=socket_options, keep_alive=True) + + options = client.transport._get_pool_options([]) + try: + assert options["socket_options"] == [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 42), + (socket.SOL_TCP, socket.TCP_KEEPINTVL, 42), + (socket.SOL_TCP, socket.TCP_KEEPIDLE, 45), + (socket.SOL_TCP, socket.TCP_KEEPCNT, 6), + ] + except AttributeError: + assert options["socket_options"] == [ + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 42), + (socket.SOL_TCP, socket.TCP_KEEPINTVL, 42), + (socket.SOL_TCP, socket.TCP_KEEPCNT, 6), + ] + + +def test_socket_options_override_defaults(make_client): + # If socket_options are set to [], this doesn't mean the user doesn't want + # any custom socket_options, but rather that they want to disable the urllib3 + # socket option defaults, so we need to set this and not ignore it. + client = make_client(socket_options=[]) + + options = client.transport._get_pool_options([]) + assert options["socket_options"] == [] + + def test_transport_infinite_loop(capturing_server, request, make_client): client = make_client( debug=True,