add totalCores

ilibx · Dec 4, 2018 · e3ea622 · e3ea622
1 parent 720bcc4
commit e3ea622
Showing 1 changed file with 13 additions and 2 deletions.
diff --git a/streamingpro-spark-2.3.0-adaptor/src/main/java/org/apache/spark/SparkInstanceService.scala b/streamingpro-spark-2.3.0-adaptor/src/main/java/org/apache/spark/SparkInstanceService.scala
@@ -1,5 +1,8 @@
 package org.apache.spark
 
+import net.csdn.common.reflect.ReflectHelper
+import org.apache.spark.scheduler.cluster.{CoarseGrainedSchedulerBackend, StandaloneSchedulerBackend}
+import org.apache.spark.scheduler.local.LocalSchedulerBackend
 import org.apache.spark.sql.SparkSession
 
 /**
@@ -17,8 +20,16 @@ class SparkInstanceService(session: SparkSession) {
       totalMemory += -1
 
     }
-    SparkInstanceResource(totalTasks, totalUsedMemory, totalMemory)
+    val totalCores = session.sparkContext.schedulerBackend match {
+      case sb: CoarseGrainedSchedulerBackend =>
+        ReflectHelper.field(sb, "totalCoreCount").asInstanceOf[Int]
+      case sb: LocalSchedulerBackend =>
+        java.lang.Runtime.getRuntime.availableProcessors
+      case sb: StandaloneSchedulerBackend =>
+        -1
+    }
+    SparkInstanceResource(totalCores.toLong, totalTasks, totalUsedMemory, totalMemory)
   }
 }
 
-case class SparkInstanceResource(totalTasks: Long, totalUsedMemory: Long, totalMemory: Long)
+case class SparkInstanceResource(totalCores: Long, totalTasks: Long, totalUsedMemory: Long, totalMemory: Long)