Make HealthGrpc.HealthImplBase a @GRpcService compatible (#242)

Author: Alexander Furer

grpc-spring-boot-starter-demo/src/test/java/org/lognet/springboot/grpc/health/CustomGRpcHealthStatusManagerTest.java renamed to grpc-spring-boot-starter-demo/src/test/java/org/lognet/springboot/grpc/health/CustomManagedHealthStatusServiceTest.java

@@ -6,12 +6,13 @@
 import io.grpc.health.v1.HealthGrpc;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.lognet.springboot.grpc.GRpcService;
 import org.lognet.springboot.grpc.GrpcServerTestBase;
 import org.lognet.springboot.grpc.demo.DemoApp;
-import org.springframework.beans.factory.annotation.Autowired;
+import org.mockito.Mockito;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.boot.test.context.TestConfiguration;
-import org.springframework.context.annotation.Bean;
+import org.springframework.boot.test.mock.mockito.SpyBean;
 import org.springframework.test.context.ActiveProfiles;
 import org.springframework.test.context.junit4.SpringRunner;
 
@@ -23,24 +24,22 @@
 import static org.springframework.boot.test.context.SpringBootTest.WebEnvironment.NONE;
 
 @RunWith(SpringRunner.class)
-@SpringBootTest(classes = {DemoApp.class, CustomGRpcHealthStatusManagerTest.Cfg.class}, webEnvironment = NONE)
+@SpringBootTest(classes = {DemoApp.class, CustomManagedHealthStatusServiceTest.Cfg.class}, webEnvironment = NONE)
 @ActiveProfiles("disable-security")
-public class CustomGRpcHealthStatusManagerTest extends GrpcServerTestBase {
+public class CustomManagedHealthStatusServiceTest extends GrpcServerTestBase {
     @TestConfiguration
     static class Cfg{
-        static class MyCustomHealthStatusManager extends DefaultHealthStatusManager{}
-        @Bean
-        public GRpcHealthStatusManager myCustom(){
-            return new MyCustomHealthStatusManager();
-        }
+        @GRpcService
+        static class MyCustomHealthStatusService extends DefaultHealthStatusService {}
+
     }
 
-    @Autowired
-    private  GRpcHealthStatusManager healthStatusManager;
+    @SpyBean
+    private ManagedHealthStatusService healthStatusManager;
 
     @Test
     public void contextLoads() {
-        assertThat(healthStatusManager,isA(Cfg.MyCustomHealthStatusManager.class));
+        assertThat(healthStatusManager,isA(Cfg.MyCustomHealthStatusService.class));
     }
 
     @Test
@@ -50,5 +49,8 @@ public void testHealthCheck() throws ExecutionException, InterruptedException {
         final HealthCheckResponse.ServingStatus servingStatus = healthFutureStub.check(healthCheckRequest).get().getStatus();
 
         assertThat(servingStatus, is(HealthCheckResponse.ServingStatus.SERVING));
+
+        Mockito.verify(healthStatusManager,Mockito.atLeast(1))
+                .setStatus(Mockito.any(String.class),Mockito.eq(HealthCheckResponse.ServingStatus.SERVING));
     }
 }

grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/GRpcServerRunner.java

@@ -7,11 +7,13 @@
 import io.grpc.ServerInterceptors;
 import io.grpc.ServerServiceDefinition;
 import io.grpc.health.v1.HealthCheckResponse;
+import io.grpc.health.v1.HealthGrpc;
 import io.grpc.protobuf.services.ProtoReflectionService;
 import lombok.extern.slf4j.Slf4j;
 import org.lognet.springboot.grpc.autoconfigure.GRpcServerProperties;
 import org.lognet.springboot.grpc.context.GRpcServerInitializedEvent;
-import org.lognet.springboot.grpc.health.GRpcHealthStatusManager;
+import org.lognet.springboot.grpc.health.ManagedHealthStatusService;
+import org.springframework.beans.FatalBeanException;
 import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.config.BeanDefinition;
@@ -42,8 +44,8 @@
 public class GRpcServerRunner implements SmartLifecycle {
 
     private final AtomicBoolean isRunning = new AtomicBoolean(false);
-    @Autowired
-    private Optional<GRpcHealthStatusManager> healthStatusManager;
+
+    private Optional<ManagedHealthStatusService> healthStatusManager = Optional.empty();
 
     @Autowired
     private AbstractApplicationContext applicationContext;
@@ -77,28 +79,39 @@ public void start() {
                     .map(name -> applicationContext.getBeanFactory().getBean(name, ServerInterceptor.class))
                     .collect(Collectors.toList());
 
-            // Adding health service
-            healthStatusManager
-                    .map(GRpcHealthStatusManager::getHealthService)
-                    .ifPresent(serverBuilder::addService);
 
             // find and register all GRpcService-enabled beans
+            List<String> serviceNames = new ArrayList<>();
             getBeanNamesByTypeWithAnnotation(GRpcService.class, BindableService.class)
                     .forEach(name -> {
                         BindableService srv = applicationContext.getBeanFactory().getBean(name, BindableService.class);
                         ServerServiceDefinition serviceDefinition = srv.bindService();
                         GRpcService gRpcServiceAnn = applicationContext.findAnnotationOnBean(name, GRpcService.class);
                         serviceDefinition = bindInterceptors(serviceDefinition, gRpcServiceAnn, globalInterceptors);
                         serverBuilder.addService(serviceDefinition);
-                        String serviceName = serviceDefinition.getServiceDescriptor().getName();
-                        healthStatusManager.ifPresent(m ->
-                                m.setStatus(serviceName, HealthCheckResponse.ServingStatus.SERVING)
-                        );
+
+                        if (srv instanceof HealthGrpc.HealthImplBase) {
+                            if(!(srv instanceof ManagedHealthStatusService)){
+                                throw new FatalBeanException(String.format("Please inherit %s from %s rather than directly from %s",
+                                        srv.getClass().getName(),
+                                        ManagedHealthStatusService.class.getName(),
+                                        HealthGrpc.HealthImplBase.class.getName()
+                                        ));
+                            }
+                            healthStatusManager = Optional.of((ManagedHealthStatusService)srv);
+                        } else {
+                            serviceNames.add(serviceDefinition.getServiceDescriptor().getName());
+                        }
+
 
                         log.info("'{}' service has been registered.", srv.getClass().getName());
 
                     });
 
+            healthStatusManager.ifPresent(h ->
+                    serviceNames.forEach(serviceName -> h.setStatus(serviceName, HealthCheckResponse.ServingStatus.SERVING))
+            );
+
             if (gRpcServerProperties.isEnableReflection()) {
                 serverBuilder.addService(ProtoReflectionService.newInstance());
                 log.info("'{}' service has been registered.", ProtoReflectionService.class.getName());
@@ -180,7 +193,7 @@ private void startDaemonAwaitThread() {
     public void stop() {
         Optional.ofNullable(server).ifPresent(s -> {
             log.info("Shutting down gRPC server ...");
-            healthStatusManager.ifPresent(GRpcHealthStatusManager::onShutdown);
+            healthStatusManager.ifPresent(ManagedHealthStatusService::onShutdown);
 
             s.shutdown();
             int shutdownGrace = gRpcServerProperties.getShutdownGrace();

grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/autoconfigure/GRpcAutoConfiguration.java

@@ -5,8 +5,8 @@
 import org.lognet.springboot.grpc.GRpcServerBuilderConfigurer;
 import org.lognet.springboot.grpc.GRpcServerRunner;
 import org.lognet.springboot.grpc.GRpcService;
-import org.lognet.springboot.grpc.health.DefaultHealthStatusManager;
-import org.lognet.springboot.grpc.health.GRpcHealthStatusManager;
+import org.lognet.springboot.grpc.health.DefaultHealthStatusService;
+import org.lognet.springboot.grpc.health.ManagedHealthStatusService;
 import org.springframework.beans.factory.BeanCreationException;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
@@ -60,8 +60,8 @@ public GRpcServerRunner grpcInprocessServerRunner(@Qualifier("grpcInternalConfig
 
     @Bean
     @ConditionalOnMissingBean
-    public GRpcHealthStatusManager  healthStatusManager() {
-        return new DefaultHealthStatusManager();
+    public ManagedHealthStatusService healthStatusManager() {
+        return new DefaultHealthStatusService();
     }
 
     @Bean

grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/health/DefaultHealthStatusManager.java

@@ -0,0 +1,32 @@
+package org.lognet.springboot.grpc.health;
+
+import io.grpc.health.v1.HealthCheckRequest;
+import io.grpc.health.v1.HealthCheckResponse;
+import io.grpc.health.v1.HealthGrpc;
+import io.grpc.protobuf.services.HealthStatusManager;
+import io.grpc.stub.StreamObserver;
+
+public class DefaultHealthStatusService extends ManagedHealthStatusService {
+    private final HealthStatusManager healthStatusManager = new HealthStatusManager();
+    private final HealthGrpc.HealthImplBase service = (HealthGrpc.HealthImplBase) healthStatusManager.getHealthService();
+
+    @Override
+    public void  onShutdown() {
+         healthStatusManager.enterTerminalState();
+    }
+
+    @Override
+    public void setStatus(String service, HealthCheckResponse.ServingStatus status) {
+        healthStatusManager.setStatus(service,status);
+    }
+
+    @Override
+    public void check(HealthCheckRequest request, StreamObserver<HealthCheckResponse> responseObserver) {
+        service.check(request, responseObserver);
+    }
+
+    @Override
+    public void watch(HealthCheckRequest request, StreamObserver<HealthCheckResponse> responseObserver) {
+        service.watch(request, responseObserver);
+    }
+}

grpc-spring-boot-starter/src/main/java/org/lognet/springboot/grpc/health/DefaultHealthStatusService.java

@@ -3,21 +3,20 @@
 import io.grpc.health.v1.HealthCheckResponse;
 import io.grpc.health.v1.HealthGrpc;
 
-public interface GRpcHealthStatusManager {
+public abstract class ManagedHealthStatusService extends HealthGrpc.HealthImplBase{
 
     /**
      * Invoked on server shutdown. Implementation is advised to set status of all services as ServingStatus.NOT_SERVING
      */
-    void onShutdown();
+    public abstract void onShutdown();
 
     /**
      * Invoked on startup with {@link io.grpc.health.v1.HealthCheckResponse.ServingStatus#SERVING } for each discovered grpc service name
      * @param service - grpc service name
      * @param status - new status
      */
-    void setStatus(String service, HealthCheckResponse.ServingStatus status);
+    public abstract  void setStatus(String service, HealthCheckResponse.ServingStatus status);
 
 
-    HealthGrpc.HealthImplBase getHealthService() ;
 
 }