|
18 | 18 | package org.apache.spark.rpc |
19 | 19 |
|
20 | 20 | import java.net.URI |
| 21 | +import java.util.concurrent.TimeoutException |
21 | 22 |
|
22 | | -import scala.concurrent.{Await, Future} |
| 23 | +import scala.concurrent.{Awaitable, Await, Future} |
| 24 | +import scala.concurrent.duration._ |
23 | 25 | import scala.language.postfixOps |
24 | 26 |
|
25 | 27 | import org.apache.spark.{SecurityManager, SparkConf} |
@@ -66,7 +68,7 @@ private[spark] object RpcEnv { |
66 | 68 | */ |
67 | 69 | private[spark] abstract class RpcEnv(conf: SparkConf) { |
68 | 70 |
|
69 | | - private[spark] val defaultLookupTimeout = RpcUtils.lookupTimeout(conf) |
| 71 | + private[spark] val defaultLookupTimeout = RpcUtils.lookupRpcTimeout(conf) |
70 | 72 |
|
71 | 73 | /** |
72 | 74 | * Return RpcEndpointRef of the registered [[RpcEndpoint]]. Will be used to implement |
@@ -94,7 +96,7 @@ private[spark] abstract class RpcEnv(conf: SparkConf) { |
94 | 96 | * Retrieve the [[RpcEndpointRef]] represented by `uri`. This is a blocking action. |
95 | 97 | */ |
96 | 98 | def setupEndpointRefByURI(uri: String): RpcEndpointRef = { |
97 | | - Await.result(asyncSetupEndpointRefByURI(uri), defaultLookupTimeout) |
| 99 | + defaultLookupTimeout.awaitResult(asyncSetupEndpointRefByURI(uri)) |
98 | 100 | } |
99 | 101 |
|
100 | 102 | /** |
@@ -184,3 +186,107 @@ private[spark] object RpcAddress { |
184 | 186 | RpcAddress(host, port) |
185 | 187 | } |
186 | 188 | } |
| 189 | + |
| 190 | + |
| 191 | +/** |
| 192 | + * An exception thrown if RpcTimeout modifies a [[TimeoutException]]. |
| 193 | + */ |
| 194 | +private[rpc] class RpcTimeoutException(message: String, cause: TimeoutException) |
| 195 | + extends TimeoutException(message) { initCause(cause) } |
| 196 | + |
| 197 | + |
| 198 | +/** |
| 199 | + * Associates a timeout with a description so that a when a TimeoutException occurs, additional |
| 200 | + * context about the timeout can be amended to the exception message. |
| 201 | + * @param duration timeout duration in seconds |
| 202 | + * @param timeoutProp the configuration property that controls this timeout |
| 203 | + */ |
| 204 | +private[spark] class RpcTimeout(val duration: FiniteDuration, val timeoutProp: String) |
| 205 | + extends Serializable { |
| 206 | + |
| 207 | + /** Amends the standard message of TimeoutException to include the description */ |
| 208 | + private def createRpcTimeoutException(te: TimeoutException): RpcTimeoutException = { |
| 209 | + new RpcTimeoutException(te.getMessage() + ". This timeout is controlled by " + timeoutProp, te) |
| 210 | + } |
| 211 | + |
| 212 | + /** |
| 213 | + * PartialFunction to match a TimeoutException and add the timeout description to the message |
| 214 | + * |
| 215 | + * @note This can be used in the recover callback of a Future to add to a TimeoutException |
| 216 | + * Example: |
| 217 | + * val timeout = new RpcTimeout(5 millis, "short timeout") |
| 218 | + * Future(throw new TimeoutException).recover(timeout.addMessageIfTimeout) |
| 219 | + */ |
| 220 | + def addMessageIfTimeout[T]: PartialFunction[Throwable, T] = { |
| 221 | + // The exception has already been converted to a RpcTimeoutException so just raise it |
| 222 | + case rte: RpcTimeoutException => throw rte |
| 223 | + // Any other TimeoutException get converted to a RpcTimeoutException with modified message |
| 224 | + case te: TimeoutException => throw createRpcTimeoutException(te) |
| 225 | + } |
| 226 | + |
| 227 | + /** |
| 228 | + * Wait for the completed result and return it. If the result is not available within this |
| 229 | + * timeout, throw a [[RpcTimeoutException]] to indicate which configuration controls the timeout. |
| 230 | + * @param awaitable the `Awaitable` to be awaited |
| 231 | + * @throws RpcTimeoutException if after waiting for the specified time `awaitable` |
| 232 | + * is still not ready |
| 233 | + */ |
| 234 | + def awaitResult[T](awaitable: Awaitable[T]): T = { |
| 235 | + try { |
| 236 | + Await.result(awaitable, duration) |
| 237 | + } catch addMessageIfTimeout |
| 238 | + } |
| 239 | +} |
| 240 | + |
| 241 | + |
| 242 | +private[spark] object RpcTimeout { |
| 243 | + |
| 244 | + /** |
| 245 | + * Lookup the timeout property in the configuration and create |
| 246 | + * a RpcTimeout with the property key in the description. |
| 247 | + * @param conf configuration properties containing the timeout |
| 248 | + * @param timeoutProp property key for the timeout in seconds |
| 249 | + * @throws NoSuchElementException if property is not set |
| 250 | + */ |
| 251 | + def apply(conf: SparkConf, timeoutProp: String): RpcTimeout = { |
| 252 | + val timeout = { conf.getTimeAsSeconds(timeoutProp) seconds } |
| 253 | + new RpcTimeout(timeout, timeoutProp) |
| 254 | + } |
| 255 | + |
| 256 | + /** |
| 257 | + * Lookup the timeout property in the configuration and create |
| 258 | + * a RpcTimeout with the property key in the description. |
| 259 | + * Uses the given default value if property is not set |
| 260 | + * @param conf configuration properties containing the timeout |
| 261 | + * @param timeoutProp property key for the timeout in seconds |
| 262 | + * @param defaultValue default timeout value in seconds if property not found |
| 263 | + */ |
| 264 | + def apply(conf: SparkConf, timeoutProp: String, defaultValue: String): RpcTimeout = { |
| 265 | + val timeout = { conf.getTimeAsSeconds(timeoutProp, defaultValue) seconds } |
| 266 | + new RpcTimeout(timeout, timeoutProp) |
| 267 | + } |
| 268 | + |
| 269 | + /** |
| 270 | + * Lookup prioritized list of timeout properties in the configuration |
| 271 | + * and create a RpcTimeout with the first set property key in the |
| 272 | + * description. |
| 273 | + * Uses the given default value if property is not set |
| 274 | + * @param conf configuration properties containing the timeout |
| 275 | + * @param timeoutPropList prioritized list of property keys for the timeout in seconds |
| 276 | + * @param defaultValue default timeout value in seconds if no properties found |
| 277 | + */ |
| 278 | + def apply(conf: SparkConf, timeoutPropList: Seq[String], defaultValue: String): RpcTimeout = { |
| 279 | + require(timeoutPropList.nonEmpty) |
| 280 | + |
| 281 | + // Find the first set property or use the default value with the first property |
| 282 | + val itr = timeoutPropList.iterator |
| 283 | + var foundProp: Option[(String, String)] = None |
| 284 | + while (itr.hasNext && foundProp.isEmpty){ |
| 285 | + val propKey = itr.next() |
| 286 | + conf.getOption(propKey).foreach { prop => foundProp = Some(propKey, prop) } |
| 287 | + } |
| 288 | + val finalProp = foundProp.getOrElse(timeoutPropList.head, defaultValue) |
| 289 | + val timeout = { Utils.timeStringAsSeconds(finalProp._2) seconds } |
| 290 | + new RpcTimeout(timeout, finalProp._1) |
| 291 | + } |
| 292 | +} |
0 commit comments