perf: share repeated server-list strings

proton/vpn/core/cache_handler.py

@@ -23,6 +23,7 @@
 import json
 import os
 from pathlib import Path
+from typing import Callable, Optional
 from proton.vpn import logging
 
 
@@ -31,8 +32,15 @@
 
 class CacheHandler:
     """Used to save, load, and remove cache files."""
-    def __init__(self, filepath: str):
+    def __init__(
+            self,
+            filepath: str,
+            object_hook_factory: Optional[
+                Callable[[], Callable[[dict], dict]]
+            ] = None
+    ):
         self._fp = Path(filepath)
+        self._object_hook_factory = object_hook_factory
 
     @property
     def exists(self):
@@ -55,7 +63,11 @@ def load(self):
 
         try:
             with open(self._fp, "r", encoding="utf-8") as f:  # pylint: disable=C0103
-                return json.load(f)  # pylint: disable=C0103
+                object_hook = self._object_hook_factory() \
+                    if self._object_hook_factory else None
+                return json.load(  # pylint: disable=C0103
+                    f, object_hook=object_hook
+                )
         except (json.decoder.JSONDecodeError, UnicodeDecodeError):
             filename = os.path.basename(self._fp)
             logger.warning(

proton/vpn/session/servers/logicals.py

@@ -35,6 +35,43 @@
 
 UNIX_EPOCH = "Thu, 01 Jan 1970 00:00:00 GMT"
 
+_SHARED_SERVER_STRING_FIELDS = frozenset({
+    "EntryCountry", "ExitCountry", "HostCountry", "City", "State", "Domain",
+    "EntryIP", "ExitIP", "Generation", "Label", "ServicesDownReason",
+    "X25519PublicKey",
+})
+
+
+def _server_string_object_hook() -> Callable[[dict], dict]:
+    """Creates a JSON object hook that shares immutable server strings."""
+    shared_strings = {}
+
+    def share_strings(item: dict) -> dict:
+        for field, value in item.items():
+            if (
+                    field in _SHARED_SERVER_STRING_FIELDS
+                    and isinstance(value, str)
+            ):
+                item[field] = shared_strings.setdefault(value, value)
+        return item
+
+    return share_strings
+
+
+def _deduplicate_server_strings(logicals: List[dict]) -> None:
+    """Shares common immutable strings within a decoded server list.
+
+    This remains necessary for freshly downloaded responses, which have
+    already been decoded by the HTTP client. Cached lists use the same helper
+    as a JSON object hook so duplicate strings need not coexist in memory.
+    """
+    share_strings = _server_string_object_hook()
+
+    for logical in logicals:
+        share_strings(logical)
+        for physical in logical.get("Servers", ()):
+            share_strings(physical)
+
 
 class PersistenceKeys(Enum):
     """JSON Keys used to persist the ServerList to disk."""
@@ -426,7 +463,9 @@ def from_dict(
         """
         try:
             user_tier = data[PersistenceKeys.USER_TIER.value]
-            logicals = [LogicalServer(logical_dict) for logical_dict in data["LogicalServers"]]
+            logical_dicts = data[PersistenceKeys.LOGICALS.value]
+            _deduplicate_server_strings(logical_dicts)
+            logicals = [LogicalServer(logical_dict) for logical_dict in logical_dicts]
         except KeyError as error:
             raise ServerListDecodeError("Error building server list from dict") from error
 

proton/vpn/session/servers/server_list_fetcher.py

@@ -26,7 +26,9 @@
 from proton.vpn.core.cache_handler import CacheHandler
 from proton.vpn.session.exceptions import ServerListDecodeError
 from proton.vpn.session.servers.types import ServerLoad
-from proton.vpn.session.servers.logicals import ServerList, PersistenceKeys
+from proton.vpn.session.servers.logicals import (
+    ServerList, PersistenceKeys, _server_string_object_hook
+)
 from proton.vpn.session.dataclasses import VPNLocation
 from proton.vpn.session.utils import rest_api_request
 from proton.vpn.platform.core import ServerStatus  # pylint: disable=E0401, E0611
@@ -173,7 +175,10 @@ def __init__(
     ):
         self._session = session
         self._server_list = server_list
-        self._cache_file = cache_file or CacheHandler(self.CACHE_PATH)
+        self._cache_file = cache_file or CacheHandler(
+            self.CACHE_PATH,
+            object_hook_factory=_server_string_object_hook,
+        )
 
     def clear_cache(self):
         """Discards the cache, if existing."""

tests/python/core/test_cachehandler.py

@@ -52,6 +52,49 @@ def test_load_stored_cache(self, cache_filepath):
         assert "load_cache" in data
         assert "dummy-data" == data["load_cache"]
 
+    def test_load_stored_cache_with_object_hook(self, cache_filepath):
+        def object_hook_factory():
+            def object_hook(item):
+                item["decoded"] = True
+                return item
+
+            return object_hook
+
+        cache_handler = CacheHandler(
+            cache_filepath,
+            object_hook_factory=object_hook_factory,
+        )
+        with open(cache_filepath, "w") as f:
+            json.dump({"load_cache": "dummy-data"}, f)
+
+        assert cache_handler.load() == {
+            "load_cache": "dummy-data",
+            "decoded": True,
+        }
+
+    def test_object_hook_factory_is_scoped_to_each_load(self, cache_filepath):
+        load_count = 0
+
+        def object_hook_factory():
+            nonlocal load_count
+            load_count += 1
+
+            def object_hook(item):
+                item["load"] = load_count
+                return item
+
+            return object_hook
+
+        cache_handler = CacheHandler(
+            cache_filepath,
+            object_hook_factory=object_hook_factory,
+        )
+        with open(cache_filepath, "w") as f:
+            json.dump({"load_cache": "dummy-data"}, f)
+
+        assert cache_handler.load()["load"] == 1
+        assert cache_handler.load()["load"] == 2
+
     def test_load_cache_with_missing_file(self, cache_filepath):
         cache_handler = CacheHandler(cache_filepath)
         assert not cache_handler.load()

tests/python/session/servers/test_logicals.py

@@ -17,18 +17,21 @@
 along with ProtonVPN.  If not, see <https://www.gnu.org/licenses/>.
 """
 import functools
+import json
+from copy import deepcopy
 from typing import List
 
 import pytest
-from unittest.mock import Mock
 
 from proton.vpn.session.servers import LogicalServer, ServerFeatureEnum
 from proton.vpn.session.servers.logicals import (
+    _server_string_object_hook,
     sort_servers_alphabetically_by_country_and_server_name,
     sort_servers_by_country_and_location_and_enabled_and_load,
     ServerList
 )
 from proton.vpn.session.location_names_fetcher import LocationTranslations
+from proton.vpn.core.cache_handler import CacheHandler
 
 
 def _compact_features(features: List[ServerFeatureEnum]) -> ServerFeatureEnum:
@@ -130,6 +133,125 @@ def fixture_api_response() -> str:
     }
 
 
+def test_from_dict_deduplicates_common_server_strings():
+    country_1 = bytes((67, 72)).decode()
+    country_2 = bytes((67, 72)).decode()
+    city_1 = bytes((90, 117, 114, 105, 99, 104)).decode()
+    city_2 = bytes((90, 117, 114, 105, 99, 104)).decode()
+    label_1 = bytes((110, 111, 110, 101)).decode()
+    label_2 = bytes((110, 111, 110, 101)).decode()
+    domain_bytes = (110, 111, 100, 101, 46, 101, 120, 97, 109, 112, 108, 101)
+    domain_1 = bytes(domain_bytes).decode()
+    domain_2 = bytes(domain_bytes).decode()
+    entry_ip_1 = bytes((49, 57, 50, 46, 48, 46, 50, 46, 49)).decode()
+    entry_ip_2 = bytes((49, 57, 50, 46, 48, 46, 50, 46, 49)).decode()
+    assert country_1 is not country_2
+    assert city_1 is not city_2
+    assert label_1 is not label_2
+    assert domain_1 is not domain_2
+    assert entry_ip_1 is not entry_ip_2
+
+    payload = {
+        "MaxTier": 2,
+        "LogicalServers": [
+            {
+                "ID": "1", "Name": "CH#1", "ExitCountry": country_1,
+                "City": city_1, "Domain": domain_1,
+                "Servers": [{
+                    "Label": label_1, "Domain": domain_1,
+                    "EntryIP": entry_ip_1,
+                }],
+            },
+            {
+                "ID": "2", "Name": "CH#2", "ExitCountry": country_2,
+                "City": city_2, "Domain": domain_2,
+                "Servers": [{
+                    "Label": label_2, "Domain": domain_2,
+                    "EntryIP": entry_ip_2,
+                }],
+            },
+        ],
+    }
+    expected_values = deepcopy(payload)
+
+    server_list = ServerList.from_dict(payload)
+
+    assert server_list[0].exit_country is server_list[1].exit_country
+    assert server_list[0].city is server_list[1].city
+    assert (
+        server_list[0].physical_servers[0].label
+        is server_list[1].physical_servers[0].label
+    )
+    assert (
+        server_list[0].physical_servers[0].domain
+        is server_list[1].physical_servers[0].domain
+    )
+    assert (
+        server_list[0].physical_servers[0].entry_ip
+        is server_list[1].physical_servers[0].entry_ip
+    )
+    assert server_list[0].id is not server_list[1].id
+    assert payload == expected_values
+
+
+def test_json_object_hook_deduplicates_strings_during_decode():
+    data = json.loads(
+        '{"LogicalServers": ['
+        '{"ExitCountry": "CH", "Servers": [{"Domain": "node.example"}]},'
+        '{"ExitCountry": "CH", "Servers": [{"Domain": "node.example"}]}'
+        ']}',
+        object_hook=_server_string_object_hook(),
+    )
+
+    first, second = data["LogicalServers"]
+    assert first["ExitCountry"] is second["ExitCountry"]
+    assert first["Servers"][0]["Domain"] is second["Servers"][0]["Domain"]
+
+
+def test_string_sharing_preserves_missing_and_non_string_fields():
+    payload = {
+        "MaxTier": 2,
+        "LogicalServers": [{
+            "ID": "1",
+            "Name": "CH#1",
+            "ExitCountry": "CH",
+            "City": None,
+            "HostCountry": 42,
+            "Servers": [{"Label": False}],
+        }],
+    }
+
+    server_list = ServerList.from_dict(payload)
+
+    assert server_list[0].city is None
+    assert server_list[0].host_country == 42
+    assert server_list[0].physical_servers[0].label is False
+    assert "State" not in payload["LogicalServers"][0]
+
+
+def test_cached_string_pool_is_not_retained_between_loads(tmp_path):
+    shared_value = "shared-value-too-long-for-python-interning"
+    payload = {
+        "MaxTier": 2,
+        "LogicalServers": [
+            {"ID": "1", "Name": "CH#1", "ExitCountry": shared_value},
+            {"ID": "2", "Name": "CH#2", "ExitCountry": shared_value},
+        ],
+    }
+    cache = CacheHandler(
+        tmp_path / "server-list.json",
+        object_hook_factory=_server_string_object_hook,
+    )
+    cache.save(payload)
+
+    first = cache.load()["LogicalServers"]
+    second = cache.load()["LogicalServers"]
+
+    assert first[0]["ExitCountry"] is first[1]["ExitCountry"]
+    assert second[0]["ExitCountry"] is second[1]["ExitCountry"]
+    assert first[0]["ExitCountry"] is not second[0]["ExitCountry"]
+
+
 def test_set_location_translations_applies_to_every_logical(api_response: str):
     server_list = ServerList(
         user_tier=2,