Fix runtime review findings

clients/go/mxgateway/client.go

@@ -16,6 +16,7 @@ import (
 const (
 	defaultDialTimeout = 10 * time.Second
 	defaultCallTimeout = 30 * time.Second
+	defaultMaxGrpcMessageBytes = 16 * 1024 * 1024
 )
 
 // Client owns a gateway gRPC connection and exposes session-oriented helpers.
@@ -50,6 +51,10 @@ func Dial(ctx context.Context, opts Options) (*Client, error) {
 		grpc.WithTransportCredentials(transportCredentials),
 		grpc.WithUnaryInterceptor(unaryAuthInterceptor(opts.APIKey)),
 		grpc.WithStreamInterceptor(streamAuthInterceptor(opts.APIKey)),
+		grpc.WithDefaultCallOptions(
+			grpc.MaxCallRecvMsgSize(resolveMaxGrpcMessageBytes(opts)),
+			grpc.MaxCallSendMsgSize(resolveMaxGrpcMessageBytes(opts)),
+		),
 		grpc.WithBlock(),
 	}
 	dialOptions = append(dialOptions, opts.DialOptions...)
@@ -62,6 +67,13 @@ func Dial(ctx context.Context, opts Options) (*Client, error) {
 	return NewClient(conn, opts), nil
 }
 
+func resolveMaxGrpcMessageBytes(opts Options) int {
+	if opts.MaxGrpcMessageBytes > 0 {
+		return opts.MaxGrpcMessageBytes
+	}
+	return defaultMaxGrpcMessageBytes
+}
+
 // NewClient wraps an existing gRPC connection. The caller owns closing conn
 // unless it calls Close on the returned Client.
 func NewClient(conn *grpc.ClientConn, opts Options) *Client {

clients/go/mxgateway/galaxy.go

@@ -13,6 +13,8 @@ import (
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
+const discoverHierarchyPageSize = 5000
+
 // RawGalaxyRepositoryClient is the generated gRPC client interface for the
 // Galaxy Repository service exposed for callers that need direct contract
 // access.
@@ -70,6 +72,10 @@ func DialGalaxy(ctx context.Context, opts Options) (*GalaxyClient, error) {
 		grpc.WithTransportCredentials(transportCredentials),
 		grpc.WithUnaryInterceptor(unaryAuthInterceptor(opts.APIKey)),
 		grpc.WithStreamInterceptor(streamAuthInterceptor(opts.APIKey)),
+		grpc.WithDefaultCallOptions(
+			grpc.MaxCallRecvMsgSize(resolveMaxGrpcMessageBytes(opts)),
+			grpc.MaxCallSendMsgSize(resolveMaxGrpcMessageBytes(opts)),
+		),
 		grpc.WithBlock(),
 	}
 	dialOptions = append(dialOptions, opts.DialOptions...)
@@ -141,11 +147,23 @@ func (c *GalaxyClient) DiscoverHierarchy(ctx context.Context) ([]*GalaxyObject,
 	callCtx, cancel := c.callContext(ctx)
 	defer cancel()
 
-	reply, err := c.raw.DiscoverHierarchy(callCtx, &pb.DiscoverHierarchyRequest{})
-	if err != nil {
-		return nil, &GatewayError{Op: "galaxy discover hierarchy", Err: err}
+	var objects []*GalaxyObject
+	pageToken := ""
+	for {
+		reply, err := c.raw.DiscoverHierarchy(callCtx, &pb.DiscoverHierarchyRequest{
+			PageSize:  discoverHierarchyPageSize,
+			PageToken: pageToken,
+		})
+		if err != nil {
+			return nil, &GatewayError{Op: "galaxy discover hierarchy", Err: err}
+		}
+		objects = append(objects, reply.GetObjects()...)
+		pageToken = reply.GetNextPageToken()
+		if pageToken == "" {
+			break
+		}
 	}
-	return reply.GetObjects(), nil
+	return objects, nil
 }
 
 // WatchDeployEventsRaw starts the generated WatchDeployEvents stream for callers

clients/go/mxgateway/galaxy_test.go

@@ -95,7 +95,9 @@ func TestGalaxyGetLastDeployTimeReturnsAbsentWhenTimestampNil(t *testing.T) {
 
 func TestGalaxyDiscoverHierarchyReturnsObjects(t *testing.T) {
 	fake := &fakeGalaxyServer{
-		discoverReply: &pb.DiscoverHierarchyReply{
+		discoverReplies: []*pb.DiscoverHierarchyReply{{
+			NextPageToken:    "page-2",
+			TotalObjectCount: 2,
 			Objects: []*pb.GalaxyObject{
 				{
 					GobjectId:     1,
@@ -114,6 +116,10 @@ func TestGalaxyDiscoverHierarchyReturnsObjects(t *testing.T) {
 						},
 					},
 				},
+			},
+		}, {
+			TotalObjectCount: 2,
+			Objects: []*pb.GalaxyObject{
 				{
 					GobjectId:       2,
 					TagName:         "TestMachine_002",
@@ -121,7 +127,7 @@ func TestGalaxyDiscoverHierarchyReturnsObjects(t *testing.T) {
 					ParentGobjectId: 1,
 				},
 			},
-		},
+		}},
 	}
 	client, cleanup := newGalaxyBufconnClient(t, fake)
 	defer cleanup()
@@ -133,6 +139,15 @@ func TestGalaxyDiscoverHierarchyReturnsObjects(t *testing.T) {
 	if len(objects) != 2 {
 		t.Fatalf("len(objects) = %d, want 2", len(objects))
 	}
+	if len(fake.discoverRequests) != 2 {
+		t.Fatalf("len(discoverRequests) = %d, want 2", len(fake.discoverRequests))
+	}
+	if fake.discoverRequests[0].GetPageSize() != 5000 || fake.discoverRequests[0].GetPageToken() != "" {
+		t.Fatalf("first request = %+v", fake.discoverRequests[0])
+	}
+	if fake.discoverRequests[1].GetPageToken() != "page-2" {
+		t.Fatalf("second page_token = %q, want page-2", fake.discoverRequests[1].GetPageToken())
+	}
 	if objects[0].GetTagName() != "TestMachine_001" {
 		t.Fatalf("objects[0].TagName = %q", objects[0].GetTagName())
 	}
@@ -375,6 +390,8 @@ type fakeGalaxyServer struct {
 	failTest          bool
 	deployReply       *pb.GetLastDeployTimeReply
 	discoverReply     *pb.DiscoverHierarchyReply
+	discoverReplies   []*pb.DiscoverHierarchyReply
+	discoverRequests  []*pb.DiscoverHierarchyRequest
 	watchEvents       []*pb.DeployEvent
 	watchRequest      *pb.WatchDeployEventsRequest
 	watchSendInterval time.Duration
@@ -400,6 +417,12 @@ func (s *fakeGalaxyServer) GetLastDeployTime(ctx context.Context, req *pb.GetLas
 }
 
 func (s *fakeGalaxyServer) DiscoverHierarchy(ctx context.Context, req *pb.DiscoverHierarchyRequest) (*pb.DiscoverHierarchyReply, error) {
+	s.discoverRequests = append(s.discoverRequests, req)
+	if len(s.discoverReplies) > 0 {
+		reply := s.discoverReplies[0]
+		s.discoverReplies = s.discoverReplies[1:]
+		return reply, nil
+	}
 	if s.discoverReply != nil {
 		return s.discoverReply, nil
 	}

clients/go/mxgateway/options.go

@@ -18,6 +18,7 @@ type Options struct {
 	ServerNameOverride   string
 	DialTimeout          time.Duration
 	CallTimeout          time.Duration
+	MaxGrpcMessageBytes  int
 	TLSConfig            *tls.Config
 	TransportCredentials credentials.TransportCredentials
 	DialOptions          []grpc.DialOption

clients/go/mxgateway/version.go

@@ -7,7 +7,7 @@ const (
 
 	// GatewayProtocolVersion matches GatewayContractInfo.GatewayProtocolVersion
 	// in the shared .NET contracts.
-	GatewayProtocolVersion uint32 = 1
+	GatewayProtocolVersion uint32 = 2
 
 	// WorkerProtocolVersion matches GatewayContractInfo.WorkerProtocolVersion
 	// and is exposed for fake-worker and parity tests.