Wire Broker listeners to healthchecks

pkg/agent/agent.go

@@ -280,18 +280,16 @@ func (a *Agent) Run(ctx context.Context) error {
 		storeService.Run,
 		catalog.ReconfigureTask(a.c.Log.WithField(telemetry.SubsystemName, "reconfigurer"), cat),
 	}
+	var apiReadyChannels []chan struct{}
 
 	if a.c.BindAddress != nil {
 		agentEndpoints := a.newEndpoints(metrics, mgr, workloadAttestor)
-		go func() {
-			agentEndpoints.WaitForListening(readyForHealthChecks)
-			a.started = true
-		}()
+		listening := make(chan struct{})
+		apiReadyChannels = append(apiReadyChannels, listening)
+		go agentEndpoints.WaitForListening(listening)
 		tasks = append(tasks, agentEndpoints.ListenAndServe)
 	} else {
 		a.c.Log.WithField("apis", "Workload and SDS APIs").Info("Skipping agent APIs because public endpoint is disabled")
-		a.started = true
-		close(readyForHealthChecks)
 	}
 
 	if a.c.AdminBindAddress != nil {
@@ -314,9 +312,20 @@ func (a *Agent) Run(ctx context.Context) error {
 		if err != nil {
 			return fmt.Errorf("failed to create broker endpoints: %w", err)
 		}
+		listening := make(chan struct{})
+		apiReadyChannels = append(apiReadyChannels, listening)
+		go brokerEndpoints.WaitForListening(listening)
 		tasks = append(tasks, brokerEndpoints.ListenAndServe)
 	}
 
+	go func() {
+		for _, listening := range apiReadyChannels {
+			<-listening
+		}
+		a.started = true
+		close(readyForHealthChecks)
+	}()
+
 	if a.c.LogReopener != nil {
 		tasks = append(tasks, a.c.LogReopener)
 	}

pkg/agent/broker/endpoints.go

@@ -86,6 +86,10 @@ type AllowedReferenceType struct {
 
 type Endpoints struct {
 	c *Config
+
+	hooks struct {
+		listening chan struct{}
+	}
 }
 
 func New(c *Config) (*Endpoints, error) {
@@ -107,9 +111,21 @@ func New(c *Config) (*Endpoints, error) {
 	}
 	return &Endpoints{
 		c: c,
+		hooks: struct {
+			listening chan struct{}
+		}{
+			listening: make(chan struct{}),
+		},
 	}, nil
 }
 
+// WaitForListening signals on listening after all configured Broker API
+// listeners have been created.
+func (e *Endpoints) WaitForListening(listening chan struct{}) {
+	<-e.hooks.listening
+	listening <- struct{}{}
+}
+
 func (e *Endpoints) ListenAndServe(ctx context.Context) error {
 	unaryInterceptor, streamInterceptor := middleware.Interceptors(
 		middleware.Chain(
@@ -178,7 +194,6 @@ func (e *Endpoints) ListenAndServe(ctx context.Context) error {
 			telemetry.Address: l.Addr().String(),
 		}).Info("Starting SPIFFE Broker Endpoint")
 	}
-
 	// Fan one gRPC server out across every listener with an errgroup. The
 	// first goroutine to error (or context cancellation) cancels the
 	// errgroup's context, which the watcher goroutine uses to call
@@ -199,6 +214,7 @@ func (e *Endpoints) ListenAndServe(ctx context.Context) error {
 		return nil
 	})
 
+	close(e.hooks.listening)
 	if err := g.Wait(); err != nil {
 		e.c.Log.WithError(err).Error("SPIFFE Broker Endpoint stopped prematurely")
 		return err