Add GrpcHealthCheckedEndpointGroupBuilder

grpc/src/main/java/com/linecorp/armeria/client/grpc/endpoint/healthcheck/GrpcHealthCheckMethod.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2025 LY Corporation
+ *
+ * LY Corporation 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:
+ *
+ *   https://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 com.linecorp.armeria.client.grpc.endpoint.healthcheck;
+
+/**
+ * Represents a gRPC health check method.
+ */
+public enum GrpcHealthCheckMethod {
+    CHECK, WATCH
+}

grpc/src/main/java/com/linecorp/armeria/client/grpc/endpoint/healthcheck/GrpcHealthCheckedEndpointGroupBuilder.java

@@ -0,0 +1,95 @@
+/*
+ * Copyright 2025 LY Corporation
+ *
+ * LY Corporation 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:
+ *
+ *   https://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 com.linecorp.armeria.client.grpc.endpoint.healthcheck;
+
+import static java.util.Objects.requireNonNull;
+
+import java.util.function.Function;
+
+import com.linecorp.armeria.client.endpoint.EndpointGroup;
+import com.linecorp.armeria.client.endpoint.healthcheck.AbstractHealthCheckedEndpointGroupBuilder;
+import com.linecorp.armeria.client.endpoint.healthcheck.HealthCheckerContext;
+import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.util.AsyncCloseable;
+import com.linecorp.armeria.internal.client.grpc.GrpcHealthCheckWatcher;
+import com.linecorp.armeria.internal.client.grpc.GrpcHealthChecker;
+
+/**
+ * Builds a health checked endpoint group whose health comes from a standard gRPC health check service.
+ */
+public final class GrpcHealthCheckedEndpointGroupBuilder
+        extends AbstractHealthCheckedEndpointGroupBuilder<GrpcHealthCheckedEndpointGroupBuilder> {
+
+    private @Nullable String service;
+    private final GrpcHealthCheckMethod healthCheckMethod;
+
+    GrpcHealthCheckedEndpointGroupBuilder(EndpointGroup delegate, GrpcHealthCheckMethod healthCheckMethod) {
+        super(delegate);
+        this.healthCheckMethod = healthCheckMethod;
+    }
+
+    /**
+     * Returns a {@link GrpcHealthCheckedEndpointGroupBuilder} that builds a health checked
+     * endpoint group with the specified {@link EndpointGroup} and {@link GrpcHealthCheckMethod}.
+     */
+    public static GrpcHealthCheckedEndpointGroupBuilder builder(EndpointGroup delegate,
+                                                                GrpcHealthCheckMethod healthCheckMethod) {
+        return new GrpcHealthCheckedEndpointGroupBuilder(requireNonNull(delegate),
+                requireNonNull(healthCheckMethod));
+    }
+
+    /**
+     * Sets the optional service field of the gRPC health check request.
+     */
+    public GrpcHealthCheckedEndpointGroupBuilder service(@Nullable String service) {
+        this.service = service;
+        return this;
+    }
+
+    @Override
+    protected Function<? super HealthCheckerContext, ? extends AsyncCloseable> newCheckerFactory() {
+        return new GrpcHealthCheckerFactory(service, healthCheckMethod);
+    }
+
+    private static final class GrpcHealthCheckerFactory
+            implements Function<HealthCheckerContext, AsyncCloseable> {
+
+        private final @Nullable String service;
+        private final GrpcHealthCheckMethod healthCheckMethod;
+
+        private GrpcHealthCheckerFactory(@Nullable String service, GrpcHealthCheckMethod healthCheckMethod) {
+            this.service = service;
+            this.healthCheckMethod = healthCheckMethod;
+        }
+
+        @Override
+        public AsyncCloseable apply(HealthCheckerContext ctx) {
+            if (healthCheckMethod == GrpcHealthCheckMethod.CHECK) {
+                final GrpcHealthChecker healthChecker = new GrpcHealthChecker(ctx, ctx.endpoint(),
+                        ctx.protocol(), service);
+                healthChecker.start();
+                return healthChecker;
+            } else if (healthCheckMethod == GrpcHealthCheckMethod.WATCH) {
+                final GrpcHealthCheckWatcher healthChecker = new GrpcHealthCheckWatcher(ctx, ctx.endpoint(),
+                        ctx.protocol(), service);
+                healthChecker.start();
+                return healthChecker;
+            }
+            // should not get here
+            throw new IllegalArgumentException("Invalid health check method");
+        }
+    }
+}

grpc/src/main/java/com/linecorp/armeria/client/grpc/endpoint/healthcheck/package-info.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright 2025 LY Corporation
+ *
+ * LY Corporation 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:
+ *
+ *   https://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.
+ */
+
+/**
+ * gRPC health checked endpoint.
+ */
+@NonNullByDefault
+package com.linecorp.armeria.client.grpc.endpoint.healthcheck;
+
+import com.linecorp.armeria.common.annotation.NonNullByDefault;

