[SPARK-58021][CONNECT] Add local server pool member claiming

python/pyspark/sql/connect/local_server.py

@@ -62,19 +62,61 @@
 # A fixed SPARK_IDENT_STRING keeps the spark-daemon.sh pid and log file names stable
 # regardless of $USER.
 _SPARK_IDENT = "local-connect"
+_LINUX_ZOMBIE_STATE = "Z"
 
 
 def _pid_alive(pid: int) -> bool:
-    """Whether ``pid`` exists (POSIX only). A process we cannot signal counts as alive."""
+    """Whether ``pid`` is running. A process we cannot signal counts as alive. Linux zombies
+    count as terminated: they remain signalable until their parent reaps them, but cannot own
+    or serve a managed server.
+
+    Off POSIX this returns ``True`` without probing: ``os.kill`` there terminates the target for
+    any signal other than ``CTRL_C_EVENT`` / ``CTRL_BREAK_EVENT``, so signal 0 is not a safe
+    liveness probe, and callers fall through to the port check instead. Guarding here rather than
+    at each call site keeps every caller (reuse and pool) safe. (The pool path needs ``fcntl`` and
+    so never runs off POSIX regardless.)
+    """
+    if os.name != "posix":
+        return True
+    if pid <= 0:
+        return False
     try:
         os.kill(pid, 0)
     except ProcessLookupError:
         return False
+    except OverflowError:
+        return False
     except OSError:
         pass
+    if sys.platform.startswith("linux"):
+        try:
+            with open(f"/proc/{pid}/status", encoding="utf-8") as status_file:
+                for line in status_file:
+                    key, separator, value = line.partition(":")
+                    if separator and key == "State":
+                        state, _, _ = value.strip().partition(" ")
+                        if state == _LINUX_ZOMBIE_STATE:
+                            return False
+                        break
+        except FileNotFoundError:
+            return False
+        except OSError:
+            pass
     return True
 
 
+def _port_open(host: str, port: int, timeout: float = 0.5) -> bool:
+    """Whether a TCP connection to ``host``:``port`` succeeds within ``timeout`` seconds. A
+    socket error or a host that fails to resolve counts as closed.
+    """
+    try:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+            sock.settimeout(timeout)
+            return sock.connect_ex((host, port)) == 0
+    except (OSError, UnicodeError):
+        return False
+
+
 def _is_local_connect_server(pid: int) -> Optional[bool]:
     """Whether ``pid`` is still the managed Connect server recorded in discovery.
 
@@ -223,16 +265,14 @@ def url(self) -> str:
     def is_listening(self) -> bool:
         if self.host is None or self.port is None:
             return False
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-            sock.settimeout(0.5)
-            return sock.connect_ex((self.host, self.port)) == 0
+        return _port_open(self.host, self.port)
 
     def is_reusable(self) -> bool:
         from pyspark.version import __version__
 
         if self.spark_version != __version__ or self.pid is None:
             return False
-        if os.name == "posix" and not _pid_alive(self.pid):
+        if not _pid_alive(self.pid):
             return False
         return self.is_listening()
 
@@ -488,10 +528,7 @@ def _await_ready(self, port: int, token: str) -> None:
         while time.time() < deadline:
             pid = self._discovery.daemon_pid()
             if pid is not None:
-                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-                    sock.settimeout(0.5)
-                    listening = sock.connect_ex(("localhost", port)) == 0
-                if listening:
+                if _port_open("localhost", port):
                     self._discovery.save(
                         {
                             "host": "localhost",

python/pyspark/sql/connect/local_server_pool.py

@@ -15,17 +15,155 @@
 # limitations under the License.
 #
 
-"""Filesystem-backed state storage for local Spark Connect server pools.
+"""Filesystem-backed state model for local Spark Connect server pools.
 
-This internal foundation owns the pool directory layout, locking, and JSON state-file access.
+This internal foundation owns member identity, directory locking, state-file access, and
+claiming. Process lifecycle and server acquisition are layered on top in follow-up changes.
 """
 
 import contextlib
+import hashlib
 import json
+import math
 import os
 import shutil
+import sys
 from typing import Any, Dict, List, Optional, Tuple
 
