grpc: implement OnHealthy (#8686)

This allows service implementations to report their health as a
callback, rather than by polling. That, in turn, allows services to
become healthy more quickly after startup.

Author: jsha

docs/health.md

@@ -0,0 +1,20 @@
+# gRPC Health checking
+
+We automatically implement the [gRPC health service] automatically for all our
+gRPC servers.
+
+There are two ways a service implementation can offer health information:
+
+ - By implementing `Health(context.Context) error`, which will be called every
+   5s. If it returns `nil`, the service is set to healthy. If it returns
+   non-`nil`, the service is set to unhealthy. The health check interval can be
+   controlled with `grpc.serverBuilder.WithCheckInterval` at build time.
+ - By implementing `OnHealthy(func())`. This will be called by
+   `grpc.serverBuilder.Build`, passing in a closure that sets the service status
+   to healthy. This is useful for services that start unhealthy and then become
+   healthy exactly once. At a protocol layer, setting the service healthy [pushes]
+   out a message to clients immediately, so they don't need to wait on the next
+   health check poll.
+
+[gRPC health service]: https://pkg.go.dev/google.golang.org/grpc/health
+[pushes]: https://github.com/grpc/grpc/blob/5b6492ea90b2b867a6adad1b10a6edda28e860d1/src/proto/grpc/health/v1/health.proto#L47-L62

grpc/server.go

@@ -31,15 +31,21 @@ var CodedError = status.Errorf
 
 var errNilTLS = errors.New("boulder/grpc: received nil tls.Config")
 
-// checker is an interface for checking the health of a grpc service
-// implementation.
+// checker can be implemented by services to receive health checks via polling.
 type checker interface {
 	// Health returns nil if the service is healthy, or an error if it is not.
 	// If the passed context is canceled, it should return immediately with an
 	// error.
 	Health(context.Context) error
 }
 
+// onHealthy can be implemented by services to report health to gRPC by callbacks.
+type onHealthy interface {
+	// During server setup, OnHealthy will be passed a callback. The implementation
+	// should arrange for this callback to be invoked once the service is healthy.
+	OnHealthy(func())
+}
+
 // service represents a single gRPC service that can be registered with a gRPC
 // server.
 type service struct {
@@ -214,12 +220,17 @@ func (sb *serverBuilder) Build(tlsConfig *tls.Config, statsRegistry prometheus.R
 	healthCtx, stopHealthChecks := context.WithCancel(context.Background())
 	for _, s := range sb.services {
 		check, ok := s.impl.(checker)
-		if !ok {
-			continue
+		if ok {
+			sb.initLongRunningCheck(healthCtx, s.desc.ServiceName, check.Health)
 		}
-		sb.initLongRunningCheck(healthCtx, s.desc.ServiceName, check.Health)
-	}
 
+		registerReady, ok := s.impl.(onHealthy)
+		if ok {
+			registerReady.OnHealthy(func() {
+				sb.healthSrv.SetServingStatus(s.desc.ServiceName, healthpb.HealthCheckResponse_SERVING)
+			})
+		}
+	}
 	// Start a goroutine which listens for a termination signal, and then
 	// gracefully stops the gRPC server. This in turn causes the start() function
 	// to exit, allowing its caller (generally a main() function) to exit.

ra/ra.go

@@ -105,13 +105,11 @@ type RegistrationAuthorityImpl struct {
 
 var _ rapb.RegistrationAuthorityServer = (*RegistrationAuthorityImpl)(nil)
 
-// Health implements our grpc.checker interface. This method will be called
-// periodically to set the gRPC service's healthpb.Health.Check() status.
-func (ra *RegistrationAuthorityImpl) Health(ctx context.Context) error {
-	if ra.txnBuilder.Ready() {
-		return nil
-	}
-	return errors.New("waiting for overrides")
+// OnHealthy registers the callback to be invoked once the txnBuilder is healthy.
+//
+// That happens when the overrides are loaded.
+func (ra *RegistrationAuthorityImpl) OnHealthy(cb func()) {
+	ra.txnBuilder.OnHealthy(cb)
 }
 
 // NewRegistrationAuthorityImpl constructs a new RA object.

ratelimits/limit.go

@@ -322,6 +322,9 @@ type limitRegistry struct {
 	// overrides stores override limits by 'name:id'.
 	overrides Limits
 
+	// Called once the overrides are loaded.
+	healthyCallbacks []func()
+
 	// overridesLoaded is true if at least one loadOverrides attempt has
 	// completed successfully. Callers should check this using the Ready()
 	// method.
@@ -372,7 +375,14 @@ func (l *limitRegistry) loadOverrides(ctx context.Context) error {
 
 	l.Lock()
 	defer l.Unlock()
-	l.overridesLoaded = true
+
+	if !l.overridesLoaded {
+		l.overridesLoaded = true
+		for _, cb := range l.healthyCallbacks {
+			cb()
+			l.healthyCallbacks = nil
+		}
+	}
 
 	if len(newOverrides) < 1 {
 		// If it's an empty set, don't replace any current overrides.
@@ -403,6 +413,14 @@ func (l *limitRegistry) Ready() bool {
 	return l.overridesLoaded
 }
 
+// OnHealthy registers the callback to be invoked once the overrides are loaded.
+func (l *limitRegistry) OnHealthy(cb func()) {
+	l.Lock()
+	defer l.Unlock()
+
+	l.healthyCallbacks = append(l.healthyCallbacks, cb)
+}
+
 // loadOverridesWithRetry tries to loadOverrides, retrying at least every 30
 // seconds upon failure.
 func (l *limitRegistry) loadOverridesWithRetry(ctx context.Context) error {