Github user squito commented on a diff in the pull request: https://github.com/apache/spark/pull/6205#discussion_r32259066 --- Diff: core/src/main/scala/org/apache/spark/rpc/RpcEnv.scala --- @@ -182,3 +185,117 @@ private[spark] object RpcAddress { RpcAddress(host, port) } } + + +/** + * An exception thrown if RpcTimeout modifies a [[TimeoutException]]. + */ +private[rpc] class RpcTimeoutException(message: String, cause: TimeoutException) + extends TimeoutException(message) { initCause(cause) } + + +/** + * Associates a timeout with a description so that a when a TimeoutException occurs, additional + * context about the timeout can be amended to the exception message. + * @param timeout timeout duration in seconds + * @param description description to be displayed in a timeout exception + */ +private[spark] class RpcTimeout(timeout: FiniteDuration, description: String) { + + /** Get the timeout duration */ + def duration: FiniteDuration = timeout + + /** Get the message associated with this timeout */ + def message: String = description + + /** Amends the standard message of TimeoutException to include the description */ + private def createRpcTimeoutException(te: TimeoutException): RpcTimeoutException = { + new RpcTimeoutException(te.getMessage() + " " + description, te) + } + + /** + * PartialFunction to match a TimeoutException and add the timeout description to the message + * + * @note This can be used in the recover callback of a Future to add to a TimeoutException + * Example: + * val timeout = new RpcTimeout(5 millis, "short timeout") + * Future(throw new TimeoutException).recover(timeout.addMessageIfTimeout) + */ + def addMessageIfTimeout[T]: PartialFunction[Throwable, T] = { + // The exception has already been converted to a RpcTimeoutException so just raise it + case rte: RpcTimeoutException => throw rte + // Any other TimeoutException get converted to a RpcTimeoutException with modified message + case te: TimeoutException => throw createRpcTimeoutException(te) + } + + /** + * Waits for a completed result to catch and amend a TimeoutException message + * @param awaitable the `Awaitable` to be awaited + * @throws RpcTimeoutException if after waiting for the specified time `awaitable` + * is still not ready + */ + def awaitResult[T](awaitable: Awaitable[T]): T = { + try { + Await.result(awaitable, duration) + } catch addMessageIfTimeout + } +} + + +/** + * Create an RpcTimeout using a configuration property that controls the timeout duration so when + * a TimeoutException is thrown, the property key will be indicated in the message. + */ --- End diff -- I think this comment can be deleted, it doesn't really apply to `object RpcTimeout`, really its for the individual `apply` methods which you've already got good docs for.
--- If your project is set up for it, you can reply to this email and have your reply appear on GitHub as well. If your project does not have this feature enabled and wishes so, or if the feature is enabled but not working, please contact infrastructure at infrastruct...@apache.org or file a JIRA ticket with INFRA. --- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org