Add registry service readiness check (#3194)

Co-authored-by: Codex <codex@openai.com>

health/health.go

@@ -314,7 +314,8 @@ func DNSCheck(host string) CheckFunc {
 //
 // The lookup runs under the check's timeout, so an unreachable registry
 // (for example an etcd that has gone away) is reported as down rather
-// than blocking the probe.
+// than blocking the probe. Registries that honor ListContext also receive
+// the check context directly.
 func RegistryCheck(reg registry.Registry) CheckFunc {
 	return func(ctx context.Context) error {
 		if reg == nil {
@@ -322,7 +323,7 @@ func RegistryCheck(reg registry.Registry) CheckFunc {
 		}
 		errc := make(chan error, 1)
 		go func() {
-			_, err := reg.ListServices()
+			_, err := reg.ListServices(registry.ListContext(ctx))
 			errc <- err
 		}()
 		select {
@@ -337,6 +338,54 @@ func RegistryCheck(reg registry.Registry) CheckFunc {
 	}
 }
 
+// RegistryServiceCheck creates a check that verifies the local service
+// registration is visible in the registry. RegistryCheck only proves the
+// registry is reachable; this check also catches the common failure mode
+// where the process is alive and the registry answers, but this service's
+// node lease/record has disappeared and other services can no longer
+// discover it.
+//
+// Pass the service name and node ID used during registration. The default
+// RPC server node ID is service.Options().Name + "-" + service.Options().Id.
+func RegistryServiceCheck(reg registry.Registry, serviceName, nodeID string) CheckFunc {
+	return func(ctx context.Context) error {
+		if reg == nil {
+			return errors.New("no registry configured")
+		}
+		if serviceName == "" {
+			return errors.New("no service name configured")
+		}
+		if nodeID == "" {
+			return errors.New("no service node configured")
+		}
+
+		errc := make(chan error, 1)
+		go func() {
+			services, err := reg.GetService(serviceName, registry.GetContext(ctx))
+			if err != nil {
+				errc <- fmt.Errorf("registry %s lookup %s failed: %w", reg.String(), serviceName, err)
+				return
+			}
+			for _, service := range services {
+				for _, node := range service.Nodes {
+					if node.Id == nodeID {
+						errc <- nil
+						return
+					}
+				}
+			}
+			errc <- fmt.Errorf("registry %s missing node %s for service %s", reg.String(), nodeID, serviceName)
+		}()
+
+		select {
+		case err := <-errc:
+			return err
+		case <-ctx.Done():
+			return fmt.Errorf("registry %s service check timed out: %w", reg.String(), ctx.Err())
+		}
+	}
+}
+
 // CustomCheck creates a check from any function returning an error
 func CustomCheck(fn func() error) CheckFunc {
 	return func(ctx context.Context) error {

health/registry_check_test.go

@@ -89,3 +89,60 @@ func TestRegistryCheckMarksNotReady(t *testing.T) {
 		t.Error("service should be not-ready when the registry check is down")
 	}
 }
+
+func TestRegistryServiceCheckHealthy(t *testing.T) {
+	reg := registry.NewMemoryRegistry()
+	service := &registry.Service{
+		Name:  "orders",
+		Nodes: []*registry.Node{{Id: "orders-1"}},
+	}
+	if err := reg.Register(service); err != nil {
+		t.Fatalf("register service: %v", err)
+	}
+
+	check := RegistryServiceCheck(reg, "orders", "orders-1")
+	if err := check(context.Background()); err != nil {
+		t.Fatalf("registered service node should pass: %v", err)
+	}
+}
+
+func TestRegistryServiceCheckMissingNode(t *testing.T) {
+	reg := registry.NewMemoryRegistry()
+	service := &registry.Service{
+		Name:  "orders",
+		Nodes: []*registry.Node{{Id: "orders-1"}},
+	}
+	if err := reg.Register(service); err != nil {
+		t.Fatalf("register service: %v", err)
+	}
+
+	check := RegistryServiceCheck(reg, "orders", "orders-2")
+	err := check(context.Background())
+	if err == nil {
+		t.Fatal("missing service node should fail")
+	}
+	if !strings.Contains(err.Error(), "missing node orders-2") {
+		t.Errorf("error should describe the missing node: %v", err)
+	}
+}
+
+func TestRegistryServiceCheckMissingService(t *testing.T) {
+	check := RegistryServiceCheck(registry.NewMemoryRegistry(), "orders", "orders-1")
+	err := check(context.Background())
+	if err == nil {
+		t.Fatal("missing service should fail")
+	}
+	if !strings.Contains(err.Error(), registry.ErrNotFound.Error()) {
+		t.Errorf("error should include registry lookup failure: %v", err)
+	}
+}
+
+func TestRegistryServiceCheckMarksNotReady(t *testing.T) {
+	Reset()
+	defer Reset()
+
+	Register("registry-service", RegistryServiceCheck(registry.NewMemoryRegistry(), "orders", "orders-1"))
+	if IsReady(context.Background()) {
+		t.Error("service should be not-ready when its registry node is missing")
+	}
+}