Skip to content

Commit ed95123

Browse files
ericm-dbdtenedor
authored andcommitted
[SPARK-58021][CONNECT] Reap orphaned local pool members
### What changes were proposed in this pull request? This is layer 5 of the ten-PR local Connect pool stack: #57684 -> #57685 -> #57907 -> #57686 -> #58247 -> #58367 -> #57687 -> #58248 -> #57102 -> #57688 This layer adds recovery after a pool member has been published or claimed: - a janitor for ready, claimed, and retiring members; - retirement of dead, unreachable, incompatible, idle, and malformed ready members; - retirement of claims whose client or server process generation has disappeared; - persisted claimant process identities so PID reuse cannot strand an orphaned claim; - recovery of independently valid server handles from malformed records; and - bounded, PID-reuse-safe completion of retirement. It also addresses the two follow-ups from the review of #57686: successful SIGTERM delivery is persisted so young retirement passes do not repeat the process inspection, and the daemon-pid fallback now has coverage for removing state after that process exits. Pending-launch, attendant, conf-file, and unreferenced member-directory recovery remain in #58367. ### Why are the changes needed? Pool members outlive individual Python call frames. A client can be killed without releasing its claim, a ready server can die or become unreachable, and persisted process IDs can be reused. Without a conservative janitor these cases can strand JVMs or make unusable members count toward the pool indefinitely. Separating post-publication recovery from launch recovery gives each process-ownership model its own review unit: this layer authorizes server cleanup from persisted process generations, while the next layer handles attendant commands and launch process groups. ### Does this PR introduce _any_ user-facing change? No. The pool is not wired into SparkSession in this layer. ### How was this patch tested? The focused pool suite's 50 tests passed: ```bash PYTHONPATH=python:python/lib/pyspark.zip:python/lib/py4j-0.10.9.9-src.zip \ SPARK_TESTING=1 \ .venv/bin/python -m unittest -v \ pyspark.sql.tests.connect.test_connect_local_server_pool ``` The changed files also passed Python AST parsing, Ruff checking and formatting, targeted mypy, `git diff --check`, and changed-file ASCII and 100-column checks. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Fable 5) and OpenAI Codex (GPT-5) Closes #58247 from ericm-db/local-connect-pool-recovery. Authored-by: Eric Marnadi <eric.marnadi@databricks.com> Signed-off-by: Daniel Tenedorio <daniel.tenedorio@databricks.com>
1 parent 7fa054b commit ed95123

2 files changed

Lines changed: 454 additions & 39 deletions

File tree

‎python/pyspark/sql/connect/local_server_pool.py‎

Lines changed: 180 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@ class RetiredState(_PoolStateRecord):
149149
pid: int
150150
process_start_id: str
151151
retired: float
152+
signalled: bool = False
152153

153154
@classmethod
154155
def pid_from_data(cls, data: Optional[Dict[str, Any]]) -> Optional[int]:
@@ -166,22 +167,32 @@ def retired_from_data(cls, data: Optional[Dict[str, Any]]) -> Optional[float]:
166167
"""Recover a valid retirement timestamp from a malformed state record."""
167168
return cls._timestamp(data.get("retired")) if data is not None else None
168169

170+
@staticmethod
171+
def signalled_from_data(data: Optional[Dict[str, Any]]) -> Optional[bool]:
172+
"""Recover SIGTERM delivery state, defaulting legacy records to unsignalled."""
173+
if data is None:
174+
return None
175+
value = data.get("signalled", False)
176+
return value if isinstance(value, bool) else None
177+
169178
@classmethod
170179
def from_data(cls, data: Optional[Dict[str, Any]]) -> Optional["RetiredState"]:
171180
if data is None:
172181
return None
173182
pid = cls.pid_from_data(data)
174183
process_start_id = cls.process_start_id_from_data(data)
175184
retired = cls.retired_from_data(data)
176-
if pid is None or process_start_id is None or retired is None:
185+
signalled = cls.signalled_from_data(data)
186+
if pid is None or process_start_id is None or retired is None or signalled is None:
177187
return None
178-
return cls(pid, process_start_id, retired)
188+
return cls(pid, process_start_id, retired, signalled)
179189

180190
def as_data(self) -> Dict[str, Any]:
181191
return {
182192
"pid": self.pid,
183193
"process_start_id": self.process_start_id,
184194
"retired": self.retired,
195+
"signalled": self.signalled,
185196
}
186197

187198

@@ -228,6 +239,23 @@ def from_data(cls, data: Dict[str, Any]) -> Optional["PoolMember"]:
228239
except (KeyError, TypeError, ValueError, OverflowError):
229240
return None
230241