grpc/src/main/java/com/linecorp/armeria/internal/client/grpc/AbstractGrpcHealthChecker.java

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2025 LY Corporation
+ *
+ * LY Corporation 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:
+ *
+ *   https://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 com.linecorp.armeria.internal.client.grpc;
+
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.locks.ReentrantLock;
+
+import com.linecorp.armeria.common.util.AsyncCloseable;
+import com.linecorp.armeria.common.util.AsyncCloseableSupport;
+import com.linecorp.armeria.internal.common.util.ReentrantShortLock;
+
+/**
+ * Abstract class that provides common structure for {@link GrpcHealthChecker} and
+ * {@link GrpcHealthCheckWatcher}.
+ */
+abstract class AbstractGrpcHealthChecker implements AsyncCloseable {
+
+    static final double HEALTHY = 1d;
+    static final double UNHEALTHY = 0d;
+
+    private final ReentrantLock lock = new ReentrantShortLock();
+    private final AsyncCloseableSupport closeable = AsyncCloseableSupport.of(this::closeAsync);
+
+    public void start() {
+        check();
+    }
+
+    protected abstract void check();
+
+    @Override
+    public CompletableFuture<?> closeAsync() {
+        return closeable.closeAsync();
+    }
+
+    private synchronized void closeAsync(CompletableFuture<?> future) {
+        future.complete(null);
+    }
+
+    @Override
+    public void close() {
+        closeable.close();
+    }
+
+    protected void lock() {
+        lock.lock();
+    }
+
+    protected void unlock() {
+        lock.unlock();
+    }
+}

grpc/src/main/java/com/linecorp/armeria/internal/client/grpc/GrpcHealthCheckWatcher.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2025 LY Corporation
+ *
+ * LY Corporation 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:
+ *
+ *   https://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 com.linecorp.armeria.internal.client.grpc;
+
+import java.time.Duration;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.client.ClientRequestContextCaptor;
+import com.linecorp.armeria.client.Clients;
+import com.linecorp.armeria.client.Endpoint;
+import com.linecorp.armeria.client.endpoint.healthcheck.HealthCheckerContext;
+import com.linecorp.armeria.client.grpc.GrpcClients;
+import com.linecorp.armeria.common.ResponseHeaders;
+import com.linecorp.armeria.common.SessionProtocol;
+import com.linecorp.armeria.common.annotation.Nullable;
+import com.linecorp.armeria.common.logging.RequestLogProperty;
+
+import io.grpc.health.v1.HealthCheckRequest;
+import io.grpc.health.v1.HealthCheckResponse;
+import io.grpc.health.v1.HealthGrpc;
+import io.grpc.stub.StreamObserver;
+
+/**
+ * Performs gRPC health checking using the Watch rpc endpoint.
+ */
+public class GrpcHealthCheckWatcher extends AbstractGrpcHealthChecker {
+
+    private static final Logger LOGGER = LoggerFactory.getLogger(GrpcHealthCheckWatcher.class);
+
+    private final HealthCheckerContext ctx;
+    @Nullable private final String service;
+    private final HealthGrpc.HealthStub stub;
+
+    public GrpcHealthCheckWatcher(HealthCheckerContext ctx, Endpoint endpoint, SessionProtocol sessionProtocol,
+                                  @Nullable String service) {
+        this.ctx = ctx;
+        this.service = service;
+
+        this.stub = GrpcClients.builder(sessionProtocol, endpoint)
+                .options(ctx.clientOptions())
+                .responseTimeout(Duration.ZERO) // disable timeout for streaming watch rpc
+                .build(HealthGrpc.HealthStub.class);
+    }
+
+    @Override
+    protected void check() {
+        lock();
+        try {
+            final HealthCheckRequest.Builder builder = HealthCheckRequest.newBuilder();
+            if (service != null) {
+                builder.setService(service);
+            }
+
+            try (ClientRequestContextCaptor reqCtxCaptor = Clients.newContextCaptor()) {
+                stub.watch(builder.build(), new StreamObserver<HealthCheckResponse>() {
+                    @Override
+                    public void onNext(HealthCheckResponse healthCheckResponse) {
+                        final ClientRequestContext reqCtx = reqCtxCaptor.get();
+                        // extract the headers from the ctx log
+                        ResponseHeaders responseHeaders = null;
+                        if (reqCtx.log().isAvailable(RequestLogProperty.RESPONSE_HEADERS)) {
+                            responseHeaders = reqCtx.log().partial().responseHeaders();
+                        }
+                        // update health
+                        if (healthCheckResponse.getStatus() == HealthCheckResponse.ServingStatus.SERVING) {
+                            LOGGER.debug("Health check returned healthy from endpoint {}",
+                                    ctx.endpoint());
+                            ctx.updateHealth(HEALTHY, reqCtx, responseHeaders, null);
+                        } else {
+                            LOGGER.debug("Health check returned unhealthy from endpoint {}",
+                                    ctx.endpoint());
+                            ctx.updateHealth(UNHEALTHY, reqCtx, responseHeaders, null);
+                        }
+                    }
+
+                    @Override
+                    public void onError(Throwable throwable) {
+                        final ClientRequestContext reqCtx = reqCtxCaptor.get();
+                        // extract the headers from the ctx log
+                        ResponseHeaders responseHeaders = null;
+                        if (reqCtx.log().isAvailable(RequestLogProperty.RESPONSE_HEADERS)) {
+                            responseHeaders = reqCtx.log().partial().responseHeaders();
+                        }
+                        // update health
+                        LOGGER.debug("Failed streaming health check on endpoint {}", ctx.endpoint(), throwable);
+                        ctx.updateHealth(UNHEALTHY, reqCtx, responseHeaders, throwable);
+
+                        // schedule next watch request
+                        ctx.executor().execute(GrpcHealthCheckWatcher.this::check);
+                    }
+
+                    @Override
+                    public void onCompleted() {
+                        final ClientRequestContext reqCtx = reqCtxCaptor.get();
+                        // update health
+                        LOGGER.debug("Streaming health check complete from endpoint {}", ctx.endpoint());
+                        ctx.updateHealth(UNHEALTHY, reqCtx, null, null);
+
+                        // schedule next watch request
+                        ctx.executor().execute(GrpcHealthCheckWatcher.this::check);
+                    }
+                });
+            }
+        } finally {
+            unlock();
+        }
+    }
+}