Re-implement rate limiter using postgres (frain-dev#1937)

* feat: reimplement rate limiter using postgres

* feat: add bucket size (time duration) for rate limit

* chore: update tests

* chore: revert postgres pkg changes

* chore: go mod tidy

* feat: update rate when endpoint is updated

* feat: set default rate limit value to 0

* chore: update function name, add benchmarks

* chore: update test

* chore: update changes from main, remove dependency on the subscription repo

Author: jirevwe

cmd/agent/agent.go

@@ -14,10 +14,10 @@ import (
 	"github.com/frain-dev/convoy/database/postgres"
 	"github.com/frain-dev/convoy/internal/pkg/cli"
 	"github.com/frain-dev/convoy/internal/pkg/fflag"
+	"github.com/frain-dev/convoy/internal/pkg/limiter"
 	"github.com/frain-dev/convoy/internal/pkg/memorystore"
 	"github.com/frain-dev/convoy/internal/pkg/pubsub"
 	"github.com/frain-dev/convoy/internal/pkg/server"
-	"github.com/frain-dev/convoy/limiter"
 	"github.com/frain-dev/convoy/pkg/log"
 	"github.com/frain-dev/convoy/worker"
 	"github.com/frain-dev/convoy/worker/task"
@@ -202,16 +202,12 @@ func startWorkerComponent(ctx context.Context, a *cli.App) error {
 	subRepo := postgres.NewSubscriptionRepo(a.DB, a.Cache)
 	deviceRepo := postgres.NewDeviceRepo(a.DB, a.Cache)
 
-	rateLimiter, err := limiter.NewLimiter(cfg.Redis)
-	if err != nil {
-		a.Logger.Debug("Failed to initialise rate limiter")
-	}
+	rateLimiter := limiter.NewLimiter(a.DB)
 
 	consumer.RegisterHandlers(convoy.EventProcessor, task.ProcessEventDelivery(
 		endpointRepo,
 		eventDeliveryRepo,
 		projectRepo,
-		subRepo,
 		a.Queue,
 		rateLimiter))
 

const.go

@@ -8,9 +8,8 @@ const (
 )
 
 const (
-	RATE_LIMIT                      = 5000
-	RATE_LIMIT_DURATION             = 60
-	RATE_LIMIT_DURATION_IN_DURATION = time.Duration(RATE_LIMIT_DURATION) * time.Second
-	HTTP_TIMEOUT                    = 10
-	HTTP_TIMEOUT_IN_DURATION        = time.Duration(HTTP_TIMEOUT) * time.Second
+	RATE_LIMIT               = 0 // should be deleted
+	RATE_LIMIT_DURATION      = 0 // should be deleted
+	HTTP_TIMEOUT             = 10
+	HTTP_TIMEOUT_IN_DURATION = time.Duration(HTTP_TIMEOUT) * time.Second
 )

database/postgres/postgres.go

@@ -4,18 +4,17 @@ import (
 	"database/sql"
 	"errors"
 	"fmt"
+	"github.com/uptrace/opentelemetry-go-extra/otelsql"
+	"github.com/uptrace/opentelemetry-go-extra/otelsqlx"
 	"io"
 	"time"
 
+	"github.com/frain-dev/convoy/config"
 	"github.com/frain-dev/convoy/database/hooks"
 	"github.com/frain-dev/convoy/pkg/log"
-	"github.com/uptrace/opentelemetry-go-extra/otelsql"
-	"github.com/uptrace/opentelemetry-go-extra/otelsqlx"
-	semconv "go.opentelemetry.io/otel/semconv/v1.10.0"
-
-	"github.com/frain-dev/convoy/config"
 	"github.com/jmoiron/sqlx"
 	_ "github.com/newrelic/go-agent/v3/integrations/nrpq"
+	semconv "go.opentelemetry.io/otel/semconv/v1.10.0"
 )
 
 const pkgName = "postgres"

datastore/models.go

@@ -398,13 +398,13 @@ type Endpoint struct {
 	SupportEmail       string  `json:"support_email,omitempty" db:"support_email"`
 	AppID              string  `json:"-" db:"app_id"` // Deprecated but necessary for backward compatibility
 
-	HttpTimeout uint64         `json:"http_timeout" db:"http_timeout"`
-	RateLimit   int            `json:"rate_limit" db:"rate_limit"`
-	Events      int64          `json:"events,omitempty" db:"event_count"`
-	Status      EndpointStatus `json:"status" db:"status"`
+	Status         EndpointStatus          `json:"status" db:"status"`
+	HttpTimeout    uint64                  `json:"http_timeout" db:"http_timeout"`
+	Events         int64                   `json:"events,omitempty" db:"event_count"`
+	Authentication *EndpointAuthentication `json:"authentication" db:"authentication"`
 
-	RateLimitDuration uint64                  `json:"rate_limit_duration" db:"rate_limit_duration"`
-	Authentication    *EndpointAuthentication `json:"authentication" db:"authentication"`
+	RateLimit         int    `json:"rate_limit" db:"rate_limit"`
+	RateLimitDuration uint64 `json:"rate_limit_duration" db:"rate_limit_duration"`
 
 	CreatedAt time.Time `json:"created_at,omitempty" db:"created_at,omitempty" swaggertype:"string"`
 	UpdatedAt time.Time `json:"updated_at,omitempty" db:"updated_at,omitempty" swaggertype:"string"`

generate.go

@@ -2,7 +2,7 @@ package convoy
 
 //go:generate mockgen --source datastore/repository.go --destination mocks/repository.go -package mocks
 //go:generate mockgen --source queue/queue.go --destination mocks/queue.go -package mocks
-//go:generate mockgen --source limiter/limiter.go --destination mocks/limiter.go -package mocks
+//go:generate mockgen --source internal/pkg/limiter/limiter.go --destination mocks/limiter.go -package mocks
 //go:generate mockgen --source cache/cache.go --destination mocks/cache.go -package mocks
 //go:generate mockgen --source internal/pkg/smtp/smtp.go --destination mocks/smtp.go -package mocks
 //go:generate mockgen --source internal/pkg/socket/socket.go --destination mocks/socket.go -package mocks

go.mod

@@ -20,7 +20,6 @@ require (
 	github.com/go-chi/chi/v5 v5.0.10
 	github.com/go-chi/render v1.0.1
 	github.com/go-redis/cache/v9 v9.0.0
-	github.com/go-redis/redis_rate/v10 v10.0.1
 	github.com/go-redsync/redsync/v4 v4.8.1
 	github.com/golang-jwt/jwt v3.2.2+incompatible
 	github.com/golang/mock v1.6.0

go.sum

@@ -215,8 +215,6 @@ github.com/go-redis/redis/v8 v8.11.2/go.mod h1:DLomh7y2e3ggQXQLd1YgmvIfecPJoFl7W
 github.com/go-redis/redis/v8 v8.11.4/go.mod h1:2Z2wHZXdQpCDXEGzqMockDpNyYvi2l4Pxt6RJr792+w=
 github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI=
 github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo=
-github.com/go-redis/redis_rate/v10 v10.0.1 h1:calPxi7tVlxojKunJwQ72kwfozdy25RjA0bCj1h0MUo=
-github.com/go-redis/redis_rate/v10 v10.0.1/go.mod h1:EMiuO9+cjRkR7UvdvwMO7vbgqJkltQHtwbdIQvaBKIU=
 github.com/go-redsync/redsync/v4 v4.8.1 h1:rq2RvdTI0obznMdxKUWGdmmulo7lS9yCzb8fgDKOlbM=
 github.com/go-redsync/redsync/v4 v4.8.1/go.mod h1:LmUAsQuQxhzZAoGY7JS6+dNhNmZyonMZiiEDY9plotM=
 github.com/go-sourcemap/sourcemap v2.1.3+incompatible h1:W1iEw64niKVGogNgBN3ePyLFfuisuzeidWPMPWmECqU=

internal/pkg/limiter/limiter.go

@@ -0,0 +1,17 @@
+package limiter
+
+import (
+	"context"
+	"github.com/frain-dev/convoy/database"
+	"github.com/frain-dev/convoy/internal/pkg/limiter/pg"
+)
+
+type RateLimiter interface {
+	// Allow rate limits outgoing events to endpoints based on a rate in a specified time duration by the endpoint id
+	Allow(ctx context.Context, key string, rate int, duration int) error
+}
+
+func NewLimiter(db database.Database) RateLimiter {
+	ra := pg.NewRateLimiter(db)
+	return ra
+}

internal/pkg/limiter/pg/client.go

@@ -0,0 +1,56 @@
+package pg
+
+import (
+	"context"
+	"errors"
+	"github.com/frain-dev/convoy/database"
+	"github.com/frain-dev/convoy/pkg/log"
+)
+
+type SlidingWindowRateLimiter struct {
+	db database.Database
+}
+
+func NewRateLimiter(db database.Database) *SlidingWindowRateLimiter {
+	return &SlidingWindowRateLimiter{db: db}
+}
+
+func (p *SlidingWindowRateLimiter) Allow(ctx context.Context, key string, rate int, bucketSize int) error {
+	return p.takeToken(ctx, key, rate, bucketSize)
+}
+
+// TakeToken is a sliding window rate limiter that tries to take a token from the bucket
+//
+// Creates the bucket if it doesn't exist and returns false if it is not successful.
+// Returns true otherwise
+func (p *SlidingWindowRateLimiter) takeToken(ctx context.Context, key string, rate int, windowSize int) error {
+	// if one of rate and bucket size if zero, we skip processing
+	if rate == 0 || windowSize == 0 {
+		return nil
+	}
+
+	tx, err := p.db.GetDB().BeginTxx(ctx, nil)
+	if err != nil {
+		return err
+	}
+
+	var allowed bool
+	err = tx.QueryRowContext(ctx, `select convoy.take_token($1, $2, $3)::bool;`, key, rate, windowSize).Scan(&allowed)
+	if err != nil {
+		return err
+	}
+
+	err = tx.Commit()
+	if err != nil {
+		if rollbackErr := tx.Rollback(); rollbackErr != nil {
+			log.Infof("update failed: %v, unable to rollback: %v", err, rollbackErr)
+		}
+		return err
+	}
+
+	if !allowed {
+		return errors.New("rate limit error")
+	}
+
+	return nil
+}

internal/pkg/limiter/pg/client_test.go

@@ -0,0 +1,64 @@
+package pg
+
+import (
+	"context"
+	"github.com/frain-dev/convoy/config"
+	"github.com/frain-dev/convoy/database/postgres"
+	"github.com/frain-dev/convoy/pkg/log"
+	"github.com/stretchr/testify/require"
+	"os"
+	"testing"
+)
+
+func getConfig() config.Configuration {
+	_ = os.Setenv("CONVOY_DB_HOST", os.Getenv("TEST_REDIS_HOST"))
+	_ = os.Setenv("CONVOY_REDIS_SCHEME", os.Getenv("TEST_REDIS_SCHEME"))
+	_ = os.Setenv("CONVOY_REDIS_PORT", os.Getenv("TEST_REDIS_PORT"))
+
+	_ = os.Setenv("CONVOY_DB_HOST", os.Getenv("TEST_DB_HOST"))
+	_ = os.Setenv("CONVOY_DB_SCHEME", os.Getenv("TEST_DB_SCHEME"))
+	_ = os.Setenv("CONVOY_DB_USERNAME", os.Getenv("TEST_DB_USERNAME"))
+	_ = os.Setenv("CONVOY_DB_PASSWORD", os.Getenv("TEST_DB_PASSWORD"))
+	_ = os.Setenv("CONVOY_DB_DATABASE", os.Getenv("TEST_DB_DATABASE"))
+	_ = os.Setenv("CONVOY_DB_PORT", os.Getenv("TEST_DB_PORT"))
+
+	err := config.LoadConfig("")
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	cfg, err := config.Get()
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return cfg
+}
+
+func BenchmarkTakeToken_TakeOneToken(b *testing.B) {
+	db, err := postgres.NewDB(getConfig())
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		rateLimiter := NewRateLimiter(db)
+		tokenErr := rateLimiter.takeToken(context.Background(), "test", 100_000_000, 1)
+		require.NoError(b, tokenErr)
+	}
+}
+
+func BenchmarkTakeToken_TakeNone(b *testing.B) {
+	db, err := postgres.NewDB(getConfig())
+	require.NoError(b, err)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+
+	for i := 0; i < b.N; i++ {
+		rateLimiter := NewRateLimiter(db)
+		tokenErr := rateLimiter.takeToken(context.Background(), "test-2", 0, 0)
+		require.NoError(b, tokenErr)
+	}
+}