feat: concurrent resource fetching per provider

internal/internal.go

@@ -11,6 +11,7 @@ import (
 	"runtime"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/getsentry/sentry-go"
@@ -54,7 +55,7 @@ var Arch = runtime.GOARCH
 var db *bun.DB
 var analytics utils.Analytics
 
-func Exec(address string, port int, configPath string, telemetry bool, a utils.Analytics, regions []string, cmd *cobra.Command) error {
+func Exec(address string, port int, configPath string, telemetry bool, a utils.Analytics, regions []string, _ *cobra.Command) error {
 	analytics = a
 
 	ctx := context.Background()
@@ -79,10 +80,15 @@ func Exec(address string, port int, configPath string, telemetry bool, a utils.A
 
 		_, err = cron.Every(1).Hours().Do(func() {
 			log.Info("Fetching resources workflow has started")
-			err = fetchResources(ctx, clients, regions, telemetry)
+
+			numWorkers := 64
+			wp := providers.NewWorkerPool(numWorkers)
+			wp.Start()
+			err = fetchResources(ctx, clients, regions, telemetry, wp)
 			if err != nil {
 				log.Fatal(err)
 			}
+			wp.Wait()
 		})
 
 		if err != nil {
@@ -209,7 +215,7 @@ func setupDBConnection(c *models.Config) error {
 	return nil
 }
 
-func triggerFetchingWorfklow(ctx context.Context, client providers.ProviderClient, provider string, telemetry bool, regions []string) {
+func triggerFetchingWorfklow(ctx context.Context, client providers.ProviderClient, provider string, telemetry bool, regions []string, wp *providers.WorkerPool) {
 	localHub := sentry.CurrentHub().Clone()
 
 	defer func() {
@@ -233,7 +239,7 @@ func triggerFetchingWorfklow(ctx context.Context, client providers.ProviderClien
 
 	switch provider {
 	case "AWS":
-		aws.FetchResources(ctx, client, regions, db, telemetry, analytics)
+		aws.FetchResources(ctx, client, regions, db, telemetry, analytics, wp)
 	case "DigitalOcean":
 		do.FetchResources(ctx, client, db, telemetry, analytics)
 	case "OCI":
@@ -257,32 +263,44 @@ func triggerFetchingWorfklow(ctx context.Context, client providers.ProviderClien
 	}
 }
 
-func fetchResources(ctx context.Context, clients []providers.ProviderClient, regions []string, telemetry bool) error {
+func fetchResources(ctx context.Context, clients []providers.ProviderClient, regions []string, telemetry bool, wp *providers.WorkerPool) error {
+
+	var wwg sync.WaitGroup
+	workflowTrigger := func(client providers.ProviderClient, provider string) {
+		wwg.Add(1)
+		go func() {
+			defer wwg.Done()
+			triggerFetchingWorfklow(ctx, client, provider, telemetry, regions, wp)
+		}()
+	}
+
 	for _, client := range clients {
 		if client.AWSClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "AWS", telemetry, regions)
+			workflowTrigger(client, "AWS")
 		} else if client.DigitalOceanClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "DigitalOcean", telemetry, regions)
+			workflowTrigger(client, "DigitalOcean")
 		} else if client.OciClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "OCI", telemetry, regions)
+			workflowTrigger(client, "OCI")
 		} else if client.CivoClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "Civo", telemetry, regions)
+			workflowTrigger(client, "Civo")
 		} else if client.K8sClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "Kubernetes", telemetry, regions)
+			workflowTrigger(client, "Kubernetes")
 		} else if client.LinodeClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "Linode", telemetry, regions)
+			workflowTrigger(client, "Linode")
 		} else if client.TencentClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "Tencent", telemetry, regions)
+			workflowTrigger(client, "Tencent")
 		} else if client.AzureClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "Azure", telemetry, regions)
+			workflowTrigger(client, "Azure")
 		} else if client.ScalewayClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "Scaleway", telemetry, regions)
