haiyangsun-db commented on code in PR #55683:
URL: https://github.com/apache/spark/pull/55683#discussion_r3376308636
##########
udf/worker/core/src/main/scala/org/apache/spark/udf/worker/core/direct/DirectWorkerProcess.scala:
##########
@@ -72,20 +81,62 @@ class DirectWorkerProcess(
/** Number of sessions currently using this worker. */
def activeSessions: Int = activeSessionCount.get()
- /** Increments the active session count. */
- def acquireSession(): Unit = activeSessionCount.incrementAndGet()
+ /**
+ * Increments the active session count.
+ *
+ * Preconditions: must not be called after [[close]]. The dispatcher
+ * arbitrates acquire-vs-close ordering; calling this on a closed worker
+ * indicates a dispatcher-side bug.
+ *
+ * @throws IllegalStateException if the worker has already been closed.
+ */
+ def acquireSession(): Unit = {
+ if (closed.get()) {
+ throw new IllegalStateException(
+ s"cannot acquire session: worker $id is already closed")
+ }
+ activeSessionCount.incrementAndGet()
+ }
+
+ /**
+ * Returns true if a session has marked this worker as unsafe to reuse.
+ * Pool implementations MUST treat [[isInvalid]] as a hard rejection in
+ * [[onLastSessionReleased]] -- the worker must be torn down, not
+ * recycled. Once set, the flag is sticky for the worker's lifetime.
+ */
+ def isInvalid: Boolean = invalid.get()
+
+ /**
+ * Marks this worker as unsafe to return to a reuse pool. Idempotent;
+ * sticky once set. Sessions call this when they observe a transport
+ * failure, hung worker, or any termination path that leaves the worker
+ * in an unknown state.
+ */
+ override def markInvalid(): Unit = invalid.set(true)
/**
* Decrements the active session count. Fires [[onLastSessionReleased]]
- * on the 0-transition. A negative count indicates an unbalanced
- * acquire/release; we log and reset to 0 rather than silently mask it.
+ * on the 0-transition.
+ *
+ * A negative count indicates an unbalanced acquire/release and is
+ * surfaced as a warning. We deliberately do NOT reset the counter:
+ * a concurrent `acquireSession()` between the decrement here and a
+ * `set(0)` would be silently overwritten, and the next `releaseSession`
+ * would prematurely fire `onLastSessionReleased` and tear the worker
+ * down underneath the racing caller. The counter recovers naturally on
+ * subsequent balanced calls.
*/
- def releaseSession(): Unit = {
+ override def releaseSession(): Unit = {
val c = activeSessionCount.decrementAndGet()
if (c < 0) {
- logger.warn(
- s"releaseSession called without a matching acquireSession (count=$c)")
- activeSessionCount.set(0)
+ // An unbalanced acquire/release is a dispatcher-side bug. Assert so
+ // it fails loudly under test (-ea is on for Spark's test JVMs); in
+ // production we only warn rather than throw, because releaseSession
+ // runs on the WorkerSession.close() path and an exception here would
+ // mask any in-flight failure the caller is already handling.
+ val msg = s"releaseSession called without a matching acquireSession
(count=$c)"
+ assert(c >= 0, msg)
Review Comment:
fixed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
