Repository: MystenLabs/sui
File: crates/sui-indexer-alt-consistent-store/src/rpc/mod.rs (Rust)

Code before:

        let reflection_v1alpha = reflection_v1alpha
            .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
            .build_v1alpha()
            .unwrap();

        let (health_reporter, health_service) = tonic_health::server::health_reporter();

        service_names.extend([
            service_name(&reflection_v1),
            service_name(&reflection_v1alpha),
            service_name(&health_service),
        ]);

        router = add_service(router, reflection_v1);
        router = add_service(router, reflection_v1alpha);
        router = add_service(router, health_service);
        router = router
            .layer(CallbackLayer::new(MakeMetricsHandler::new(metrics.clone())))
            .layer(axum::middleware::from_fn_with_state(
                Version(version),
                middleware::version::set_version,
            ))
            .layer(CatchPanicLayer::new(metrics));

        for service_name in &service_names {
            health_reporter
                .set_service_status(service_name, ServingStatus::NotServing)
                .await;
        }
        // Set empty string service to NotServing (this is what the health probe will check)
        health_reporter
            .set_service_status("", ServingStatus::NotServing)
            .await;

        // Flip to Serving once snapshots exist
        {
            let reporter = health_reporter.clone();
            tokio::spawn(async move {
                loop {
                    if db.snapshot_range(u64::MAX).is_some() {
                        reporter
                            .set_service_status("", ServingStatus::Serving)
                            .await;
                        for name in &service_names {
                            reporter
                                .set_service_status(*name, ServingStatus::Serving)
                                .await;
                        }
                        reporter
                            .set_service_status("", ServingStatus::Serving)

Review comment (suggestion):
I don't think this is technically right -- the health check is itself a service, and by doing this, we are saying that it is not ready until the Db has a snapshot, which is not true. Similar for the introspection service.

The right way to do this is to supply a `ready` future with each service when we add it, and then here, we spawn a separate task for each (user) service, that awaits the `ready` service and then sets it as serving.

Then we can add this health check just for the consistent service, and the health probe should check the health of the consistent service.

Code after:

        let reflection_v1alpha = reflection_v1alpha
            .register_encoded_file_descriptor_set(tonic_health::pb::FILE_DESCRIPTOR_SET)
            .build_v1alpha()
            .unwrap();

        let (health_reporter, health_service) = tonic_health::server::health_reporter();

        let internal_services = vec![
            service_name(&reflection_v1),
            service_name(&reflection_v1alpha),
            service_name(&health_service),
        ];

        router = add_service(router, reflection_v1);
        router = add_service(router, reflection_v1alpha);
        router = add_service(router, health_service);
        router = router
            .layer(CallbackLayer::new(MakeMetricsHandler::new(metrics.clone())))
            .layer(axum::middleware::from_fn_with_state(
                Version(version),
                middleware::version::set_version,
            ))
            .layer(CatchPanicLayer::new(metrics));

        for service_name in internal_services {
            health_reporter
                .set_service_status(service_name, ServingStatus::Serving)
                .await;
        }

        // Create a Service to be attached as secondary to the main service
        let mut readiness_checks = Service::new();

        for (name, ready) in service_futures {
            health_reporter
                .set_service_status(name, ServingStatus::NotServing)
                .await;

            let reporter = health_reporter.clone();
            readiness_checks = readiness_checks.spawn(async move {
                ready.await;
                reporter
                    .set_service_status(name, ServingStatus::Serving)
                    .await;
                info!("gRPC service {name} is now SERVING");
                Ok(())
            });
        }

        let mut service = Service::new();