242+
@classmethod
243+
def pid_from_data(cls, data: Optional[Dict[str, Any]]) -> Optional[int]:
244+
"""Recover a valid server pid even when another member field is malformed."""
245+
return cls._positive_pid(data.get("pid")) if data is not None else None
246+
247+
@staticmethod
248+
def process_start_id_from_data(data: Optional[Dict[str, Any]]) -> Optional[str]:
249+
"""Recover a process generation identifier from a malformed member record."""
250+
value = data.get("process_start_id") if data is not None else None
251+
return value if isinstance(value, str) and value else None
252+
253+
@staticmethod
254+
def client_process_start_id_from_data(data: Optional[Dict[str, Any]]) -> Optional[str]:
255+
"""Recover the claiming client's process generation from a claimed record."""
256+
value = data.get("client_process_start_id") if data is not None else None
257+
return value if isinstance(value, str) and value else None
258+
231259
@property
232260
def url(self) -> str:
233261
return f"sc://{self.host}:{self.port}"
@@ -486,13 +514,14 @@ def remove_member_dir(self, uid: str) -> None:
486514

487515

488516
class ServerPool:
489-
"""Claims and retires members of one pool directory."""
517+
"""Claims, reaps, and retires members of one pool directory."""
490518

491-
# A retired server still alive after the grace period is hard-killed. After the give-up
492-
# age, tracking is removed only once the process is gone, replaced, or successfully
493-
# signalled.
519+
# A retired server still alive after the grace period is hard-killed. With a process handle,
520+
# tracking is removed only once it is gone, replaced, or successfully signalled. PID-less
521+
# malformed state uses the give-up age as its bounded recovery window.
494522
_RETIRE_KILL_AFTER_SECONDS = 30
495523
_RETIRE_GIVE_UP_AFTER_SECONDS = 600
524+
_DEFAULT_IDLE_TIMEOUT_SECONDS = 1800
496525
_PROCESS_INSPECTION_TIMEOUT_SECONDS = 5
497526
_PROC_STAT_START_TIME_INDEX = 19
498527

@@ -543,13 +572,25 @@ def _process_start_id(cls, pid: int) -> Optional[str]:
543572
return f"ps:{started}" if result.returncode == 0 and started else None
544573

545574
@classmethod
546-
def _same_server_instance(cls, pid: int, process_start_id: str) -> Optional[bool]:
547-
"""Whether ``pid`` is still the recorded Connect server process generation."""
575+
def _same_process_generation(
576+
cls, pid: Optional[int], process_start_id: Optional[str]
577+
) -> Optional[bool]:
578+
"""Whether ``pid`` is alive and still has its recorded process generation."""
579+
if pid is None or not _pid_alive(pid):
580+
return False
581+
if process_start_id is None:
582+
return None
548583
current_start_id = cls._process_start_id(pid)
549584
if current_start_id is None:
550585
return None
551-
if current_start_id != process_start_id:
552-
return False
586+
return current_start_id == process_start_id
587+
588+
@classmethod
589+
def _same_server_instance(cls, pid: int, process_start_id: str) -> Optional[bool]:
590+
"""Whether ``pid`` is still the recorded Connect server process generation."""
591+
same_generation = cls._same_process_generation(pid, process_start_id)
592+
if same_generation is not True:
593+
return same_generation
553594
return _is_local_connect_server(pid)
554595

555596
@staticmethod
@@ -568,6 +609,18 @@ def _signal_server(cls, pid: int, process_start_id: str, sig: int) -> bool:
568609
"""Signal only the recorded generation of the managed Connect server."""
569610
return cls._same_server_instance(pid, process_start_id) is True and cls._signal(pid, sig)
570611