+			workflowTrigger(client, "Scaleway")
 		} else if client.MongoDBAtlasClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "MongoDBAtlas", telemetry, regions)
+			workflowTrigger(client, "MongoDBAtlas")
 		} else if client.GCPClient != nil {
-			go triggerFetchingWorfklow(ctx, client, "GCP", telemetry, regions)
+			workflowTrigger(client, "GCP")
 		}
 	}
+
+	wwg.Wait()
 	return nil
 }
 

providers/aws/aws.go

@@ -98,7 +98,7 @@ func listOfSupportedServices() []providers.FetchDataFunction {
 	}
 }
 
-func FetchResources(ctx context.Context, client providers.ProviderClient, regions []string, db *bun.DB, telemetry bool, analytics utils.Analytics) {
+func FetchResources(ctx context.Context, client providers.ProviderClient, regions []string, db *bun.DB, telemetry bool, analytics utils.Analytics, wp *providers.WorkerPool) {
 	listOfSupportedRegions := getRegions()
 	if len(regions) > 0 {
 		log.Infof("Komiser will fetch resources from the following regions: %s", strings.Join(regions, ","))
@@ -108,23 +108,25 @@ func FetchResources(ctx context.Context, client providers.ProviderClient, region
 	for _, region := range listOfSupportedRegions {
 		client.AWSClient.Region = region
 		for _, fetchResources := range listOfSupportedServices() {
-			resources, err := fetchResources(ctx, client)
-			if err != nil {
-				log.Warnf("[%s][AWS] %s", client.Name, err)
-			} else {
-				for _, resource := range resources {
-				_, err = db.NewInsert().Model(&resource).On("CONFLICT (resource_id) DO UPDATE").Set("cost = EXCLUDED.cost, relations=EXCLUDED.relations").Exec(context.Background())
-					if err != nil {
-						log.WithError(err).Errorf("db trigger failed")
+			wp.SubmitTask(func() {
+				resources, err := fetchResources(ctx, client)
+				if err != nil {
+					log.Warnf("[%s][AWS] %s", client.Name, err)
+				} else {
+					for _, resource := range resources {
+						_, err = db.NewInsert().Model(&resource).On("CONFLICT (resource_id) DO UPDATE").Set("cost = EXCLUDED.cost, relations=EXCLUDED.relations").Exec(context.Background())
+						if err != nil {
+							log.WithError(err).Errorf("db trigger failed")
+						}
+					}
+					if telemetry {
+						analytics.TrackEvent("discovered_resources", map[string]interface{}{
+							"provider":  "AWS",
+							"resources": len(resources),
+						})
 					}
 				}
-				if telemetry {
-					analytics.TrackEvent("discovered_resources", map[string]interface{}{
-						"provider":  "AWS",
-						"resources": len(resources),
-					})
-				}
-			}
+			})
 		}
 	}
 }

providers/providers.go

@@ -2,6 +2,7 @@ package providers
 
 import (
 	"context"
+	"sync"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azidentity"
 	"github.com/aws/aws-sdk-go-v2/aws"
@@ -47,3 +48,39 @@ type K8sClient struct {
 	Client          *kubernetes.Clientset
 	OpencostBaseUrl string
 }
+
+type WorkerPool struct {
+	numWorkers int
+	tasks      chan func()
+	wg         sync.WaitGroup
+}
+
+func NewWorkerPool(numWorkers int) *WorkerPool {
+	return &WorkerPool{
+		numWorkers: numWorkers,
+		tasks:      make(chan func()),
+	}
+}
+
+func (wp *WorkerPool) Start() {
+	for i := 0; i < wp.numWorkers; i++ {
+		go wp.worker()
+	}
+}
+
+func (wp *WorkerPool) SubmitTask(task func()) {
+	wp.wg.Add(1)
+	wp.tasks <- task
+}
+
+func (wp *WorkerPool) Wait() {
+	wp.wg.Wait()
+	close(wp.tasks)
+}
+
+func (wp *WorkerPool) worker() {
+	for task := range wp.tasks {
+		task()
+		wp.wg.Done()
+	}
+}