+from pyspark.errors import PySparkValueError
+
+
+# Environment variables that shape the JVM the launcher boots through
+# sbin/start-connect-server.sh -> spark-daemon.sh -> load-spark-env.sh / spark-submit.
+# SPARK_CONF_DIR selects the spark-env.sh and spark-defaults.conf that seed the server; the
+# rest feed the classpath, heap, and JVM options. This is a curated set, not an exhaustive
+# one: the launcher inherits the whole environment, so it names the inputs that most commonly
+# differ between runs rather than every variable a server could read.
+#
+# PATH is a deliberate omission: bin/spark-class prefers ${JAVA_HOME}/bin/java and only falls
+# back to the first java on PATH, so two runs with different JDKs first on PATH and no JAVA_HOME
+# would share a member. The identity already tracks PATH indirectly through shutil.which for the
+# interpreters, and PATH is too volatile to fingerprint whole; a run needing a specific JDK
+# should set JAVA_HOME.
+_JVM_ENV_VARS = (
+    "SPARK_CONF_DIR",
+    "JAVA_HOME",
+    "SPARK_DIST_CLASSPATH",
+    "SPARK_DAEMON_MEMORY",
+    "SPARK_DRIVER_MEMORY",
+    "SPARK_SUBMIT_OPTS",
+    "SPARK_DAEMON_JAVA_OPTS",
+)
+
+
+def pool_fingerprint(master: str, seed_conf: Dict[str, Any]) -> str:
+    """The identity of a pool member: a curated set of inputs that shape the server a run would
+    have booted for itself. A run only claims members whose fingerprint equals its own, so a
+    pre-booted JVM is never handed to a run it would not have produced. The set is curated
+    rather than complete because the launcher inherits the full environment (see
+    ``_JVM_ENV_VARS``) -- it covers the inputs that most commonly differ between runs.
+
+    Besides the master and the seeded confs, this covers the working directory (unset warehouse
+    and Derby metastore locations resolve relative to it), the PySpark installation, the Python
+    interpreters the server would run UDFs and Python data sources with, and the environment
+    variables that shape the launched JVM.
+    """
+
+    def resolved(command: str) -> str:
+        # Relative commands resolve through PATH server-side; fold that in so equal command
+        # strings cannot stand for different interpreters on different PATHs.
+        return shutil.which(command) or command
+
+    # Two server code paths resolve the Python interpreter with opposite precedence, so a run
+    # changing only one of these variables would still have booted a different server. Include
+    # both resolutions: SparkConnectPlanner.pythonExec (Connect Python UDFs) prefers
+    # PYSPARK_PYTHON, while PythonUtils.defaultPythonExec (Python data sources) prefers
+    # PYSPARK_DRIVER_PYTHON. Both fall back to python3 and treat an empty value as set, matching
+    # the Scala sys.env.getOrElse chains.
+    udf_python = resolved(
+        os.environ.get("PYSPARK_PYTHON", os.environ.get("PYSPARK_DRIVER_PYTHON", "python3"))
+    )
+    data_source_python = resolved(
+        os.environ.get("PYSPARK_DRIVER_PYTHON", os.environ.get("PYSPARK_PYTHON", "python3"))
+    )
+    spark_home = os.environ.get("SPARK_HOME")
+    identity = [
+        master,
+        sorted((str(k), str(v)) for k, v in seed_conf.items()),
+        os.getcwd(),
+        sys.executable,
+        udf_python,
+        data_source_python,
+        os.path.realpath(__file__),
+        os.path.realpath(spark_home) if spark_home else "",
+        os.environ.get("PYTHONPATH", ""),
+        [os.environ.get(var, "") for var in _JVM_ENV_VARS],
+    ]
+    return hashlib.sha256(json.dumps(identity).encode("utf-8")).hexdigest()
+
+
+# The end of year 9999 UTC, as a Unix timestamp. ``created`` is a wall-clock ``time.time()``
+# reading, so no real clock reaches this for millennia; rejecting values beyond it keeps a
+# corrupt far-future timestamp from looking perpetually fresh to age-based reaping in the
+# layers above, which measure a member's age as ``time.time() - created``.
+_MAX_CREATED = 253402300799
+
+
+class PoolMember:
+    """One published pool server, wrapping its ``server-<uid>.json`` record."""
+
+    def __init__(self, data: Dict[str, Any]):
+        record = dict(data)
+        for key in ("host", "token", "spark_version", "fingerprint"):
+            if not isinstance(record[key], str) or not record[key]:
+                raise PySparkValueError(f"{key} must be a nonempty string")
+        for key in ("port", "pid"):
+            value = record[key]
+            if isinstance(value, bool) or not isinstance(value, int):
+                raise PySparkValueError(f"{key} must be an integer")
+        created = record["created"]
+        if isinstance(created, bool) or not isinstance(created, (int, float)):
+            raise PySparkValueError("created must be a number")
+        created = float(created)
+        if not 1 <= record["port"] <= 65535:
+            raise PySparkValueError("port is out of range")
+        if record["pid"] <= 0:
+            raise PySparkValueError("pid must be positive")
+        if not math.isfinite(created) or not 0 <= created <= _MAX_CREATED:
+            raise PySparkValueError(f"created must be a finite timestamp in [0, {_MAX_CREATED}]")
+        self.host: str = record["host"]
+        self.port: int = record["port"]
+        self.token: str = record["token"]
+        self.pid: int = record["pid"]
+        self.spark_version: str = record["spark_version"]
+        self.fingerprint: str = record["fingerprint"]
+        self.created: float = created
+        # Set when this process claims the member; the path of its claimed-<pid>-<uid>.json.
+        self.claim_path: Optional[str] = None
+
+    @classmethod
+    def from_data(cls, data: Dict[str, Any]) -> Optional["PoolMember"]:
+        """Parse a published member record, returning ``None`` when it is malformed."""
+        try:
+            return cls(data)
+        except (KeyError, TypeError, ValueError, OverflowError):
+            return None
+
+    @property
+    def url(self) -> str:
+        return f"sc://{self.host}:{self.port}"
+
+    def is_usable(self) -> bool:
+        """Whether this member has a matching Spark version, live process, and open port. Uses
+        the same liveness and reachability probes as the reuse path (see ``local_server``), so
+        the pool and reuse discovery agree on when a recorded server is still good."""
+        from pyspark.version import __version__
+        from pyspark.sql.connect.local_server import _pid_alive, _port_open
+
+        if self.spark_version != __version__ or not _pid_alive(self.pid):
+            return False
+        return _port_open(self.host, self.port)
+
 
 class PoolDirectory:
     """Path layout, file access, and the cross-process lock of one pool directory.
@@ -229,3 +367,43 @@ def remove(self, path: str) -> None:
     def remove_member_dir(self, uid: str) -> None:
         self._assert_locked()
         shutil.rmtree(self.member_dir(uid), ignore_errors=True)
+
+
+class ServerPool:
+    """Claims members from one pool directory; lifecycle operations are added later."""
+
+    def __init__(self, directory: Optional[PoolDirectory] = None):
+        self._directory = directory or PoolDirectory()
+
+    def claim(self, fingerprint: str) -> Optional[PoolMember]:
+        """Claim the oldest usable member with this fingerprint, or ``None``. The rename to
+        ``claimed-<pid>-<uid>.json`` marks the member as owned by this process; the reaping
+        rules use that pid to retire members whose client died without releasing them. The
+        caller must hold the directory lock so selection and rename form one transition.
+
+        Ordering is by ``created``, a wall-clock ``time.time()`` reading. It is comparable
+        across the independent processes that publish members, which ``time.monotonic()`` is
+        not, at the cost that a backward clock step (NTP, suspend/resume) can perturb the order.
+        Ties break by the stable ``sorted()`` over the sorted directory listing, so the order is
+        well defined but only approximately FIFO, not guaranteed.
+
+        ``is_usable`` runs under the held lock and does blocking network I/O -- up to a 0.5s
+        connect for each candidate that is live but not accepting connections. The candidate
+        count is bounded by ``spark.local.connect.pool.size``, which is user-tunable, so a large
+        pool widens the window the lock is held; the reaping rules keep stale members from
+        accumulating without bound."""
+        candidates = []
+        for uid, path in self._directory.paths_of_kind("server"):
+            data = self._directory.read_json(path)
+            member = PoolMember.from_data(data) if data is not None else None
+            if member is not None and member.fingerprint == fingerprint:
+                candidates.append((member, uid, path))
+        candidates.sort(key=lambda c: c[0].created)
+        for member, uid, path in candidates:
+            if not member.is_usable():
+                continue  # left for the reaping rules to retire
+            claim_path = self._directory.claimed_path(os.getpid(), uid)
+            self._directory.rename(path, claim_path)
+            member.claim_path = claim_path
+            return member
+        return None