[CELEBORN-2415] Add Celeborn Spark UI extension

client-spark/spark-3-shaded/pom.xml

@@ -34,6 +34,11 @@
       <artifactId>celeborn-client-spark-3_${scala.binary.version}</artifactId>
       <version>${project.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.celeborn</groupId>
+      <artifactId>celeborn-client-spark-3-ui_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+    </dependency>
   </dependencies>
 
   <build>

client-spark/spark-3-ui/pom.xml

@@ -0,0 +1,91 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+  ~ Licensed to the Apache Software Foundation (ASF) under one or more
+  ~ contributor license agreements.  See the NOTICE file distributed with
+  ~ this work for additional information regarding copyright ownership.
+  ~ The ASF licenses this file to You under the Apache License, Version 2.0
+  ~ (the "License"); you may not use this file except in compliance with
+  ~ the License.  You may obtain a copy of the License at
+  ~
+  ~    http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <modelVersion>4.0.0</modelVersion>
+  <parent>
+    <groupId>org.apache.celeborn</groupId>
+    <artifactId>celeborn-parent_${scala.binary.version}</artifactId>
+    <version>${project.version}</version>
+    <relativePath>../../pom.xml</relativePath>
+  </parent>
+
+  <artifactId>celeborn-client-spark-3-ui_${scala.binary.version}</artifactId>
+  <packaging>jar</packaging>
+  <name>Celeborn Spark UI Plugin</name>
+
+  <dependencies>
+    <dependency>
+      <groupId>org.apache.celeborn</groupId>
+      <artifactId>celeborn-client-spark-3_${scala.binary.version}</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.spark</groupId>
+      <artifactId>spark-core_${scala.binary.version}</artifactId>
+      <scope>provided</scope>
+      <exclusions>
+        <exclusion>
+          <groupId>org.xerial.snappy</groupId>
+          <artifactId>snappy-java</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+    <dependency>
+      <groupId>javax.servlet</groupId>
+      <artifactId>javax.servlet-api</artifactId>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>jakarta.servlet</groupId>
+      <artifactId>jakarta.servlet-api</artifactId>
+      <scope>provided</scope>
+    </dependency>
+  </dependencies>
+
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>net.alchim31.maven</groupId>
+        <artifactId>scala-maven-plugin</artifactId>
+      </plugin>
+      <plugin>
+        <groupId>com.diffplug.spotless</groupId>
+        <artifactId>spotless-maven-plugin</artifactId>
+      </plugin>
+      <plugin>
+        <groupId>org.codehaus.mojo</groupId>
+        <artifactId>build-helper-maven-plugin</artifactId>
+        <version>3.6.0</version>
+        <executions>
+          <execution>
+            <id>add-servlet-source</id>
+            <goals>
+              <goal>add-source</goal>
+            </goals>
+            <phase>generate-sources</phase>
+            <configuration>
+              <sources>
+                <source>src/main/${servlet.source.dir}</source>
+              </sources>
+            </configuration>
+          </execution>
+        </executions>
+      </plugin>
+    </plugins>
+  </build>
+</project>

client-spark/spark-3-ui/src/main/resources/META-INF/services/org.apache.spark.status.AppHistoryServerPlugin

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+org.apache.spark.shuffle.celeborn.CelebornHistoryServerPlugin

client-spark/spark-3-ui/src/main/scala-spark3/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala

@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.celeborn.ui
+
+// Spark 3.x uses javax.servlet. This type alias bridges the API difference
+// so a single CelebornShufflePage.scala compiles for both Spark 3 and Spark 4.
+object SparkServletBridge {
+  type HttpServletRequest = javax.servlet.http.HttpServletRequest
+}

client-spark/spark-3-ui/src/main/scala-spark4/org/apache/spark/shuffle/celeborn/ui/SparkServletBridge.scala

@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.celeborn.ui
+
+// Spark 4.x uses jakarta.servlet. This type alias bridges the API difference
+// so a single CelebornShufflePage.scala compiles for both Spark 3 and Spark 4.
+object SparkServletBridge {
+  type HttpServletRequest = jakarta.servlet.http.HttpServletRequest
+}

client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornHistoryServerPlugin.scala

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.celeborn
+
+import org.apache.spark.SparkConf
+import org.apache.spark.scheduler.SparkListener
+import org.apache.spark.shuffle.celeborn.ui.CelebornUITab
+import org.apache.spark.status.{AppHistoryServerPlugin, ElementTrackingStore}
+import org.apache.spark.ui.SparkUI
+
+/** Registered via SPI at META-INF/services/org.apache.spark.status.AppHistoryServerPlugin. */
+class CelebornHistoryServerPlugin extends AppHistoryServerPlugin {
+
+  override def createListeners(
+      conf: SparkConf,
+      store: ElementTrackingStore): Seq[SparkListener] = {
+    val listener = new CelebornListener(store, conf, requirePluginOptIn = true)
+    // Persist the final accumulated values when replay finishes, covering logs
+    // whose last events fall inside the throttle interval and that may lack an
+    // ApplicationEnd event.
+    store.onFlush(listener.flush())
+    Seq(listener)
+  }
+
+  override def setupUI(ui: SparkUI): Unit = {
+    val statusStore = new CelebornStatusStore(ui.store.store)
+    // Only attach the tab for applications that opted in via spark.plugins or
+    // spark.plugins.defaultList.
+    if (statusStore.extensionEnabled()) {
+      new CelebornUITab(statusStore, ui)
+    }
+  }
+
+  override def displayOrder: Int = 1
+}

client-spark/spark-3-ui/src/main/scala/org/apache/spark/shuffle/celeborn/CelebornListener.scala

@@ -0,0 +1,143 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.shuffle.celeborn
+
+import java.util.concurrent.atomic.AtomicLong
+
+import org.apache.spark.SparkConf
+import org.apache.spark.internal.Logging
+import org.apache.spark.scheduler._
+import org.apache.spark.util.kvstore.KVStore
+
+/**
+ * Collects Celeborn shuffle metrics into the Spark KVStore for live UI and
+ * HistoryServer replay.
+ *
+ * When `requirePluginOptIn` is true (History Server replay), collection stays
+ * disabled until the application's recorded `spark.plugins` or
+ * `spark.plugins.defaultList` contains [[CelebornPlugin]], and an enable marker
+ * is persisted so `setupUI` can decide whether to attach the Celeborn tab.
+ */
+private[celeborn] class CelebornListener(
+    val kvstore: KVStore,
+    val conf: SparkConf,
+    requirePluginOptIn: Boolean = false)
+  extends SparkListener with Logging {
+
+  @volatile private var pluginEnabled = !requirePluginOptIn
+
+  private val totalWriteBytes = new AtomicLong(0L)
+  private val totalWriteTimeMs = new AtomicLong(0L)
+  private val totalReadBytes = new AtomicLong(0L)
+  private val totalFetchWaitTimeMs = new AtomicLong(0L)
+  private val totalTaskDurationMs = new AtomicLong(0L)
+
+  private val lastUpdateTimestamp = new AtomicLong(-1L)
+  private val updateIntervalMillis = 5000L
+
+  def register(sc: org.apache.spark.SparkContext): Unit = {
+    sc.addSparkListener(this)
+    logInfo("CelebornListener registered successfully")
+  }
+
+  override def onTaskEnd(taskEnd: SparkListenerTaskEnd): Unit = {
+    if (!pluginEnabled) {
+      return
+    }
+    Option(taskEnd.taskMetrics).foreach { metrics =>
+      totalWriteBytes.addAndGet(metrics.shuffleWriteMetrics.bytesWritten)
+      // writeTime is in nanoseconds; normalize to ms.
+      totalWriteTimeMs.addAndGet(metrics.shuffleWriteMetrics.writeTime / 1000000L)
+      totalReadBytes.addAndGet(metrics.shuffleReadMetrics.totalBytesRead)
+      totalFetchWaitTimeMs.addAndGet(metrics.shuffleReadMetrics.fetchWaitTime)
+      totalTaskDurationMs.addAndGet(taskEnd.taskInfo.duration)
+    }
+    mayUpdate()
+  }
+
+  override def onJobEnd(jobEnd: SparkListenerJobEnd): Unit = {
+    // Flush per-job so that tasks finishing within the throttle interval of the
+    // previous flush are not lost when the application stays alive but idle.
+    mayUpdate(force = true)
+  }
+
+  override def onEnvironmentUpdate(environmentUpdate: SparkListenerEnvironmentUpdate): Unit = {
+    val sparkProps = environmentUpdate.environmentDetails
+      .getOrElse("Spark Properties", Seq.empty)
+    if (!pluginEnabled) {
+      val pluginClass = classOf[CelebornPlugin].getName
+      // Spark loads plugins from both keys: `spark.plugins.defaultList` allows a
+      // default plugin list in the config file that `spark.plugins` does not overwrite.
+      val optedIn = sparkProps.exists { case (k, v) =>
+        (k == "spark.plugins" || k == "spark.plugins.defaultList") &&
+          v.split(",").exists(_.trim == pluginClass)
+      }
+      if (optedIn) {
+        pluginEnabled = true
+        kvstore.write(new CelebornExtensionEnabledUIData())
+      } else {
+        return
+      }
+    }
+    val celebornProps = sparkProps
+      .filter { case (k, _) => k.startsWith("spark.celeborn.") }
+      .sortBy(_._1)
+    if (celebornProps.nonEmpty) {
+      kvstore.write(new CelebornPropertiesUIData(celebornProps.toList))
+    }
+  }
+
+  override def onApplicationEnd(applicationEnd: SparkListenerApplicationEnd): Unit = {
+    mayUpdate(force = true)
+    logInfo("CelebornListener: application ended, final flush completed")
+  }
+
+  /** Flushes the current aggregations immediately, bypassing the throttle. */
+  def flush(): Unit = {
+    lastUpdateTimestamp.set(System.currentTimeMillis())
+    flushAggregations()
+  }
+
+  private def mayUpdate(force: Boolean = false): Unit = {
+    val now = System.currentTimeMillis()
+    val last = lastUpdateTimestamp.get()
+    if (!force && (last != -1L && (now - last) < updateIntervalMillis)) {
+      return
+    }
+    if (lastUpdateTimestamp.compareAndSet(last, now) || force) {
+      flushAggregations()
+    }
+  }
+
+  private def flushAggregations(): Unit = {
+    if (!pluginEnabled) {
+      return
+    }
+    try {
+      kvstore.write(AggregatedTaskInfoUIData(
+        totalWriteBytes.get(),
+        totalWriteTimeMs.get(),
+        totalReadBytes.get(),
+        totalFetchWaitTimeMs.get(),
+        totalTaskDurationMs.get()))
+    } catch {
+      case e: Exception =>
+        logWarning("Failed to flush CelebornListener aggregations", e)
+    }
+  }
+}