From 23f806b64adc2d7e5750e45961ea2d4f10287026 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Mon, 10 Aug 2026 18:42:05 +0000 Subject: [PATCH 1/3] [SPARK-58021][CONNECT] Add local server pool filesystem storage --- dev/sparktestsupport/modules.py | 1 + .../pyspark/sql/connect/local_server_pool.py | 191 ++++++++++++++++++ .../connect/test_connect_local_server_pool.py | 118 +++++++++++ 3 files changed, 310 insertions(+) create mode 100644 python/pyspark/sql/connect/local_server_pool.py create mode 100644 python/pyspark/sql/tests/connect/test_connect_local_server_pool.py diff --git a/dev/sparktestsupport/modules.py b/dev/sparktestsupport/modules.py index 2faeea33b6f6b..fd4a452d9750b 100644 --- a/dev/sparktestsupport/modules.py +++ b/dev/sparktestsupport/modules.py @@ -1199,6 +1199,7 @@ def __hash__(self): "pyspark.sql.tests.connect.test_connect_retry", "pyspark.sql.tests.connect.test_connect_session", "pyspark.sql.tests.connect.test_connect_local_server", + "pyspark.sql.tests.connect.test_connect_local_server_pool", "pyspark.sql.tests.connect.test_connect_stat", "pyspark.sql.tests.connect.test_parity_geographytype", "pyspark.sql.tests.connect.test_parity_geometrytype", diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py new file mode 100644 index 0000000000000..381c2b31227d8 --- /dev/null +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -0,0 +1,191 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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. +# + +"""Filesystem-backed state storage for local Spark Connect server pools. + +This internal foundation owns the pool directory layout, locking, and JSON state-file access. +Member validation, claiming, and process lifecycle are layered on top in follow-up changes. +""" + +import contextlib +import json +import os +import shutil +from typing import Any, Dict, List, Optional, Tuple + + +class PoolDirectory: + """Path layout, file access, and the cross-process lock of one pool directory. + + Used as a context manager that holds the directory's exclusive lock:: + + with PoolDirectory() as directory: + ...read and write state files... + + Pool operations are infrequent, so one exclusive lock for every state transition is simpler + than a finer-grained scheme. A context can be entered again after exiting, allowing callers + to release the lock between polling attempts so other processes can update the directory. + """ + + def __init__(self, path: Optional[str] = None): + if path is None: + path = os.environ.get("SPARK_LOCAL_CONNECT_POOL_DIR") + if path is None: + from pyspark.sql.connect.local_server import runtime_dir + + path = os.path.join(runtime_dir(), "pool") + self.path = os.path.abspath(path) + self._lock_fd: Optional[int] = None + + def __enter__(self) -> "PoolDirectory": + import fcntl + + os.makedirs(self.path, mode=0o700, exist_ok=True) + # Re-assert privacy for an existing override directory: state files contain auth tokens, + # and directory write access would allow replacing them or bypassing the shared lock. + os.chmod(self.path, 0o700) + lock_fd = os.open(os.path.join(self.path, ".lock"), os.O_RDWR | os.O_CREAT, 0o600) + try: + os.fchmod(lock_fd, 0o600) + fcntl.flock(lock_fd, fcntl.LOCK_EX) + except BaseException: + os.close(lock_fd) + raise + self._lock_fd = lock_fd + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + assert self._lock_fd is not None + os.close(self._lock_fd) # closing releases the lock + self._lock_fd = None + + def _assert_locked(self) -> None: + assert self._lock_fd is not None, "PoolDirectory must be used as a context manager" + + # Path builders; these do not touch the filesystem and need no lock. + + def pending_path(self, uid: str) -> str: + return os.path.join(self.path, f"pending-{uid}.json") + + def conf_path(self, uid: str) -> str: + return os.path.join(self.path, f"conf-{uid}.json") + + def server_path(self, uid: str) -> str: + return os.path.join(self.path, f"server-{uid}.json") + + def claimed_path(self, client_pid: int, uid: str) -> str: + return os.path.join(self.path, f"claimed-{client_pid}-{uid}.json") + + def retired_path(self, uid: str) -> str: + return os.path.join(self.path, f"retired-{uid}.json") + + def member_dir(self, uid: str) -> str: + return os.path.join(self.path, f"member-{uid}") + + @staticmethod + def parse_entry(name: str) -> Tuple[Optional[str], Optional[str]]: + """The ``(kind, uid)`` of a pool directory entry, ``(None, None)`` for anything + else (the lock file, editor droppings, ...).""" + if name.startswith("member-"): + return "member", name[len("member-") :] + if not name.endswith(".json"): + return None, None + stem = name[: -len(".json")] + for kind in ("pending", "conf", "server", "retired"): + if stem.startswith(kind + "-"): + return kind, stem[len(kind) + 1 :] + if stem.startswith("claimed-"): + client_pid, sep, uid = stem[len("claimed-") :].partition("-") + if sep and client_pid.isdigit(): + return "claimed", uid + return None, None + + @staticmethod + def claiming_pid(claimed_path: str) -> int: + """The client pid recorded in a ``claimed--.json`` file name.""" + return int(os.path.basename(claimed_path)[len("claimed-") :].split("-", 1)[0]) + + # Locked accessors. + + def uids(self) -> List[str]: + self._assert_locked() + seen = [] + for name in self._entries(): + _, uid = self.parse_entry(name) + if uid is not None and uid not in seen: + seen.append(uid) + return seen + + def states(self, uid: str) -> Dict[str, str]: + """The state entries currently existing for ``uid``, as ``{kind: path}`` with kinds + ``pending``, ``conf``, ``server``, ``claimed``, ``retired``, and ``member`` (the + member's directory).""" + self._assert_locked() + found: Dict[str, str] = {} + for name in self._entries(): + kind, entry_uid = self.parse_entry(name) + if kind is not None and entry_uid == uid: + found[kind] = os.path.join(self.path, name) + return found + + def paths_of_kind(self, kind: str) -> List[Tuple[str, str]]: + """All ``(uid, path)`` of one state kind.""" + self._assert_locked() + return [ + (uid, os.path.join(self.path, name)) + for name in self._entries() + for entry_kind, uid in (self.parse_entry(name),) + if entry_kind == kind and uid is not None + ] + + def _entries(self) -> List[str]: + try: + return sorted(os.listdir(self.path)) + except OSError: + return [] + + def read_json(self, path: str) -> Optional[Dict[str, Any]]: + """``None`` for files that are missing or unreadable -- callers treat both like the + state not existing, and the reaping rules remove unreadable leftovers.""" + self._assert_locked() + try: + with open(path, "r") as f: + data = json.load(f) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None + + def write_json(self, path: str, data: Dict[str, Any]) -> None: + self._assert_locked() + # 0600 like the reuse discovery file: server entries hold the auth token. + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600) + with os.fdopen(fd, "w") as f: + os.fchmod(fd, 0o600) + f.write(json.dumps(data)) + + def rename(self, src: str, dst: str) -> None: + self._assert_locked() + os.rename(src, dst) + + def remove(self, path: str) -> None: + self._assert_locked() + with contextlib.suppress(FileNotFoundError): + os.remove(path) + + def remove_member_dir(self, uid: str) -> None: + self._assert_locked() + shutil.rmtree(self.member_dir(uid), ignore_errors=True) diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py new file mode 100644 index 0000000000000..da48bacf4e77c --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -0,0 +1,118 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# 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 os +import shutil +import subprocess +import sys +import tempfile +import unittest + +from pyspark.util import is_remote_only +from pyspark.testing.connectutils import should_test_connect, connect_requirement_message + +if should_test_connect: + from pyspark.sql.connect.local_server_pool import PoolDirectory + + +_SAVED_ENV_KEYS = ("SPARK_LOCAL_CONNECT_POOL_DIR",) + + +@unittest.skipIf( + not should_test_connect or is_remote_only(), + connect_requirement_message or "Requires Spark Connect test dependencies", +) +@unittest.skipUnless(os.name == "posix", "the pool relies on POSIX file locks") +class LocalConnectServerPoolUnitTests(unittest.TestCase): + """Tests for the pool filesystem storage; no real servers are started.""" + + def setUp(self) -> None: + self._tmpdir = tempfile.mkdtemp() + self._saved_env = {k: os.environ.get(k) for k in _SAVED_ENV_KEYS} + for k in _SAVED_ENV_KEYS: + os.environ.pop(k, None) + os.environ["SPARK_LOCAL_CONNECT_POOL_DIR"] = os.path.join(self._tmpdir, "pool") + self._directory = PoolDirectory() + + def tearDown(self) -> None: + for k, v in self._saved_env.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def test_pool_directory_location(self) -> None: + self.assertEqual(self._directory.path, os.path.join(self._tmpdir, "pool")) + os.environ.pop("SPARK_LOCAL_CONNECT_POOL_DIR") + default = PoolDirectory() + self.assertEqual(os.path.basename(default.path), "pool") + self.assertTrue(default.path.startswith(tempfile.gettempdir())) + + def test_pool_directory_lock_and_state_file_permissions(self) -> None: + os.makedirs(self._directory.path, mode=0o755) + os.chmod(self._directory.path, 0o755) + state_path = self._directory.server_path("secure") + with self._directory as directory: + self.assertEqual(os.stat(directory.path).st_mode & 0o777, 0o700) + lock_path = os.path.join(directory.path, ".lock") + self.assertEqual(os.stat(lock_path).st_mode & 0o777, 0o600) + directory.write_json(state_path, {"token": "secret"}) + self.assertEqual(os.stat(state_path).st_mode & 0o777, 0o600) + self.assertEqual(directory.read_json(state_path), {"token": "secret"}) + + # O_CREAT does not apply its mode to an existing file, so writes must re-assert it. + os.chmod(state_path, 0o644) + directory.write_json(state_path, {"token": "new-secret"}) + self.assertEqual(os.stat(state_path).st_mode & 0o777, 0o600) + + with open(state_path, "w") as state_file: + state_file.write("not json") + with self._directory as directory: + self.assertIsNone(directory.read_json(state_path)) + + def test_pool_directory_lock_blocks_another_process(self) -> None: + child = ( + "import errno\n" + "import fcntl\n" + "import os\n" + "import sys\n" + "fd = os.open(sys.argv[1], os.O_RDWR)\n" + "try:\n" + " fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)\n" + "except OSError as error:\n" + " if error.errno not in (errno.EACCES, errno.EAGAIN):\n" + " raise\n" + "else:\n" + " raise RuntimeError('acquired a held lock')\n" + "finally:\n" + " os.close(fd)\n" + ) + with self._directory: + result = subprocess.run( + [sys.executable, "-c", child, os.path.join(self._directory.path, ".lock")], + capture_output=True, + text=True, + timeout=10, + ) + self.assertEqual(result.returncode, 0, result.stderr) + + +if __name__ == "__main__": + from pyspark.testing import main + + main() From e85e9486543fb2f1c33676e66ff917c323a8a011 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Mon, 10 Aug 2026 23:11:40 +0000 Subject: [PATCH 2/3] [SPARK-58021][CONNECT][FOLLOWUP] Address review comments on pool storage Follow-up to the local server pool filesystem storage layer, addressing review feedback: - Guard PoolDirectory against re-entrant locking (a nested `with` on the same object would deadlock silently) with an explicit assertion. - Validate the uid shape (nonempty lowercase hex) in parse_entry for every entry kind, so editor droppings (e.g. member-.json.swp) and empty stems (server-.json) no longer surface as phantom uids. - Factor the claimed-- filename grammar into one _split_claimed helper shared by parse_entry and claiming_pid. The helper returns the pid unparsed so parse_entry, which classifies every directory entry, never raises; the isascii() guard keeps claiming_pid's int() total (str.isdigit() otherwise accepts characters int() rejects, e.g. superscripts). - Assert states() sees at most one entry per kind, so two claimed entries for one uid surface instead of silently collapsing to whichever sorts last. - Add table-driven tests for the filename grammar plus coverage for claiming_pid, the locked accessors, rename/remove/remove_member_dir, the lock-required guard, and the reentrancy guard. - Drop the over-restrictive is_remote_only() test gate (these tests start no server and use only stdlib) and correct the skip message. Co-authored-by: Isaac --- .../pyspark/sql/connect/local_server_pool.py | 61 ++++++--- .../connect/test_connect_local_server_pool.py | 117 +++++++++++++++++- 2 files changed, 161 insertions(+), 17 deletions(-) diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index 381c2b31227d8..cf6f85e66f252 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -54,6 +54,9 @@ def __init__(self, path: Optional[str] = None): def __enter__(self) -> "PoolDirectory": import fcntl + # Not reentrant: a nested enter would os.open a second fd and flock(LOCK_EX) would block + # forever against the fd this process already holds. Fail loudly instead of deadlocking. + assert self._lock_fd is None, "PoolDirectory is not reentrant" os.makedirs(self.path, mode=0o700, exist_ok=True) # Re-assert privacy for an existing override directory: state files contain auth tokens, # and directory write access would allow replacing them or bypassing the shared lock. @@ -96,28 +99,54 @@ def retired_path(self, uid: str) -> str: def member_dir(self, uid: str) -> str: return os.path.join(self.path, f"member-{uid}") - @staticmethod - def parse_entry(name: str) -> Tuple[Optional[str], Optional[str]]: + # uids are generated as ``uuid.uuid4().hex[:12]`` (see the acquisition layer), so a valid + # uid is a nonempty run of lowercase hex. Validating the shape keeps editor droppings such + # as ``member-abc.json.swp`` and empty stems like ``server-.json`` from becoming phantom uids. + _UID_CHARS = frozenset("0123456789abcdef") + + @classmethod + def _is_uid(cls, uid: str) -> bool: + return bool(uid) and all(c in cls._UID_CHARS for c in uid) + + @classmethod + def _split_claimed(cls, stem: str) -> Optional[Tuple[str, str]]: + """Split a well-formed ``claimed--`` stem (without the ``.json`` suffix) into + ``(client_pid, uid)`` as strings, or ``None`` otherwise. The pid is returned unparsed: + ``parse_entry`` classifies over every directory entry and must never raise, and + ``str.isdigit()`` accepts characters ``int()`` rejects (e.g. superscripts), so the + ``isascii()`` guard keeps the eventual ``int()`` in ``claiming_pid`` total.""" + if not stem.startswith("claimed-"): + return None + client_pid, sep, uid = stem[len("claimed-") :].partition("-") + if not sep or not (client_pid.isascii() and client_pid.isdigit()) or not cls._is_uid(uid): + return None + return client_pid, uid + + @classmethod + def parse_entry(cls, name: str) -> Tuple[Optional[str], Optional[str]]: """The ``(kind, uid)`` of a pool directory entry, ``(None, None)`` for anything - else (the lock file, editor droppings, ...).""" + else (the lock file, editor droppings, entries with a malformed uid, ...).""" if name.startswith("member-"): - return "member", name[len("member-") :] + uid = name[len("member-") :] + return ("member", uid) if cls._is_uid(uid) else (None, None) if not name.endswith(".json"): return None, None stem = name[: -len(".json")] for kind in ("pending", "conf", "server", "retired"): if stem.startswith(kind + "-"): - return kind, stem[len(kind) + 1 :] - if stem.startswith("claimed-"): - client_pid, sep, uid = stem[len("claimed-") :].partition("-") - if sep and client_pid.isdigit(): - return "claimed", uid - return None, None - - @staticmethod - def claiming_pid(claimed_path: str) -> int: + uid = stem[len(kind) + 1 :] + return (kind, uid) if cls._is_uid(uid) else (None, None) + claimed = cls._split_claimed(stem) + return ("claimed", claimed[1]) if claimed is not None else (None, None) + + @classmethod + def claiming_pid(cls, claimed_path: str) -> int: """The client pid recorded in a ``claimed--.json`` file name.""" - return int(os.path.basename(claimed_path)[len("claimed-") :].split("-", 1)[0]) + name = os.path.basename(claimed_path) + stem = name[: -len(".json")] if name.endswith(".json") else name + claimed = cls._split_claimed(stem) + assert claimed is not None, f"not a claimed entry: {claimed_path!r}" + return int(claimed[0]) # Locked accessors. @@ -139,6 +168,10 @@ def states(self, uid: str) -> Dict[str, str]: for name in self._entries(): kind, entry_uid = self.parse_entry(name) if kind is not None and entry_uid == uid: + # At most one entry per kind. Claiming renames a single file into place (see the + # claiming layer), so two claimed entries for one uid means the pid a reaper would + # read via claiming_pid is ambiguous; surface that rather than pick one silently. + assert kind not in found, f"duplicate {kind} entries for uid {uid}" found[kind] = os.path.join(self.path, name) return found diff --git a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py index da48bacf4e77c..7e8d2326fbf8f 100644 --- a/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -22,7 +22,6 @@ import tempfile import unittest -from pyspark.util import is_remote_only from pyspark.testing.connectutils import should_test_connect, connect_requirement_message if should_test_connect: @@ -32,9 +31,12 @@ _SAVED_ENV_KEYS = ("SPARK_LOCAL_CONNECT_POOL_DIR",) +# These tests start no server and exercise only stdlib filesystem code, so they do not need JVM +# access (no is_remote_only gate). should_test_connect is still required because importing +# PoolDirectory pulls in the pyspark.sql.connect package, which checks Connect dependencies. @unittest.skipIf( - not should_test_connect or is_remote_only(), - connect_requirement_message or "Requires Spark Connect test dependencies", + not should_test_connect, + connect_requirement_message or "Requires Spark Connect dependencies to import the pool module", ) @unittest.skipUnless(os.name == "posix", "the pool relies on POSIX file locks") class LocalConnectServerPoolUnitTests(unittest.TestCase): @@ -111,6 +113,115 @@ def test_pool_directory_lock_blocks_another_process(self) -> None: ) self.assertEqual(result.returncode, 0, result.stderr) + def test_parse_entry_grammar(self) -> None: + # uids match the acquisition layer's uuid4().hex[:12]: nonempty lowercase hex. + uid = "0123456789ab" + cases = [ + # Well-formed state entries of every kind. + (f"pending-{uid}.json", ("pending", uid)), + (f"conf-{uid}.json", ("conf", uid)), + (f"server-{uid}.json", ("server", uid)), + (f"retired-{uid}.json", ("retired", uid)), + (f"claimed-4321-{uid}.json", ("claimed", uid)), + (f"member-{uid}", ("member", uid)), + # Short but valid hex uids. + ("member-abc", ("member", "abc")), + ("server-abc.json", ("server", "abc")), + # The lock file and unrelated entries. + (".lock", (None, None)), + ("random.txt", (None, None)), + # Editor droppings: the finding-2 cases that used to slip through as phantom uids. + (f"member-{uid}.json.swp", (None, None)), + (f"server-{uid}.json.swp", (None, None)), + # Empty uids are rejected for every kind. + ("server-.json", (None, None)), + ("member-", (None, None)), + ("claimed-1234-.json", (None, None)), + # Non-hex uids (uppercase, out-of-range letters) are rejected. + ("server-ABCDEF.json", (None, None)), + ("server-ghij.json", (None, None)), + # Malformed claimed stems (non-numeric or missing pid, or a malformed uid). + (f"claimed-notapid-{uid}.json", (None, None)), + (f"claimed-{uid}.json", (None, None)), + ("claimed-4321-ABCDEF.json", (None, None)), + # A pid that is str.isdigit() but not int()-parsable (superscript two, U+00B2) must + # classify as "not claimed", never raise, since parse_entry runs over every entry. + (f"claimed-{chr(0xB2)}-{uid}.json", (None, None)), + ] + for name, expected in cases: + with self.subTest(name=name): + self.assertEqual(PoolDirectory.parse_entry(name), expected) + + def test_claiming_pid(self) -> None: + uid = "0123456789ab" + path = self._directory.claimed_path(4321, uid) + self.assertEqual(PoolDirectory.claiming_pid(path), 4321) + # Parses from the basename alone, independent of the directory prefix. + self.assertEqual(PoolDirectory.claiming_pid(f"claimed-7-{uid}.json"), 7) + + def test_locked_accessors_enumerate_state(self) -> None: + uid_a, uid_b = "aaaaaaaaaaaa", "bbbbbbbbbbbb" + with self._directory as directory: + directory.write_json(directory.server_path(uid_a), {"a": 1}) + directory.write_json(directory.pending_path(uid_a), {"a": 2}) + directory.write_json(directory.server_path(uid_b), {"b": 1}) + os.makedirs(directory.member_dir(uid_a), mode=0o700) + + self.assertEqual(sorted(directory.uids()), [uid_a, uid_b]) + + states_a = directory.states(uid_a) + self.assertEqual(set(states_a), {"server", "pending", "member"}) + self.assertEqual(states_a["server"], directory.server_path(uid_a)) + self.assertEqual(states_a["member"], directory.member_dir(uid_a)) + + servers = dict(directory.paths_of_kind("server")) + self.assertEqual(set(servers), {uid_a, uid_b}) + self.assertEqual(servers[uid_a], directory.server_path(uid_a)) + self.assertEqual(directory.paths_of_kind("retired"), []) + + def test_rename_remove_and_member_dir(self) -> None: + uid = "cccccccccccc" + with self._directory as directory: + src = directory.pending_path(uid) + dst = directory.server_path(uid) + directory.write_json(src, {"x": 1}) + directory.rename(src, dst) + self.assertFalse(os.path.exists(src)) + self.assertEqual(directory.read_json(dst), {"x": 1}) + + directory.remove(dst) + self.assertFalse(os.path.exists(dst)) + # Removing a missing path is a no-op. + directory.remove(dst) + + member = directory.member_dir(uid) + os.makedirs(member, mode=0o700) + with open(os.path.join(member, "inner"), "w") as f: + f.write("data") + directory.remove_member_dir(uid) + self.assertFalse(os.path.exists(member)) + # Removing a missing member directory is a no-op. + directory.remove_member_dir(uid) + + def test_states_rejects_duplicate_claimed(self) -> None: + uid = "dddddddddddd" + with self._directory as directory: + directory.write_json(directory.claimed_path(111, uid), {}) + directory.write_json(directory.claimed_path(222, uid), {}) + with self.assertRaisesRegex(AssertionError, "duplicate claimed"): + directory.states(uid) + + def test_accessors_require_the_lock(self) -> None: + # The locked accessors must refuse to run outside the context manager. + with self.assertRaisesRegex(AssertionError, "context manager"): + self._directory.uids() + + def test_not_reentrant(self) -> None: + with self._directory: + with self.assertRaisesRegex(AssertionError, "not reentrant"): + with self._directory: + pass + if __name__ == "__main__": from pyspark.testing import main From c0f73af1fa975f8ea9fd5d9fae4e1a3788c3cde9 Mon Sep 17 00:00:00 2001 From: ericm-db Date: Tue, 11 Aug 2026 23:22:44 +0000 Subject: [PATCH 3/3] [SPARK-58021][CONNECT][FOLLOWUP] Clarify pool storage documentation --- python/pyspark/sql/connect/local_server_pool.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/python/pyspark/sql/connect/local_server_pool.py b/python/pyspark/sql/connect/local_server_pool.py index cf6f85e66f252..72b13d5b52f48 100644 --- a/python/pyspark/sql/connect/local_server_pool.py +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -18,7 +18,6 @@ """Filesystem-backed state storage for local Spark Connect server pools. This internal foundation owns the pool directory layout, locking, and JSON state-file access. -Member validation, claiming, and process lifecycle are layered on top in follow-up changes. """ import contextlib @@ -31,10 +30,18 @@ class PoolDirectory: """Path layout, file access, and the cross-process lock of one pool directory. - Used as a context manager that holds the directory's exclusive lock:: + Used as a context manager that holds the directory's exclusive lock: - with PoolDirectory() as directory: - ...read and write state files... + directory = PoolDirectory() + with directory: + path = directory.pending_path(uid) + directory.write_json(path, data) + stored = directory.read_json(path) + directory.rename(path, directory.server_path(uid)) + + Entering the context creates the directory and acquires its lock. Callers then use path + builders and the locked accessors to enumerate, read, write, rename, or remove state. Exiting + the context releases the lock. Pool operations are infrequent, so one exclusive lock for every state transition is simpler than a finer-grained scheme. A context can be entered again after exiting, allowing callers