GeezerVison is showing github.com — refactor: Share one dual-stack socket builder across REST, ui, lineage · feast-dev/feast@a4d478f
File treeExpand file treeCollapse file tree Open diff view settings Expand file treeCollapse file tree Open diff view settings Collapse file Copy file name to clipboardExpand all lines: sdk/python/feast/api/registry/rest/rest_registry_server.py +2-28Lines changed: 2 additions & 28 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 1 | 1 | import json | 2 | 2 | import logging | 3 | 3 | import re | 4 | | -import socket | 5 | 4 |
| 6 | 5 | from fastapi import Depends, FastAPI, HTTPException, Request, status | 7 | 6 | from fastapi.exceptions import RequestValidationError | | 24 | 23 | str_to_auth_manager_type, | 25 | 24 | ) | 26 | 25 | from feast.registry_server import RegistryServer | 27 | | -from feast.utils import _ipv6_available, _utc_now | | 26 | +from feast.utils import _make_dual_stack_socket, _utc_now | 28 | 27 |
| 29 | 28 | logger = logging.getLogger(__name__) | 30 | 29 | logger.setLevel(logging.INFO) | 31 | 30 |
| 32 | 31 |
| 33 | | -def _make_rest_socket(port: int) -> socket.socket: | 34 | | - """Build the REST registry server's listening socket, dual-stack ("::") | 35 | | - when the host supports IPv6, or IPv4-only otherwise (e.g. ipv6.disable=1 | 36 | | - kernels), matching the gRPC registry server's "[::]" reachability. | 37 | | -
| 38 | | - uvicorn's own ``host="::"`` binds IPv6-only: asyncio's | 39 | | - ``loop.create_server`` sets ``IPV6_V6ONLY=1`` on any socket it creates | 40 | | - itself from a host string, which would silently drop IPv4 clients on | 41 | | - every host, not just IPv6-less ones. So the socket is bound here, with | 42 | | - ``IPV6_V6ONLY`` explicitly cleared, and handed to uvicorn pre-built. | 43 | | - """ | 44 | | - if _ipv6_available(): | 45 | | - sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM) | 46 | | - sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) | 47 | | - address: tuple = ("::", port) | 48 | | - else: | 49 | | - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | 50 | | - address = ("0.0.0.0", port) | 51 | | - sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | 52 | | - sock.bind(address) | 53 | | - sock.listen(socket.SOMAXCONN) | 54 | | - sock.setblocking(False) | 55 | | - return sock | 56 | | - | 57 | | - | 58 | 32 | class RestRegistryServer: | 59 | 33 | def __init__(self, store: FeatureStore): | 60 | 34 | self.store = store | @@ -357,7 +331,7 @@ def start_server( | 357 | 331 |
| 358 | 332 | _sync_protected_project_tag(self.store) | 359 | 333 |
| 360 | | - sock = _make_rest_socket(port) | | 334 | + sock = _make_dual_stack_socket(port) | 361 | 335 | if tls_key_path and tls_cert_path: | 362 | 336 | logger.info("Starting REST registry server in TLS(SSL) mode") | 363 | 337 | logger.info(f"REST registry server listening on https://localhost:{port}") | |
Collapse file Copy file name to clipboardExpand all lines: sdk/python/tests/unit/api/test_api_rest_registry_server.py +5-99Lines changed: 5 additions & 99 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 1 | | -import http.client | 2 | | -import socket | 3 | | -import threading | 4 | | -import time | 5 | | -from unittest.mock import MagicMock, patch | | 1 | +from unittest.mock import MagicMock | 6 | 2 |
| 7 | 3 | import pytest | 8 | 4 |
| 9 | | -from feast.api.registry.rest.rest_registry_server import ( | 10 | | - RestRegistryServer, | 11 | | - _make_rest_socket, | 12 | | -) | | 5 | +from feast.api.registry.rest.rest_registry_server import RestRegistryServer | 13 | 6 | from feast.feature_store import FeatureStore | 14 | 7 |
| 15 | 8 |
| @@ -85,93 +78,6 @@ def test_routes_registered_in_app(): | 85 | 78 | assert app.include_router.call_count == 16 | 86 | 79 |
| 87 | 80 |
| 88 | | -def _can_bind_ipv6_loopback() -> bool: | 89 | | - """Independent of `_ipv6_available` so a mutant in that function can't | 90 | | - also disable this skip guard and hide a killed mutant as a skip.""" | 91 | | - try: | 92 | | - with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as sock: | 93 | | - sock.bind(("::1", 0)) | 94 | | - return True | 95 | | - except OSError: | 96 | | - return False | 97 | | - | 98 | | - | 99 | | -def test_make_rest_socket_is_dual_stack_when_ipv6_available(): | 100 | | - mock_sock = MagicMock() | 101 | | - mock_sock.getsockname.return_value = ("::", 0) | 102 | | - with patch( | 103 | | - "feast.api.registry.rest.rest_registry_server._ipv6_available", | 104 | | - return_value=True, | 105 | | - ): | 106 | | - with patch("socket.socket", return_value=mock_sock) as mock_socket_cls: | 107 | | - result = _make_rest_socket(6572) | 108 | | - | 109 | | - mock_socket_cls.assert_called_once_with(socket.AF_INET6, socket.SOCK_STREAM) | 110 | | - mock_sock.setsockopt.assert_any_call(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) | 111 | | - mock_sock.setsockopt.assert_any_call(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | 112 | | - mock_sock.bind.assert_called_once_with(("::", 6572)) | 113 | | - mock_sock.listen.assert_called_once_with(socket.SOMAXCONN) | 114 | | - mock_sock.setblocking.assert_called_once_with(False) | 115 | | - assert result is mock_sock | 116 | | - | 117 | | - | 118 | | -def test_make_rest_socket_falls_back_to_ipv4(): | 119 | | - mock_sock = MagicMock() | 120 | | - with patch( | 121 | | - "feast.api.registry.rest.rest_registry_server._ipv6_available", | 122 | | - return_value=False, | 123 | | - ): | 124 | | - with patch("socket.socket", return_value=mock_sock) as mock_socket_cls: | 125 | | - _make_rest_socket(6572) | 126 | | - | 127 | | - mock_socket_cls.assert_called_once_with(socket.AF_INET, socket.SOCK_STREAM) | 128 | | - mock_sock.bind.assert_called_once_with(("0.0.0.0", 6572)) | 129 | | - | 130 | | - | 131 | | -def test_make_rest_socket_dual_stack_is_reachable_on_both_families(): | 132 | | - # Regression test: uvicorn's own host="::" binds IPv6-only, because | 133 | | - # asyncio's loop.create_server sets IPV6_V6ONLY=1 on any socket it | 134 | | - # creates itself from a host string. _make_rest_socket must hand | 135 | | - # uvicorn a pre-bound, dual-stack socket instead. | 136 | | - if not _can_bind_ipv6_loopback(): | 137 | | - pytest.skip("no IPv6 loopback on this host") | 138 | | - | 139 | | - import uvicorn | 140 | | - from fastapi import FastAPI | 141 | | - | 142 | | - sock = _make_rest_socket(0) | 143 | | - port = sock.getsockname()[1] | 144 | | - | 145 | | - app = FastAPI() | 146 | | - | 147 | | - @app.get("/health") | 148 | | - def health(): | 149 | | - return {"ok": True} | 150 | | - | 151 | | - server = uvicorn.Server(uvicorn.Config(app, log_level="critical")) | 152 | | - thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]}) | 153 | | - thread.daemon = True | 154 | | - thread.start() | 155 | | - try: | 156 | | - for _ in range(200): | 157 | | - if server.started: | 158 | | - break | 159 | | - time.sleep(0.01) | 160 | | - assert server.started, "uvicorn server did not start in time" | 161 | | - | 162 | | - for host in ("127.0.0.1", "::1"): | 163 | | - conn = http.client.HTTPConnection(host, port, timeout=5) | 164 | | - try: | 165 | | - conn.request("GET", "/health") | 166 | | - resp = conn.getresponse() | 167 | | - assert resp.status == 200 | 168 | | - finally: | 169 | | - conn.close() | 170 | | - finally: | 171 | | - server.should_exit = True | 172 | | - thread.join(timeout=5) | 173 | | - | 174 | | - | 175 | 81 | @pytest.mark.parametrize("tls", [True, False]) | 176 | 82 | def test_start_server_wires_uvicorn_server_with_prebound_socket( | 177 | 83 | mock_store_and_registry, mocker, tls | @@ -184,8 +90,8 @@ def test_start_server_wires_uvicorn_server_with_prebound_socket( | 184 | 90 | mocker.patch("feast.registry_server._sync_protected_project_tag") | 185 | 91 |
| 186 | 92 | mock_sock = MagicMock() | 187 | | - mock_make_rest_socket = mocker.patch( | 188 | | - "feast.api.registry.rest.rest_registry_server._make_rest_socket", | | 93 | + mock_make_dual_stack_socket = mocker.patch( | | 94 | + "feast.api.registry.rest.rest_registry_server._make_dual_stack_socket", | 189 | 95 | return_value=mock_sock, | 190 | 96 | ) | 191 | 97 | mock_config_cls = mocker.patch("uvicorn.Config") | @@ -203,7 +109,7 @@ def test_start_server_wires_uvicorn_server_with_prebound_socket( | 203 | 109 | else: | 204 | 110 | server.start_server(port=6572) | 205 | 111 |
| 206 | | - mock_make_rest_socket.assert_called_once_with(6572) | | 112 | + mock_make_dual_stack_socket.assert_called_once_with(6572) | 207 | 113 | mock_server_cls.assert_called_once_with(mock_config_cls.return_value) | 208 | 114 | mock_server_instance.run.assert_called_once_with(sockets=[mock_sock]) | 209 | 115 |
| |
Collapse file Copy file name to clipboardExpand all lines: sdk/python/tests/unit/test_utils.py +62Lines changed: 62 additions & 0 deletions | Original file line number | Diff line number | Diff line change |
|---|
| 6 | 6 | populates the GetOnlineFeaturesResponse. | 7 | 7 | """ | 8 | 8 |
| | 9 | +import http.client | 9 | 10 | import socket | | 11 | +import threading | | 12 | +import time | 10 | 13 | from datetime import datetime, timezone | 11 | 14 | from unittest.mock import MagicMock, patch | 12 | 15 |
| | 16 | +import pytest | | 17 | + | 13 | 18 | from feast.protos.feast.serving.ServingService_pb2 import ( | 14 | 19 | FieldStatus, | 15 | 20 | GetOnlineFeaturesResponse, | | 18 | 23 | from feast.utils import _make_dual_stack_socket, _populate_response_from_feature_data | 19 | 24 |
| 20 | 25 |
| | 26 | +def _can_bind_ipv6_loopback() -> bool: | | 27 | + """Independent of `_ipv6_available` so a mutant in that function can't | | 28 | + also disable this skip guard and hide a killed mutant as a skip.""" | | 29 | + try: | | 30 | + with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as sock: | | 31 | + sock.bind(("::1", 0)) | | 32 | + return True | | 33 | + except OSError: | | 34 | + return False | | 35 | + | | 36 | + | 21 | 37 | def _make_table(name="test_fv"): | 22 | 38 | """Create a minimal mock FeatureView for testing.""" | 23 | 39 | table = MagicMock() | @@ -466,3 +482,49 @@ def test_make_dual_stack_socket_falls_back_to_ipv4(): | 466 | 482 |
| 467 | 483 | mock_socket_cls.assert_called_once_with(socket.AF_INET, socket.SOCK_STREAM) | 468 | 484 | mock_sock.bind.assert_called_once_with(("0.0.0.0", 6580)) | | 485 | + | | 486 | + | | 487 | +def test_make_dual_stack_socket_is_reachable_on_both_families(): | | 488 | + # Regression test: a framework's own host="::" (uvicorn.run, | | 489 | + # asyncio.loop.create_server) binds IPv6-only, because IPV6_V6ONLY=1 | | 490 | + # gets set on any socket it creates itself from a host string. | | 491 | + # _make_dual_stack_socket must hand back a pre-bound, dual-stack socket | | 492 | + # instead -- this exercises the one real socket every dual-stack | | 493 | + # consumer (REST registry, ui, lineage) shares. | | 494 | + if not _can_bind_ipv6_loopback(): | | 495 | + pytest.skip("no IPv6 loopback on this host") | | 496 | + | | 497 | + import uvicorn | | 498 | + from fastapi import FastAPI | | 499 | + | | 500 | + sock = _make_dual_stack_socket(0) | | 501 | + port = sock.getsockname()[1] | | 502 | + | | 503 | + app = FastAPI() | | 504 | + | | 505 | + @app.get("/health") | | 506 | + def health(): | | 507 | + return {"ok": True} | | 508 | + | | 509 | + server = uvicorn.Server(uvicorn.Config(app, log_level="critical")) | | 510 | + thread = threading.Thread(target=server.run, kwargs={"sockets": [sock]}) | | 511 | + thread.daemon = True | | 512 | + thread.start() | | 513 | + try: | | 514 | + for _ in range(200): | | 515 | + if server.started: | | 516 | + break | | 517 | + time.sleep(0.01) | | 518 | + assert server.started, "uvicorn server did not start in time" | | 519 | + | | 520 | + for host in ("127.0.0.1", "::1"): | | 521 | + conn = http.client.HTTPConnection(host, port, timeout=5) | | 522 | + try: | | 523 | + conn.request("GET", "/health") | | 524 | + resp = conn.getresponse() | | 525 | + assert resp.status == 200 | | 526 | + finally: | | 527 | + conn.close() | | 528 | + finally: | | 529 | + server.should_exit = True | | 530 | + thread.join(timeout=5) |
|