This is an automated email from the ASF dual-hosted git repository.
SteNicholas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 0d63a20e4c [CELEBORN-2447] Add a /healthz endpoint to master and worker
0d63a20e4c is described below
commit 0d63a20e4c08b713e1d0cf69a9d03c67cb6db6af
Author: strelok89 <[email protected]>
AuthorDate: Mon Sep 14 19:10:22 2026 +0800
[CELEBORN-2447] Add a /healthz endpoint to master and worker
### What changes were proposed in this pull request?
Add a `GET /healthz` endpoint served by both master and worker, whose HTTP
status code reflects whether the service is able to serve: `200` when healthy
and `503` when not.
- `HttpService#healthCheck` defaults to a shallow check that reports
healthy once the HTTP service is available. This is what the master uses.
- `Worker` overrides it to report healthy only when the worker is
registered with the master and its state is `Normal`.
- `/healthz` is added to the default HTTP authentication bypass paths.
The endpoint is mounted at the root rather than under `/api/v1` so that a
probe is not coupled to the API version, consistent with the other operational
paths (`/ping`, `/metrics/prometheus`). The checked-in OpenAPI spec covers only
`/api/v1`, so the generated client is unaffected.
Example:
```
$ curl -s -o /dev/null -w '%{http_code}\n' http://worker:9096/healthz
200
$ curl -s http://worker:9096/healthz
{"service":"worker","healthy":false,"reason":"worker is not registered with
master"}
```
### Why are the changes needed?
Celeborn currently exposes no endpoint whose status code reflects whether
the process can serve, so a meaningful Kubernetes probe cannot be written.
`/ping` always returns `200` regardless of state, and `/api/v1/workers` returns
`200` even when `isRegistered` is false, so operators have to fall back to an
`exec` probe that shells out and greps the JSON body.
The practical cost is rolling updates. Without a readiness signal, a
StatefulSet marks a worker pod Ready as soon as its container process starts
and immediately proceeds to the next ordinal, before the restarted worker has
re-registered with the master. On a large cluster this can remove a meaningful
fraction of workers from service before any of them rejoin.
Two design points worth calling out for review:
**The worker check mirrors the master's own definition.**
`AbstractMetaManager#isWorkerAvailable` accepts only `Normal`, so a worker that
is idle, decommissioning or exiting is already excluded from `availableWorkers`
and never selected in `offerSlots`. Reporting such a worker as healthy would
put `/healthz` at odds with the master. Whether `Idle` specifically should fail
readiness is discussed on CELEBORN-2447 — it is reachable only through an
explicit `DecommissionThenIdle` event, s [...]
**The master check is deliberately shallow.** Masters are fronted by a
headless Service that does not set `publishNotReadyAddresses`. A quorum-aware
or leader-aware check would report every master as unhealthy during a cold
start, withholding their DNS records and preventing them from discovering each
other to form quorum. A follower is also a healthy replica. Quorum and
leadership remain observable through `/api/v1/ratis` and `/api/v1/masters`.
`/healthz` bypasses authentication by default because a kubelet cannot
present credentials, and `celeborn.http.auth.bypass.api.paths` defaults to
empty (CELEBORN-2278).
Wiring probes into the Helm chart, which currently defines no
`livenessProbe`, `readinessProbe` or `startupProbe` for either role, is
intended as a follow-up so the two changes can be reviewed independently.
### Does this PR resolve a correctness bug?
- [ ] Yes
### Does this PR introduce _any_ user-facing change?
- [x] Yes
A new `/healthz` endpoint on master and worker, documented in
`docs/restapi.md`.
### How was this patch tested?
New tests:
- `ApiBaseResourceSuite`: `/healthz` returns `200` and reports the correct
service name, for both master and worker.
- `ApiBaseResourceAuthenticationSuite`: `/healthz` is reachable without
credentials.
- `ApiWorkerResourceSuite`: `/healthz` returns `503` when the worker is not
registered.
Verified locally on JDK 11: `celeborn-service/Test/compile` and
`celeborn-worker/Test/compile` pass, `ApiMasterResourceSuite` (19/19) and
`ApiMasterResourceAuthenticationSuite` (9/9) pass, and `spotless:check` passes
for the `service` and `worker` modules.
The worker-side suites were not run locally: they were developed on
Windows, where any `MiniClusterFeature` test fails during setup because
`CelebornConf#workerBaseDirs` splits each storage dir on `:` and rejects
`C:\...` paths. That is pre-existing and unrelated to this change, and those
suites are covered by CI.
Closes #3832 from strelok89/CELEBORN-2447.
Authored-by: strelok89 <[email protected]>
Signed-off-by: Nicholas Jiang <[email protected]>
---
docs/restapi.md | 32 ++++++++++
.../celeborn/server/common/HttpService.scala | 11 ++++
.../server/common/http/api/HealthResource.scala | 71 ++++++++++++++++++++++
.../http/authentication/AuthenticationFilter.scala | 2 +-
.../http/ApiBaseResourceAuthenticationSuite.scala | 7 +++
.../server/common/http/ApiBaseResourceSuite.scala | 10 +++
.../celeborn/service/deploy/worker/Worker.scala | 33 ++++++++++
.../deploy/worker/WorkerStatusManager.scala | 4 +-
.../service/deploy/MiniClusterFeature.scala | 3 +
.../service/deploy/worker/WorkerSuite.scala | 49 ++++++++++++++-
.../worker/http/api/ApiWorkerResourceSuite.scala | 65 ++++++++++++++++++++
11 files changed, 284 insertions(+), 3 deletions(-)
diff --git a/docs/restapi.md b/docs/restapi.md
index 4f2892aec4..73e00b9a45 100644
--- a/docs/restapi.md
+++ b/docs/restapi.md
@@ -34,6 +34,38 @@ The configuration of `<master-http-host>`,
`<master-http-port>`, `<worker-http-h
| celeborn.worker.http.host | 0.0.0.0 | Worker's http host. | 0.4.0 |
| celeborn.worker.http.port | 9096 | Worker's http port. | 0.4.0 |
+### Health Check API (Since 1.0.0)
+
+Both master and worker serve a health check endpoint at `/healthz`, intended
to back a Kubernetes
+readiness probe. It responds `200` when the service is able to serve and `503`
when it is not,
+so that it can be consumed with a plain `httpGet` probe.
+
+| Path | Method | Meaning
|
+|------------|--------|-----------------------------------------------------------------------------------|
+| `/healthz` | GET | Whether the service is able to serve. Returns `200`
when healthy, `503` when not. |
+
+The response body reports the reason when the service is not able to serve:
+
+```json
+{"service":"worker","healthy":false,"reason":"worker is not registered with
master"}
+```
+
+The check differs by role:
+
+- **Master**: healthy once its HTTP service is available. The check is
deliberately shallow and
+ does not depend on Ratis quorum or leadership, since a follower master is a
healthy replica and
+ a quorum-aware check would report every master as unhealthy during a cold
start. Use
+ `/api/v1/ratis` and `/api/v1/masters` to monitor quorum and leadership.
+- **Worker**: healthy only when it is registered with the master and its state
is `Normal`. This
+ matches the master's own definition of an available worker, so a worker that
is idle,
+ decommissioning or exiting is reported as not able to serve, consistent with
the fact that the
+ master no longer offers slots to it. A worker also reports not able to serve
while it is
+ re-registering after a heartbeat response told it that the master no longer
knows about it.
+ Note that the worker keeps serving push and fetch requests throughout that
window, since it
+ still holds the data clients hold locations for; only readiness is withheld.
+
+`/healthz` bypasses HTTP authentication by default, since a kubelet cannot
present credentials.
+
### Deprecated REST APIs
Since 0.6.0, the legacy REST APIs are deprecated and will be removed in the
future.
diff --git
a/service/src/main/scala/org/apache/celeborn/server/common/HttpService.scala
b/service/src/main/scala/org/apache/celeborn/server/common/HttpService.scala
index a66f53305d..4f06075986 100644
--- a/service/src/main/scala/org/apache/celeborn/server/common/HttpService.scala
+++ b/service/src/main/scala/org/apache/celeborn/server/common/HttpService.scala
@@ -185,6 +185,17 @@ abstract class HttpService extends Service with Logging {
def exit(exitType: String): String = throw new
UnsupportedOperationException()
+ /**
+ * Whether the service is able to serve, along with the reason when it is
not. Intended to back
+ * the `/healthz` endpoint consumed by readiness probes.
+ *
+ * The default is a shallow check that only reports that the HTTP service is
available, which is
+ * what the master needs: a follower master is a healthy replica, and gating
on Ratis quorum or
+ * leadership would leave every master unhealthy during a cold start. The
worker overrides this
+ * with a check of its registration and state.
+ */
+ def healthCheck(): HandleResponse = (true, "")
+
def handleWorkerEvent(
workerEventType: WorkerEventType,
workers: Seq[WorkerInfo]): HandleResponse =
diff --git
a/service/src/main/scala/org/apache/celeborn/server/common/http/api/HealthResource.scala
b/service/src/main/scala/org/apache/celeborn/server/common/http/api/HealthResource.scala
new file mode 100644
index 0000000000..51d7ea8dc3
--- /dev/null
+++
b/service/src/main/scala/org/apache/celeborn/server/common/http/api/HealthResource.scala
@@ -0,0 +1,71 @@
+/*
+ * 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.celeborn.server.common.http.api
+
+import javax.ws.rs.{GET, Path, Produces}
+import javax.ws.rs.core.{MediaType, Response}
+
+import io.swagger.v3.oas.annotations.Operation
+import io.swagger.v3.oas.annotations.media.{Content, Schema}
+import io.swagger.v3.oas.annotations.responses.ApiResponse
+import io.swagger.v3.oas.annotations.tags.Tag
+
+/**
+ * The response of the health check, see [[HealthResource]].
+ *
+ * @param service the name of the service, either `master` or `worker`.
+ * @param healthy whether the service is able to serve.
+ * @param reason the reason why the service is not able to serve, empty when
healthy.
+ */
+case class HealthCheckResponse(service: String, healthy: Boolean, reason:
String)
+
+@Tag(name = "Health")
+@Path("/healthz")
+@Produces(Array(MediaType.APPLICATION_JSON))
+private[api] class HealthResource extends ApiRequestContext {
+
+ @Operation(description =
+ "Check whether the service is able to serve. Returns 200 when the service
is healthy " +
+ "and 503 otherwise, which is intended to be consumed by a readiness
probe. " +
+ "The master reports healthy once its HTTP service is available,
including when it is " +
+ "a follower. The worker reports healthy only when it is registered with
the master and " +
+ "its state is Normal, so a worker that is decommissioning, idle or
exiting is reported " +
+ "as not able to serve.")
+ @ApiResponse(
+ responseCode = "200",
+ description = "The service is able to serve.",
+ content = Array(new Content(
+ mediaType = MediaType.APPLICATION_JSON,
+ schema = new Schema(implementation = classOf[HealthCheckResponse]))))
+ @ApiResponse(
+ responseCode = "503",
+ description = "The service is not able to serve.",
+ content = Array(new Content(
+ mediaType = MediaType.APPLICATION_JSON,
+ schema = new Schema(implementation = classOf[HealthCheckResponse]))))
+ @GET
+ def health(): Response = {
+ val (healthy, reason) = httpService.healthCheck()
+ val status =
+ if (healthy) Response.Status.OK else Response.Status.SERVICE_UNAVAILABLE
+ Response.status(status)
+ .`type`(MediaType.APPLICATION_JSON)
+ .entity(HealthCheckResponse(httpService.serviceName, healthy, reason))
+ .build()
+ }
+}
diff --git
a/service/src/main/scala/org/apache/celeborn/server/common/http/authentication/AuthenticationFilter.scala
b/service/src/main/scala/org/apache/celeborn/server/common/http/authentication/AuthenticationFilter.scala
index 55274fc8a4..9f994ccebc 100644
---
a/service/src/main/scala/org/apache/celeborn/server/common/http/authentication/AuthenticationFilter.scala
+++
b/service/src/main/scala/org/apache/celeborn/server/common/http/authentication/AuthenticationFilter.scala
@@ -221,7 +221,7 @@ class AuthenticationFilter(conf: CelebornConf, serviceName:
String) extends Filt
}
object AuthenticationFilter {
- private val BYPASS_DEFAULT_API_PATHS = Set("/openapi.json", "/openapi.yaml")
+ private val BYPASS_DEFAULT_API_PATHS = Set("/healthz", "/openapi.json",
"/openapi.yaml")
final val HTTP_CLIENT_IP_ADDRESS = new ThreadLocal[String]() {
override protected def initialValue: String = null
diff --git
a/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceAuthenticationSuite.scala
b/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceAuthenticationSuite.scala
index 7e62a331d4..8fb95c7ef7 100644
---
a/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceAuthenticationSuite.scala
+++
b/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceAuthenticationSuite.scala
@@ -26,6 +26,7 @@ import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.authentication.HttpAuthSchemes
import org.apache.celeborn.common.network.TestHelper
import
org.apache.celeborn.server.common.http.HttpAuthUtils.AUTHORIZATION_HEADER
+import org.apache.celeborn.server.common.http.api.HealthCheckResponse
import
org.apache.celeborn.server.common.http.authentication.{UserDefinedPasswordAuthenticationProviderImpl,
UserDefineTokenAuthenticationProviderImpl}
abstract class ApiBaseResourceAuthenticationSuite extends HttpTestHelper {
@@ -93,6 +94,12 @@ abstract class ApiBaseResourceAuthenticationSuite extends
HttpTestHelper {
}
}
+ test("health api does not need authentication") {
+ val response =
webTarget.path("healthz").request(MediaType.APPLICATION_JSON).get()
+ assert(HttpServletResponse.SC_OK == response.getStatus)
+ assert(response.readEntity(classOf[HealthCheckResponse]).healthy)
+ }
+
test("swagger api do not need authentication") {
Seq("swagger", "docs", "help").foreach { path =>
val response = webTarget.path(path).request(MediaType.TEXT_HTML).get()
diff --git
a/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceSuite.scala
b/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceSuite.scala
index 49d5138517..22e77d7a71 100644
---
a/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceSuite.scala
+++
b/service/src/test/scala/org/apache/celeborn/server/common/http/ApiBaseResourceSuite.scala
@@ -22,6 +22,7 @@ import javax.ws.rs.core.MediaType
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.network.TestHelper
+import org.apache.celeborn.server.common.http.api.HealthCheckResponse
abstract class ApiBaseResourceSuite extends HttpTestHelper {
celebornConf.set(CelebornConf.METRICS_ENABLED.key, "true")
@@ -35,6 +36,15 @@ abstract class ApiBaseResourceSuite extends HttpTestHelper {
assert(response.readEntity(classOf[String]) == "pong")
}
+ test("health") {
+ val response =
webTarget.path("healthz").request(MediaType.APPLICATION_JSON).get()
+ assert(HttpServletResponse.SC_OK == response.getStatus)
+ val health = response.readEntity(classOf[HealthCheckResponse])
+ assert(health.healthy)
+ assert(health.reason.isEmpty)
+ assert(health.service == httpService.serviceName)
+ }
+
test("conf") {
val response = webTarget.path("conf").request(MediaType.TEXT_PLAIN).get()
assert(HttpServletResponse.SC_OK == response.getStatus)
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index 09999bad98..79f3f20b32 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -302,6 +302,15 @@ private[celeborn] class Worker(
// whether this Worker registered to Master successfully
val registered = new AtomicBoolean(false)
+
+ // Whether the master still knows about this worker, per the last heartbeat
response. This is
+ // separate from `registered`, which gates RPC serving and stays true across
a re-registration
+ // so that clients can keep fetching data this worker still holds.
+ val registeredInMasterView = new AtomicBoolean(true)
+
+ // Whether `initialize()` finished wiring up the push/fetch/replicate
handlers and the
+ // controller endpoint. Registration alone is not enough to serve traffic,
see the health check.
+ val initialized = new AtomicBoolean(false)
val shuffleMapperAttempts: ConcurrentHashMap[String, AtomicIntegerArray] =
JavaUtils.newConcurrentHashMap[String, AtomicIntegerArray]()
val shufflePartitionType: ConcurrentHashMap[String, PartitionType] =
@@ -537,6 +546,10 @@ private[celeborn] class Worker(
highWorkload,
workerStatusManager.currentWorkerStatus),
classOf[HeartbeatFromWorkerResponse])
+ handleHeartbeatResponse(response)
+ }
+
+ private[worker] def handleHeartbeatResponse(response:
HeartbeatFromWorkerResponse): Unit = {
response.expiredShuffleKeys.asScala.foreach(shuffleKey =>
workerInfo.releaseSlots(shuffleKey))
cleanTaskQueue.put(response.expiredShuffleKeys)
@@ -544,6 +557,7 @@ private[celeborn] class Worker(
workerStatusManager.doTransition(workerEvent)
if (!response.registered) {
logError("Worker not registered in master, clean expired shuffle data
and register again.")
+ registeredInMasterView.set(false)
try {
registerWithMaster()
} catch {
@@ -612,6 +626,7 @@ private[celeborn] class Worker(
controller.init(this)
rpcEnv.setupEndpoint(RpcNameConstants.WORKER_EP, controller)
+ initialized.set(true)
logInfo("Worker started.")
rpcEnv.awaitTermination()
@@ -724,6 +739,7 @@ private[celeborn] class Worker(
// Register successfully
if (null != resp && resp.getSuccess) {
registered.set(true)
+ registeredInMasterView.set(true)
logInfo("Register worker successfully.")
return
}
@@ -935,6 +951,23 @@ private[celeborn] class Worker(
sb.toString()
}
+ override def healthCheck(): HandleResponse = {
+ val state = workerStatusManager.currentWorkerStatus.getState
+ if (!initialized.get()) {
+ // The HTTP server starts and registration completes before the
push/fetch handlers and the
+ // controller endpoint are set up, so a probe in that interval must not
report healthy.
+ (false, "worker is still initializing")
+ } else if (!registered.get() || !registeredInMasterView.get()) {
+ (false, "worker is not registered with master")
+ } else if (state != State.Normal) {
+ // Only workers in Normal state are selected when the master offers
slots, see
+ // AbstractMetaManager#isWorkerAvailable.
+ (false, s"worker state is $state instead of ${State.Normal}")
+ } else {
+ (true, "")
+ }
+ }
+
override def listPartitionLocationInfo: String = {
val sb = new StringBuilder
sb.append("==================== Partition Location Info
=========================\n")
diff --git
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
index e31fcf576a..52dfe075ac 100644
---
a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
+++
b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerStatusManager.scala
@@ -34,7 +34,9 @@ import
org.apache.celeborn.service.deploy.worker.storage.StorageManager
private[celeborn] class WorkerStatusManager(conf: CelebornConf) extends
Logging {
- var currentWorkerStatus = WorkerStatus.normalWorkerStatus()
+ // Written under this manager's monitor but read without it, for example by
the /healthz
+ // endpoint on an HTTP thread, so it is published volatile.
+ @volatile var currentWorkerStatus = WorkerStatus.normalWorkerStatus()
var exitEventType = WorkerEventType.Immediately
private var worker: Worker = _
private var shutdown: AtomicBoolean = _
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
index e4b26a1c51..7904f6bb71 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/MiniClusterFeature.scala
@@ -257,6 +257,9 @@ trait MiniClusterFeature extends Logging with
RandomPortSupport {
if (!worker.registered.get()) {
throw new IllegalStateException(s"worker $i hasn't been
registered")
}
+ if (!worker.initialized.get()) {
+ throw new IllegalStateException(s"worker $i hasn't finished
initialization")
+ }
if (!workerInfos.contains(worker)) {
workerInfos.put(worker, threads(i))
}
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
index 26a1cb1b6d..09c4c78b51 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/WorkerSuite.scala
@@ -21,6 +21,7 @@ import java.io.File
import java.nio.file.{Files, Paths}
import java.util
import java.util.{HashSet => JHashSet}
+import java.util.concurrent.atomic.AtomicBoolean
import scala.collection.JavaConverters._
@@ -29,6 +30,8 @@ import org.mockito.{ArgumentCaptor, ArgumentMatchers,
MockedConstruction, Mockit
import org.mockito.MockedConstruction.MockInitializer
import org.mockito.Mockito.mockConstruction
import org.mockito.MockitoSugar._
+import org.mockito.invocation.InvocationOnMock
+import org.mockito.stubbing.Answer
import org.scalatest.BeforeAndAfterEach
import org.scalatest.funsuite.AnyFunSuite
@@ -36,7 +39,7 @@ import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.client.MasterClient
import org.apache.celeborn.common.identity.UserIdentifier
import org.apache.celeborn.common.protocol._
-import
org.apache.celeborn.common.protocol.message.ControlMessages.CommitFilesResponse
+import
org.apache.celeborn.common.protocol.message.ControlMessages.{CommitFilesResponse,
HeartbeatFromWorkerResponse}
import org.apache.celeborn.common.protocol.message.StatusCode
import org.apache.celeborn.common.quota.ResourceConsumption
import org.apache.celeborn.common.rpc.RpcCallContext
@@ -307,6 +310,50 @@ class WorkerSuite extends AnyFunSuite with
BeforeAndAfterEach with MiniClusterFe
assert(epochCommitMap.get(epoch2).response.status == StatusCode.SUCCESS)
}
+ test("CELEBORN-2447: heartbeat reporting not-registered clears the
master-view flag " +
+ "before re-registering") {
+ conf.set(CelebornConf.WORKER_STORAGE_DIRS.key, "/tmp")
+ // Recorded at the moment the re-registration RPC is issued, so the test
asserts the
+ // ordering of the flag clear rather than only its final value.
+ val flagWhenReRegistering = new AtomicBoolean(true)
+ val mockInitializer = {
+ // Old syntax needed for scala 2.11
+ new MockInitializer[MasterClient] {
+ override def prepare(instance: MasterClient, context:
MockedConstruction.Context): Unit = {
+ val answer = new Answer[PbRegisterWorkerResponse] {
+ override def answer(invocation: InvocationOnMock):
PbRegisterWorkerResponse = {
+ flagWhenReRegistering.set(worker.registeredInMasterView.get())
+ PbRegisterWorkerResponse.newBuilder().setSuccess(true).build()
+ }
+ }
+ Mockito.doAnswer(answer)
+ .when(instance)
+ .askSync(
+ ArgumentMatchers.any(classOf[PbRegisterWorker]),
+ ArgumentMatchers.eq(classOf[PbRegisterWorkerResponse]))
+ }
+ }
+ }
+ val mockedMasterClient = mockConstruction(classOf[MasterClient],
mockInitializer)
+ try {
+ worker = new Worker(conf, workerArgs)
+ worker.registered.set(true)
+ assert(worker.registeredInMasterView.get())
+
+ // WorkerEventType.None keeps doTransition a no-op; other events spawn
an exit thread.
+ worker.handleHeartbeatResponse(
+ HeartbeatFromWorkerResponse(new JHashSet[String](), registered =
false))
+
+ assert(!flagWhenReRegistering.get(), "master-view flag must be cleared
before re-registering")
+ // `registered` gates RPC serving and must survive the re-registration.
+ assert(worker.registered.get())
+ // Re-registration succeeded, so readiness is restored.
+ assert(worker.registeredInMasterView.get())
+ } finally {
+ mockedMasterClient.close()
+ }
+ }
+
test("CELEBORN-2257: Properly reports remote disks on worker registration") {
val mockInitializer = {
// Old syntax needed for scala 2.11
diff --git
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/http/api/ApiWorkerResourceSuite.scala
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/http/api/ApiWorkerResourceSuite.scala
index 408e394360..1f8c548ddf 100644
---
a/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/http/api/ApiWorkerResourceSuite.scala
+++
b/worker/src/test/scala/org/apache/celeborn/service/deploy/worker/http/api/ApiWorkerResourceSuite.scala
@@ -20,8 +20,11 @@ package org.apache.celeborn.service.deploy.worker.http.api
import javax.servlet.http.HttpServletResponse
import javax.ws.rs.core.MediaType
+import org.apache.celeborn.common.meta.WorkerStatus
+import org.apache.celeborn.common.protocol.PbWorkerStatus.State
import org.apache.celeborn.server.common.HttpService
import org.apache.celeborn.server.common.http.ApiBaseResourceSuite
+import org.apache.celeborn.server.common.http.api.HealthCheckResponse
import org.apache.celeborn.service.deploy.MiniClusterFeature
import org.apache.celeborn.service.deploy.worker.Worker
@@ -43,6 +46,68 @@ class ApiWorkerResourceSuite extends ApiBaseResourceSuite
with MiniClusterFeatur
shutdownMiniCluster()
}
+ test("health reports unavailable during the startup interval before
initialization completes") {
+ // `initialize()` starts the HTTP server and registers with the master
before the push/fetch
+ // handlers and the controller endpoint are set up. Clearing the flag
reproduces that interval:
+ // the worker is registered and in Normal state, but cannot serve traffic
yet.
+ assert(worker.registered.get())
+ assert(worker.registeredInMasterView.get())
+ assert(worker.workerStatusManager.getWorkerState() == State.Normal)
+ worker.initialized.set(false)
+ try {
+ val response =
webTarget.path("healthz").request(MediaType.APPLICATION_JSON).get()
+ assert(HttpServletResponse.SC_SERVICE_UNAVAILABLE == response.getStatus)
+ val health = response.readEntity(classOf[HealthCheckResponse])
+ assert(!health.healthy)
+ assert(health.reason.contains("still initializing"))
+ } finally {
+ worker.initialized.set(true)
+ }
+ }
+
+ test("health reports unavailable when the worker is not registered") {
+ worker.registered.set(false)
+ try {
+ val response =
webTarget.path("healthz").request(MediaType.APPLICATION_JSON).get()
+ assert(HttpServletResponse.SC_SERVICE_UNAVAILABLE == response.getStatus)
+ val health = response.readEntity(classOf[HealthCheckResponse])
+ assert(!health.healthy)
+ assert(health.reason.contains("not registered"))
+ } finally {
+ worker.registered.set(true)
+ }
+ }
+
+ test("health reports unavailable when the master no longer knows the
worker") {
+ worker.registeredInMasterView.set(false)
+ try {
+ val response =
webTarget.path("healthz").request(MediaType.APPLICATION_JSON).get()
+ assert(HttpServletResponse.SC_SERVICE_UNAVAILABLE == response.getStatus)
+ val health = response.readEntity(classOf[HealthCheckResponse])
+ assert(!health.healthy)
+ assert(health.reason.contains("not registered"))
+ } finally {
+ worker.registeredInMasterView.set(true)
+ }
+ }
+
+ test("health reports unavailable when the worker state is not Normal") {
+ worker.workerStatusManager.transitionState(State.InDecommission)
+ try {
+ assert(worker.workerStatusManager.getWorkerState() ==
State.InDecommission)
+ val response =
webTarget.path("healthz").request(MediaType.APPLICATION_JSON).get()
+ assert(HttpServletResponse.SC_SERVICE_UNAVAILABLE == response.getStatus)
+ val health = response.readEntity(classOf[HealthCheckResponse])
+ assert(!health.healthy)
+ assert(health.reason.contains(State.InDecommission.toString))
+ } finally {
+ // InDecommission may only transition to Exit, so restore Normal
directly.
+ val normal = WorkerStatus.normalWorkerStatus()
+ worker.workerStatusManager.currentWorkerStatus = normal
+ worker.workerInfo.setWorkerStatus(normal)
+ }
+ }
+
test("listPartitionLocationInfo") {
val response =
webTarget.path("listPartitionLocationInfo").request(MediaType.TEXT_PLAIN).get()
assert(HttpServletResponse.SC_OK == response.getStatus)
