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..72b13d5b52f48 --- /dev/null +++ b/python/pyspark/sql/connect/local_server_pool.py @@ -0,0 +1,231 @@ +# +# 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. +""" + +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: + + 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 + 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 + + # 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. + 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}") + + # 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, entries with a malformed uid, ...).""" + if name.startswith("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 + "-"): + 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.""" + 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. + + 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: + # 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 + + 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..7e8d2326fbf8f --- /dev/null +++ b/python/pyspark/sql/tests/connect/test_connect_local_server_pool.py @@ -0,0 +1,229 @@ +# +# 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.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",) + + +# 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, + 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): + """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) + + 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 + + main()