feat(zapwire): WaitForHealthy over ZAP

zapwire/client.go

@@ -212,6 +212,16 @@ func (c *Client) SendOutboundMessage(ctx context.Context, req *types.SendOutboun
 	return resp, nil
 }
 
+// WaitForHealthy blocks server-side until the cluster reports healthy
+// on every node, then returns the snapshot.
+func (c *Client) WaitForHealthy(ctx context.Context) (*types.WaitForHealthyResponse, error) {
+	resp := &types.WaitForHealthyResponse{}
+	if err := c.callSub(ctx, OpStart, SubWaitForHealthy, &types.WaitForHealthyRequest{}, resp); err != nil {
+		return nil, err
+	}
+	return resp, nil
+}
+
 // Stop terminates the cluster and returns the final snapshot.
 func (c *Client) Stop(ctx context.Context) (*types.StopResponse, error) {
 	resp := &types.StopResponse{}

zapwire/e2e_test.go

@@ -213,6 +213,19 @@ func (b *stubBackend) SendOutboundMessage(ctx context.Context, req *types.SendOu
 	return &types.SendOutboundMessageResponse{Sent: true}, nil
 }
 
+func (b *stubBackend) WaitForHealthy(ctx context.Context) (*types.WaitForHealthyResponse, error) {
+	return &types.WaitForHealthyResponse{
+		ClusterInfo: &types.ClusterInfo{
+			NodeNames:   []string{"alpha", "beta"},
+			NodeInfos:   map[string]*types.NodeInfo{},
+			PID:         42,
+			RootDataDir: "/tmp/netrunner",
+			Healthy:     true,
+			NetworkID:   12345,
+		},
+	}, nil
+}
+
 func (b *stubBackend) Stop(ctx context.Context) (*types.StopResponse, error) {
 	return &types.StopResponse{
 		ClusterInfo: &types.ClusterInfo{
@@ -591,6 +604,33 @@ func TestE2ESendOutboundMessage(t *testing.T) {
 	}
 }
 
+func TestE2EWaitForHealthy(t *testing.T) {
+	srv, teardown := runServer(t, &stubBackend{})
+	defer teardown()
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	client, err := Dial(ctx, srv.Addr())
+	if err != nil {
+		t.Fatalf("Dial: %v", err)
+	}
+	defer client.Close()
+
+	resp, err := client.WaitForHealthy(ctx)
+	if err != nil {
+		t.Fatalf("WaitForHealthy: %v", err)
+	}
+	if resp.ClusterInfo == nil {
+		t.Fatal("ClusterInfo: nil")
+	}
+	if !resp.ClusterInfo.Healthy {
+		t.Fatal("Healthy: false")
+	}
+	if len(resp.ClusterInfo.NodeNames) != 2 {
+		t.Fatalf("NodeNames: got %d, want 2", len(resp.ClusterInfo.NodeNames))
+	}
+}
+
 func TestE2EStop(t *testing.T) {
 	srv, teardown := runServer(t, &stubBackend{})
 	defer teardown()

zapwire/server.go

@@ -30,6 +30,7 @@ type Backend interface {
 	ResumeNode(ctx context.Context, req *types.ResumeNodeRequest) (*types.ResumeNodeResponse, error)
 	AttachPeer(ctx context.Context, req *types.AttachPeerRequest) (*types.AttachPeerResponse, error)
 	SendOutboundMessage(ctx context.Context, req *types.SendOutboundMessageRequest) (*types.SendOutboundMessageResponse, error)
+	WaitForHealthy(ctx context.Context) (*types.WaitForHealthyResponse, error)
 }
 
 // Server hosts a netrunner control RPC over ZAP.
@@ -216,6 +217,12 @@ func (s *Server) handleStartSub(
 			return 0, nil, err
 		}
 		return msgType, encodeResp(resp), nil
+	case SubWaitForHealthy:
+		resp, err := s.be.WaitForHealthy(ctx)
+		if err != nil {
+			return 0, nil, err
+		}
+		return msgType, encodeResp(resp), nil
 	default:
 		return 0, nil, fmt.Errorf("zapwire: unknown SubOp 0x%02x for OpStart", uint8(sub))
 	}

zapwire/types/types.go

@@ -645,6 +645,43 @@ func (s *SendOutboundMessageResponse) Decode(rd *zap.Reader) error {
 	return nil
 }
 
+// WaitForHealthyRequest is empty (proto has an optional network_name
+// field used only by the multi-network REST gateway shim; the native
+// control protocol always operates on the single in-process network,
+// so this field is omitted).
+type WaitForHealthyRequest struct{}
+
+func (*WaitForHealthyRequest) Encode(b *zap.Buffer)       {}
+func (*WaitForHealthyRequest) Decode(r *zap.Reader) error { return nil }
+
+// WaitForHealthyResponse carries the cluster snapshot once all nodes
+// report Healthy = true.
+type WaitForHealthyResponse struct {
+	ClusterInfo *ClusterInfo
+}
+
+func (w *WaitForHealthyResponse) Encode(b *zap.Buffer) {
+	if w.ClusterInfo == nil {
+		b.WriteUint8(0)
+		return
+	}
+	b.WriteUint8(1)
+	w.ClusterInfo.Encode(b)
+}
+
+func (w *WaitForHealthyResponse) Decode(r *zap.Reader) error {
+	present, err := r.ReadUint8()
+	if err != nil {
+		return err
+	}
+	if present == 0 {
+		w.ClusterInfo = nil
+		return nil
+	}
+	w.ClusterInfo = &ClusterInfo{}
+	return w.ClusterInfo.Decode(r)
+}
+
 // StopRequest is empty.
 type StopRequest struct{}