@@ -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 }
@@ -65,6 +67,8 @@ private[kyuubi] class EngineRef(
65
67
66
68
private val timeout : Long = conf.get(ENGINE_INIT_TIMEOUT )
67
69
70
+ private val objectMapper = new ObjectMapper ().registerModule(DefaultScalaModule )
71
+
68
72
// Share level of the engine
69
73
private val shareLevel : ShareLevel = ShareLevel .withName(conf.get(ENGINE_SHARE_LEVEL ))
70
74
@@ -324,10 +328,10 @@ private[kyuubi] class EngineRef(
324
328
}
325
329
}
326
330
327
- def getAdaptivePoolId (poolSize : Int ): Int = {
331
+ private def getAdaptivePoolId (poolSize : Int ): Int = {
328
332
val sessionThreshold = conf.get(ENGINE_POOL_ADAPTIVE_SESSION_THRESHOLD )
329
333
val metricsSpace =
330
- s " /metrics/ ${serverSpace}_ ${KYUUBI_VERSION }_ ${shareLevel}_ ${ engineType} / $user"
334
+ s " /metrics/ ${serverSpace}_ ${KYUUBI_VERSION }_ ${shareLevel}_ $engineType/ $user"
331
335
DiscoveryClientProvider .withDiscoveryClient(conf) { client =>
332
336
tryWithLock(client) {
333
337
if (client.pathNonExists(metricsSpace)) {
@@ -340,25 +344,23 @@ private[kyuubi] class EngineRef(
340
344
} else {
341
345
engineType match {
342
346
case SPARK_SQL =>
343
- val engineMetricsMap = metrics.map(p =>
344
- new String (client.getData(s " $metricsSpace/ $p" ))
345
- .split(" ;" )
346
- .map(_.split(" =" , 2 ))
347
- .filter(_.length == 2 )
348
- .map(kv => (kv.head, kv.last.toInt))
349
- .toMap)
347
+ val engineMetricsMap = metrics.map { p =>
348
+ objectMapper.readValue(
349
+ new String (client.getData(s " $metricsSpace/ $p" )),
350
+ classOf [Map [String , Int ]])
351
+ }
350
352
if (engineMetricsMap.isEmpty) {
351
353
return Random .nextInt(poolSize)
352
354
} else {
353
- val sortedEngineMetrics = engineMetricsMap.sortBy { map =>
354
- (
355
- map.getOrElse( " openSessionCount " , sessionThreshold),
356
- map.getOrElse(" activeTask " , 0 ))
357
- }
358
- val candidate = sortedEngineMetrics.head
359
- if (candidate.contains( " poolId " ) && ( candidate(
360
- " openSessionCount " ) < sessionThreshold || metrics.size == poolSize) ) {
361
- candidate(" poolId " )
355
+ val candidate = engineMetricsMap.filter(_.contains( " poolID " ))
356
+ .minBy { map =>
357
+ (
358
+ map.getOrElse(" openSessionCount " , sessionThreshold),
359
+ map.getOrElse( " activeTask " , 0 ))
360
+ }
361
+ if (( candidate.nonEmpty && candidate(" openSessionCount " ) < sessionThreshold) ||
362
+ metrics.size == poolSize) {
363
+ candidate(" poolID " )
362
364
} else {
363
365
Random .nextInt(poolSize)
364
366
}
0 commit comments