configure page_size as an optional parameter

Signed-off-by: Dmitry Shmulevich <[email protected]>
NVIDIA · Dec 19, 2024 · ebc574b · ebc574b
1 parent 9ea30c0
commit ebc574b
Show file tree

Hide file tree

Showing 10 changed files with 26 additions and 13 deletions.
diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -34,7 +34,7 @@ type Config struct {
 	RequestAggregationDelay time.Duration     `yaml:"request_aggregation_delay"`
 	Provider                string            `yaml:"provider,omitempty"`
 	Engine                  string            `yaml:"engine,omitempty"`
-	PageSize                int               `yaml:"page_size,omitempty"`
+	PageSize                *int              `yaml:"page_size,omitempty"`
 	SSL                     *SSL              `yaml:"ssl,omitempty"`
 	CredsPath               *string           `yaml:"credentials_path,omitempty"`
 	FwdSvcURL               *string           `yaml:"forward_service_url,omitempty"`

diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
@@ -22,6 +22,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/agrea/ptr"
 	"github.com/stretchr/testify/require"
 )
 
@@ -36,6 +37,7 @@ http:
   port: 49021
   ssl: true
 request_aggregation_delay: 15s
+page_size: 50
 ssl:
   cert: %s
   key: %s
@@ -89,6 +91,7 @@ func TestConfig(t *testing.T) {
 			SSL:  true,
 		},
 		RequestAggregationDelay: 15 * time.Second,
+		PageSize:                ptr.Int(50),
 		SSL: &SSL{
 			Cert:   cert.Name(),
 			Key:    key.Name(),

diff --git a/pkg/providers/aws/instance_topology.go b/pkg/providers/aws/instance_topology.go
@@ -32,11 +32,21 @@ import (
 
 var defaultPageSize int32 = 100
 
-func (p *Provider) generateInstanceTopology(ctx context.Context, pageSize int32, cis []topology.ComputeInstances) ([]types.InstanceTopology, error) {
-	var err error
-	topology := []types.InstanceTopology{}
+func (p *Provider) generateInstanceTopology(ctx context.Context, pageSize *int, cis []topology.ComputeInstances) ([]types.InstanceTopology, error) {
+	var (
+		err      error
+		topology []types.InstanceTopology
+		limit    int32
+	)
+
+	if pageSize != nil {
+		limit = int32(*pageSize)
+	} else {
+		limit = defaultPageSize
+	}
+
 	for _, ci := range cis {
-		if topology, err = p.generateInstanceTopologyForRegionInstances(ctx, pageSize, &ci, topology); err != nil {
+		if topology, err = p.generateInstanceTopologyForRegionInstances(ctx, limit, &ci, topology); err != nil {
 			return nil, err
 		}
 	}

diff --git a/pkg/providers/aws/provider.go b/pkg/providers/aws/provider.go
@@ -177,8 +177,8 @@ func New(clientFactory ClientFactory, imdsClient IDMSClient) *Provider {
 	}
 }
 
-func (p *Provider) GenerateTopologyConfig(ctx context.Context, pageSize int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
-	topology, err := p.generateInstanceTopology(ctx, int32(pageSize), instances)
+func (p *Provider) GenerateTopologyConfig(ctx context.Context, pageSize *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
+	topology, err := p.generateInstanceTopology(ctx, pageSize, instances)
 	if err != nil {
 		return nil, err
 	}

diff --git a/pkg/providers/baremetal/provider.go b/pkg/providers/baremetal/provider.go
@@ -26,7 +26,7 @@ func New() (*Provider, error) {
 	return &Provider{}, nil
 }
 
-func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
+func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
 	if len(instances) > 1 {
 		return nil, ErrMultiRegionNotSupported
 	}

diff --git a/pkg/providers/cw/provider.go b/pkg/providers/cw/provider.go
@@ -44,7 +44,7 @@ func New() (*Provider, error) {
 	return &Provider{}, nil
 }
 
-func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
+func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
 	if len(instances) > 1 {
 		return nil, fmt.Errorf("CW does not support mult-region topology requests")
 	}

diff --git a/pkg/providers/gcp/provider.go b/pkg/providers/gcp/provider.go
@@ -81,7 +81,7 @@ func New(clientFactory ClientFactory) (*Provider, error) {
 	}, nil
 }
 
-func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
+func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
 	if len(instances) > 1 {
 		return nil, fmt.Errorf("GCP does not support mult-region topology requests")
 	}

diff --git a/pkg/providers/oci/provider.go b/pkg/providers/oci/provider.go
@@ -145,7 +145,7 @@ func New(ociClientFactory ClientFactory) *Provider {
 	}
 }
 
-func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
+func (p *Provider) GenerateTopologyConfig(ctx context.Context, _ *int, instances []topology.ComputeInstances) (*topology.Vertex, error) {
 	cfg, err := GenerateInstanceTopology(ctx, p.clientFactory, instances)
 	if err != nil {
 		return nil, err

diff --git a/pkg/providers/providers.go b/pkg/providers/providers.go
@@ -26,7 +26,7 @@ import (
 )
 
 type Provider interface {
-	GenerateTopologyConfig(ctx context.Context, pageSize int, instances []topology.ComputeInstances) (*topology.Vertex, error)
+	GenerateTopologyConfig(ctx context.Context, pageSize *int, instances []topology.ComputeInstances) (*topology.Vertex, error)
 }
 
 type Config struct {

diff --git a/pkg/providers/test/test.go b/pkg/providers/test/test.go
@@ -75,6 +75,6 @@ func (p *Provider) GetComputeInstances(_ context.Context) ([]topology.ComputeIns
 	}, nil
 }
 
-func (p *Provider) GenerateTopologyConfig(_ context.Context, _ int, _ []topology.ComputeInstances) (*topology.Vertex, error) {
+func (p *Provider) GenerateTopologyConfig(_ context.Context, _ *int, _ []topology.ComputeInstances) (*topology.Vertex, error) {
 	return p.tree, nil
 }