612+
@classmethod
613+
def _idle_timeout(cls) -> int:
614+
"""Seconds an unclaimed member may sit before it is retired.
615+
616+
Zero or a negative value disables idle retirement. Read the environment on each pass so
617+
every reaper uses the same source of truth.
618+
"""
619+
try:
620+
return int(os.environ["SPARK_LOCAL_CONNECT_POOL_IDLE_TIMEOUT"])
621+
except (KeyError, ValueError):
622+
return cls._DEFAULT_IDLE_TIMEOUT_SECONDS
623+
571624
def claim(self, fingerprint: str) -> Optional[PoolMember]:
572625
"""Claim the oldest usable member with this fingerprint, or ``None``. The rename to
573626
``claimed-<pid>-<uid>.json`` marks the member as owned by this process; the reaping
@@ -590,40 +643,102 @@ def claim(self, fingerprint: str) -> Optional[PoolMember]:
590643
data = self._directory.read_json(path)
591644
member = PoolMember.from_data(data) if data is not None else None
592645
if member is not None and member.fingerprint == fingerprint:
593-
candidates.append((member, uid, path))
646+
assert data is not None
647+
candidates.append((member, uid, path, data))
594648
candidates.sort(key=lambda c: c[0].created)
595-
for member, uid, path in candidates:
649+
for member, uid, path, data in candidates:
596650
if not member.is_usable():
597651
continue # left for the reaping rules to retire
652+
client_process_start_id = self._process_start_id(os.getpid())
653+
if client_process_start_id is not None:
654+
claimed_data = dict(data)
655+
claimed_data["client_process_start_id"] = client_process_start_id
656+
# Persist ownership before the atomic rename. If this process dies between the
657+
# write and rename, the extra field is harmless on the still-unclaimed record.
658+
self._directory.write_json(path, claimed_data)
598659
claim_path = self._directory.claimed_path(os.getpid(), uid)
599660
self._directory.rename(path, claim_path)
600661
member.claim_path = claim_path
601662
return member
602663
return None
603664

665+
def janitor(self) -> None:
666+
"""Reap unusable or orphaned pool members. Every rule is idempotent, so successive
667+
passes from any process are safe."""
668+
for uid in self._directory.uids():
669+
self.reap(uid)
670+
604671
def reap(self, uid: str) -> bool:
605-
"""Advance a retiring member and report whether nothing of it remains."""
672+
"""Apply the reaping rules to one member; ``True`` when nothing of it remains."""
606673
states = self._directory.states(uid)
607-
if "retired" in states:
674+
had_retired = "retired" in states
675+
if "server" in states:
676+
self._reap_server(uid, states["server"])
677+
states = self._directory.states(uid)
678+
if "claimed" in states:
679+
self._reap_claimed(uid, states["claimed"])
680+
states = self._directory.states(uid)
681+
if had_retired and "retired" in states:
608682
self._reap_retired(uid, states["retired"])
609683
return not self._directory.states(uid)
610684

685+
def _reap_server(self, uid: str, path: str) -> None:
686+
"""A ready member that is unusable (dead, unreachable, version-mismatched after an
687+
upgrade, or an unreadable record) or has sat unclaimed past the idle timeout: retire
688+
it."""
689+
data = self._directory.read_json(path)
690+
member = PoolMember.from_data(data) if data is not None else None
691+
server_pid, process_start_id = self._recover_server_handle(uid, data)
692+
idle = self._idle_timeout()
693+
expired = member is not None and idle > 0 and time.time() - member.created > idle
694+
if member is None or expired or not member.is_usable():
695+
self._retire(path, server_pid, process_start_id)
696+
697+
def _reap_claimed(self, uid: str, path: str) -> None:
698+
"""A claimed member whose client died without releasing it (e.g. SIGKILL), or whose
699+
server died under its client: retire it. Claims of this live process are its own."""
700+
data = self._directory.read_json(path)
701+
server_pid, process_start_id = self._recover_server_handle(uid, data)
702+
client_pid = self._directory.claiming_pid(path)
703+
client_process_start_id = PoolMember.client_process_start_id_from_data(data)
704+
if client_process_start_id is None:
705+
# Records written by older clients have no process generation. Preserve their
706+
# liveness-only behavior rather than risking retirement of a live claim.
707+
client_alive = client_pid == os.getpid() or _pid_alive(client_pid)
708+
else:
709+
client_alive = (
710+
self._same_process_generation(client_pid, client_process_start_id) is not False
711+
)
712+
if not client_alive or self._same_process_generation(server_pid, process_start_id) is False:
713+
self._retire(path, server_pid, process_start_id)
714+
611715
def _reap_retired(self, uid: str, path: str) -> None:
612716
"""A retiring member: drop it once its server is gone, hard-kill the server if it
613717
hangs in shutdown, and stop tracking a record whose process generation was reused."""
614718
data = self._directory.read_json(path)
615719
record_pid = RetiredState.pid_from_data(data)
616720
record_process_start_id = RetiredState.process_start_id_from_data(data)
617721
retired = RetiredState.retired_from_data(data)
722+
signalled = RetiredState.signalled_from_data(data)
618723
server_pid = self._recover_server_pid(uid, record_pid)
619724
# A generation id is meaningful only with the pid from the same record. The daemon pid
620725
# file has no companion generation id, so never pair a recovered daemon pid with
621726
# unrelated record data.
622727
process_start_id = record_process_start_id if record_pid is not None else None
623728
now = time.time()
624729
if server_pid is None:
625-
# Without a recoverable process handle, retain the state for a later read instead of
626-
# declaring a potentially live server gone.
730+
# A daemon pid may appear after a partial publication, so retain the state for one
731+
# recovery window. After that there is no process handle left to act on or observe.
732+
if retired is None or retired > now:
733+
self._directory.write_json(
734+
path,
735+
{
736+
"retired": now,
737+
"signalled": False,
738+
},
739+
)
740+
elif now - retired > self._RETIRE_GIVE_UP_AFTER_SECONDS:
741+
self._remove_retired(uid, path)
627742
return
628743
if not _pid_alive(server_pid):
629744
self._remove_retired(uid, path)
@@ -645,7 +760,13 @@ def _reap_retired(self, uid: str, path: str) -> None:
645760
# A crash while _retire rewrites the atomically renamed state can leave its old
646761
# payload. Restore a shutdown clock while preserving its process identity.
647762
self._directory.write_json(
648-
path, RetiredState(server_pid, process_start_id, now).as_data()
763+
path,
764+
RetiredState(
765+
server_pid,
766+
process_start_id,
767+
now,
768+
signalled=signalled is True,
769+
).as_data(),
649770
)
650771
return
651772
age = now - retired
@@ -662,28 +783,60 @@ def _reap_retired(self, uid: str, path: str) -> None:
662783
self._remove_retired(uid, path)
663784
elif is_server is True:
664785
self._signal(server_pid, signal.SIGKILL)
665-
else:
666-
# Retrying SIGTERM is idempotent and recovers a transient inspection failure during
667-
# _retire before the shutdown deadline escalates to SIGKILL.
668-
self._signal_server(server_pid, process_start_id, signal.SIGTERM)
786+
elif signalled is not True:
787+
# Retry a SIGTERM that was not confirmed during retirement. Persist success without
788+
# refreshing the retirement clock so repeated passes remain free and escalation is
789+
# still measured from the original transition.
790+
if self._signal_server(server_pid, process_start_id, signal.SIGTERM):
791+
self._directory.write_json(
792+
path,
793+
RetiredState(
794+
server_pid,
795+
process_start_id,
796+
retired,
797+
signalled=True,
798+
).as_data(),
799+
)
669800

670801
def _remove_retired(self, uid: str, path: str) -> None:
671802
self._directory.remove(path)
672803
self._directory.remove_member_dir(uid)
673804

674-
def _retire(self, state_path: str, server_pid: int, process_start_id: str) -> None:
805+
def _retire(
806+
self,
807+
state_path: str,
808+
server_pid: Optional[int],
809+
process_start_id: Optional[str],
810+
) -> None:
675811
"""Move a member into the retired state: signal its server and track the shutdown so
676812
:meth:`_reap_retired` can escalate if the JVM hangs."""
677813
_, uid = self._directory.parse_entry(os.path.basename(state_path))
678814
assert uid is not None
679-
self._signal_server(server_pid, process_start_id, signal.SIGTERM)
815+
signalled = False
816+
if server_pid is not None and process_start_id is not None:
817+
signalled = self._signal_server(server_pid, process_start_id, signal.SIGTERM)
680818
retired_path = self._directory.retired_path(uid)
681819
# Rename instead of removing the old state so a crash cannot leave a live server with
682820
# no state. If rewriting is interrupted, _reap_retired preserves the recoverable pid.
683821
self._directory.rename(state_path, retired_path)
684-
self._directory.write_json(
685-
retired_path, RetiredState(server_pid, process_start_id, time.time()).as_data()
686-
)
822+
retired_data: Dict[str, Any] = {
823+
"retired": time.time(),
824+
"signalled": signalled,
825+
}
826+
if server_pid is not None:
827+
retired_data["pid"] = server_pid
828+
if process_start_id is not None:
829+
retired_data["process_start_id"] = process_start_id
830+
self._directory.write_json(retired_path, retired_data)
831+
832+
def _recover_server_handle(
833+
self, uid: str, data: Optional[Dict[str, Any]]
834+
) -> Tuple[Optional[int], Optional[str]]:
835+
"""Recover a pid and only the process identity paired with that pid's record."""
836+
record_pid = PoolMember.pid_from_data(data)
837+
if record_pid is None:
838+
return self._recorded_daemon_pid(uid), None
839+
return record_pid, PoolMember.process_start_id_from_data(data)
687840

688841
def _recover_server_pid(self, uid: str, record_pid: Optional[int]) -> Optional[int]:
689842
"""Use a record pid when present, otherwise fall back to the daemon pid file.
@@ -701,7 +854,7 @@ def _recorded_daemon_pid(self, uid: str) -> Optional[int]:
701854

702855
def release(self, member: PoolMember) -> None:
703856
"""Retire this process's claimed member; the shutdown completes in the background,
704-
ready for a later lifecycle pass to finish.
857+
ready for a later janitor pass to finish.
705858
706859
This method acquires the pool-directory lock and must not be called while the same pool
707860
directory is already locked, including through a different ``PoolDirectory`` instance.

0 commit comments

Comments
 (0)