Merge pull request #7211 from concourse/fix-rate-limiter

Treat resource types as checkable in rate limiter
concourse · Jun 24, 2021 · d7346bf · d7346bf
2 parents 58526cc + 6e04339
commit d7346bf
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 22 deletions.
diff --git a/atc/db/resource_check_rate_limiter.go b/atc/db/resource_check_rate_limiter.go
@@ -3,10 +3,11 @@ package db
 import (
 	"context"
 	"fmt"
-	sq "github.com/Masterminds/squirrel"
 	"sync"
 	"time"
 
+	sq "github.com/Masterminds/squirrel"
+
 	"code.cloudfoundry.org/clock"
 	"golang.org/x/time/rate"
 )
@@ -87,17 +88,28 @@ func (limiter *ResourceCheckRateLimiter) Limit() rate.Limit {
 }
 
 func (limiter *ResourceCheckRateLimiter) refreshCheckLimiter() error {
-	var count int
+	var resourceCount int
+	var resourceTypeCount int
 	err := psql.Select("COUNT(id)").
 		From("resources").
 		Where(sq.Eq{"active": true}).
 		RunWith(limiter.refreshConn).
 		QueryRow().
-		Scan(&count)
+		Scan(&resourceCount)
+	if err != nil {
+		return err
+	}
+	err = psql.Select("COUNT(id)").
+		From("resource_types").
+		Where(sq.Eq{"active": true}).
+		RunWith(limiter.refreshConn).
+		QueryRow().
+		Scan(&resourceTypeCount)
 	if err != nil {
 		return err
 	}
 
+	count := resourceCount + resourceTypeCount
 	limit := rate.Limit(float64(count) / limiter.checkInterval.Seconds())
 	if count == 0 {
 		// don't bother waiting if there aren't any checkables

diff --git a/atc/db/resource_check_rate_limiter_test.go b/atc/db/resource_check_rate_limiter_test.go
@@ -29,6 +29,8 @@ var _ = Describe("ResourceCheckRateLimiter", func() {
 		ctx context.Context
 
 		limiter *db.ResourceCheckRateLimiter
+
+		pipelineConfig *atc.Config
 	)
 
 	BeforeEach(func() {
@@ -37,13 +39,16 @@ var _ = Describe("ResourceCheckRateLimiter", func() {
 		refreshInterval = 5 * time.Minute
 		fakeClock = fakeclock.NewFakeClock(time.Now())
 
-		_, err := dbConn.Exec("update resources set active=false")
-		Expect(err).ToNot(HaveOccurred())
-
 		checkableCount = 0
 		scenario = &dbtest.Scenario{}
 
 		ctx = context.Background()
+
+		By("Unpolluting our env")
+		err := defaultPipeline.Destroy()
+		Expect(err).NotTo(HaveOccurred())
+
+		pipelineConfig = &atc.Config{}
 	})
 
 	JustBeforeEach(func() {
@@ -64,21 +69,34 @@ var _ = Describe("ResourceCheckRateLimiter", func() {
 		return errs
 	}
 
-	createCheckable := func() {
+	createResource := func() {
+		checkableCount++
+
+		pipelineConfig.Resources = append(pipelineConfig.Resources, atc.ResourceConfig{Name: fmt.Sprintf("resource-%d", checkableCount),
+			Type:   dbtest.BaseResourceType,
+			Source: atc.Source{"some": "source"},
+		})
+
+		scenario.Run(builder.WithPipeline(*pipelineConfig))
+	}
+
+	createResourceWithCustomType := func() {
+		checkableCount++
+		customResourceType := fmt.Sprintf("resource-type-%d", checkableCount)
+
+		pipelineConfig.ResourceTypes = append(pipelineConfig.ResourceTypes, atc.ResourceType{Name: customResourceType,
+			Type:   dbtest.BaseResourceType,
+			Source: atc.Source{"some": "source"},
+		})
+
 		checkableCount++
 
-		var resources atc.ResourceConfigs
-		for i := 0; i < checkableCount; i++ {
-			resources = append(resources, atc.ResourceConfig{
-				Name:   fmt.Sprintf("resource-%d", i),
-				Type:   dbtest.BaseResourceType,
-				Source: atc.Source{"some": "source"},
-			})
-		}
-
-		scenario.Run(builder.WithPipeline(atc.Config{
-			Resources: resources,
-		}))
+		pipelineConfig.Resources = append(pipelineConfig.Resources, atc.ResourceConfig{Name: fmt.Sprintf("resource-%d", checkableCount),
+			Type:   customResourceType,
+			Source: atc.Source{"some": "source"},
+		})
+
+		scenario.Run(builder.WithPipeline(*pipelineConfig))
 	}
 
 	Context("with no static limit provided", func() {
@@ -92,7 +110,7 @@ var _ = Describe("ResourceCheckRateLimiter", func() {
 			Expect(limiter.Limit()).To(Equal(rate.Inf))
 
 			By("creating one checkable")
-			createCheckable()
+			createResource()
 
 			By("continuing to return immediately, as the refresh interval has not elapsed")
 			Expect(<-wait(limiter)).To(Succeed())
@@ -118,7 +136,11 @@ var _ = Describe("ResourceCheckRateLimiter", func() {
 
 			By("creating more checkables")
 			for i := 0; i < 10; i++ {
-				createCheckable()
+				if i%3 == 0 {
+					createResourceWithCustomType()
+				} else {
+					createResource()
+				}
 			}
 
 			By("waiting for the refresh interval")
@@ -174,7 +196,7 @@ var _ = Describe("ResourceCheckRateLimiter", func() {
 
 			By("creating a few (ignored) checkables")
 			for i := 0; i < 10; i++ {
-				createCheckable()
+				createResource()
 			}
 
 			By("waiting for the (ignored) refresh interval")