@@ -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
488516class 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