@@ -23,6 +23,8 @@ import scala.collection.JavaConverters._
23
23
import scala .util .Random
24
24
25
25
import com .codahale .metrics .MetricRegistry
26
+ import com .fasterxml .jackson .databind .ObjectMapper
27
+ import com .fasterxml .jackson .module .scala .DefaultScalaModule
26
28
import com .google .common .annotations .VisibleForTesting
27
29
28
30
import org .apache .kyuubi .{KYUUBI_VERSION , KyuubiSQLException , Logging , Utils }
@@ -66,6 +68,8 @@ private[kyuubi] class EngineRef(
66
68
67
69
private val timeout : Long = conf.get(ENGINE_INIT_TIMEOUT )
68
70
71
+ private val objectMapper = new ObjectMapper ().registerModule(DefaultScalaModule )
72
+
69
73
// Share level of the engine
70
74
private val shareLevel : ShareLevel = ShareLevel .withName(conf.get(ENGINE_SHARE_LEVEL ))
71
75
@@ -371,10 +375,10 @@ private[kyuubi] class EngineRef(
371
375
}
372
376
}
373
377
374
- def getAdaptivePoolId (poolSize : Int ): Int = {
378
+ private def getAdaptivePoolId (poolSize : Int ): Int = {
375
379
val sessionThreshold = conf.get(ENGINE_POOL_ADAPTIVE_SESSION_THRESHOLD )
376
380
val metricsSpace =
377
- s " /metrics/ ${serverSpace}_ ${KYUUBI_VERSION }_ ${shareLevel}_ ${ engineType} / $user"
381
+ s " /metrics/ ${serverSpace}_ ${KYUUBI_VERSION }_ ${shareLevel}_ $engineType/ $user"
378
382
DiscoveryClientProvider .withDiscoveryClient(conf) { client =>
379
383
tryWithLock(client) {
380
384
if (client.pathNonExists(metricsSpace)) {
@@ -387,25 +391,23 @@ private[kyuubi] class EngineRef(
387
391
} else {
388
392
engineType match {
389
393
case SPARK_SQL =>
390
- val engineMetricsMap = metrics.map(p =>
391
- new String (client.getData(s " $metricsSpace/ $p" ))
392
- .split(" ;" )
393
- .map(_.split(" =" , 2 ))
394
- .filter(_.length == 2 )
395
- .map(kv => (kv.head, kv.last.toInt))
396
- .toMap)
394
+ val engineMetricsMap = metrics.map { p =>
395
+ objectMapper.readValue(
396
+ new String (client.getData(s " $metricsSpace/ $p" )),
397
+ classOf [Map [String , Int ]])
398
+ }
397
399
if (engineMetricsMap.isEmpty) {
398
400
return Random .nextInt(poolSize)
399
401
} else {
400
- val sortedEngineMetrics = engineMetricsMap.sortBy { map =>
401
- (
402
- map.getOrElse( " openSessionCount " , sessionThreshold),
403
- map.getOrElse(" activeTask " , 0 ))
404
- }
405
- val candidate = sortedEngineMetrics.head
406
- if (candidate.contains( " poolId " ) && ( candidate(
407
- " openSessionCount " ) < sessionThreshold || metrics.size == poolSize) ) {
408
- candidate(" poolId " )
402
+ val candidate = engineMetricsMap.filter(_.contains( " poolID " ))
403
+ .minBy { map =>
404
+ (
405
+ map.getOrElse(" openSessionCount " , sessionThreshold),
406
+ map.getOrElse( " activeTask " , 0 ))
407
+ }
408
+ if (( candidate.nonEmpty && candidate(" openSessionCount " ) < sessionThreshold) ||
409
+ metrics.size == poolSize) {
410
+ candidate(" poolID " )
409
411
} else {
410
412
Random .nextInt(poolSize)
411
413
}
0 commit comments