diff --git a/Makefile b/Makefile index 0120c720..106115ae 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,8 @@ test: ginkgo .PHONY: bench bench: - go run cmd/bench/bench.go + go build -o ./.bin/bench -v github.com/flanksource/duty/cmd/bench && \ + ./.bin/bench --count 250_000 fmt: go fmt ./... diff --git a/benchmark.md b/benchmark.md new file mode 100644 index 00000000..84cb1d05 --- /dev/null +++ b/benchmark.md @@ -0,0 +1,83 @@ +# PostgreSQL RLS Benchmark + +## Running Benchmarks + +query duration to fetch 10k, 25k, 50k and 100k config items in random are recorded. + +```bash +# With RLS +go run main.go + +# Without RLS +go run main.go --disable-rls +``` + +## Results + +### Without RLS + +```json +[ + { + "config_count": 10000, + "duration": 766282108, + "rls_enabled": false + }, + { + "config_count": 25000, + "duration": 1964247115, + "rls_enabled": false + }, + { + "config_count": 50000, + "duration": 4738917435, + "rls_enabled": false + }, + { + "config_count": 100000, + "duration": 7663285526, + "rls_enabled": false + } +] +``` + +| Configuration Count | Duration (s) | +| ------------------- | ------------ | +| 10,000 | 766.3 | +| 25,000 | 1.964 | +| 50,000 | 4.738 | +| 100,000 | 7.663 | + +### With RLS + +```json +[ + { + "config_count": 10000, + "duration": 897432497, + "rls_enabled": true + }, + { + "config_count": 25000, + "duration": 1794868472, + "rls_enabled": true + }, + { + "config_count": 50000, + "duration": 4383678988, + "rls_enabled": true + }, + { + "config_count": 100000, + "duration": 7650471155, + "rls_enabled": true + } +] +``` + +| Configuration Count | Duration (s) | +| ------------------- | ------------ | +| 10,000 | 897.4 | +| 25,000 | 1.794 | +| 50,000 | 4.383 | +| 100,000 | 7.650 | diff --git a/cmd/bench/bench.go b/cmd/bench/bench.go index 0c2e8add..25511ea2 100644 --- a/cmd/bench/bench.go +++ b/cmd/bench/bench.go @@ -1,19 +1,189 @@ package main import ( + "encoding/json" + "errors" "flag" + "fmt" + "math/rand" + "os" + "path/filepath" "time" + "github.com/google/uuid" + "github.com/samber/lo" + "github.com/flanksource/commons/logger" "github.com/flanksource/duty/context" + "github.com/flanksource/duty/models" + "github.com/flanksource/duty/query" "github.com/flanksource/duty/shutdown" pkgGenerator "github.com/flanksource/duty/tests/generator" "github.com/flanksource/duty/tests/setup" ) -var count = 2_000 +var ( + dbDataPath string + count = 250_000 + disableRLS bool +) + +type BenchType string + +const ( + BenchTypeFetchConfig BenchType = "config_query" + BenchTypeFetchOutgoing BenchType = "outgoing_configs" + BenchTypeFetchIncoming BenchType = "incoming_configs" + BenchTypeFetchBoth BenchType = "incoming_and_outgoing_configs" +) + +type BenchmarkResult struct { + BenchType BenchType `json:"bench_type"` + ConfigCount int `json:"config_count"` + Duration time.Duration `json:"duration"` + RLSEnabled bool `json:"rls_enabled"` +} + +func main() { + shutdown.WaitForSignal() + + flag.IntVar(&count, "count", count, "generates at least these number of config items") + flag.BoolVar(&disableRLS, "disable-rls", false, "disable rls") + flag.StringVar(&dbDataPath, "db-data-path", "", "use existing postgres data dir to skip insertion of dummy data") + flag.Parse() + + if dbDataPath != "" { + os.Setenv(setup.DUTY_DB_DATA_DIR, dbDataPath) + } else if v, ok := os.LookupEnv(setup.DUTY_DB_DATA_DIR); ok { + dbDataPath = v + } + + if err := run(count, disableRLS); err != nil { + shutdown.ShutdownAndExit(1, err.Error()) + } + + shutdown.ShutdownAndExit(0, "exiting ...") +} + +func run(count int, disableRLS bool) error { + args := []any{setup.WithoutDummyData} // we generate the dummy dataa + if disableRLS { + args = append(args, setup.WithoutRLS) + } + + ctx, err := setup.SetupDB("test", args...) + if err != nil { + return err + } + + var allConfigIDs []uuid.UUID + if dbDataPath == "" { + generatedList, err := generateConfigItems(ctx, count) + if err != nil { + return err + } + + allConfigIDs = getAllConfigIDs(generatedList) + } else { + if err := ctx.DB().Select("id").Model(&models.ConfigItem{}).Find(&allConfigIDs).Error; err != nil { + return err + } + logger.Infof("fetched %d config ids from database", len(allConfigIDs)) + } + + if !disableRLS { + if err := setupRLSPayload(ctx); err != nil { + return err + } + } + + var benchResults []BenchmarkResult + + for _, size := range []int{10_000, 25_000, 50_000, 100_000} { + ids := shuffleAndPickNIDs(allConfigIDs, size) + + start := time.Now() + if err := fetchConfigs(ctx, ids); err != nil { + return err + } + ctx.Infof("fetched %d configs in %s", size, time.Since(start)) + + benchResults = append(benchResults, BenchmarkResult{ + BenchType: BenchTypeFetchConfig, + ConfigCount: size, + Duration: time.Since(start), + RLSEnabled: !disableRLS, + }) + } + + for _, size := range []int{50, 100, 250} { + ids := shuffleAndPickNIDs(allConfigIDs, size) + + start := time.Now() + if err := fetchRelatedConfigs(ctx, query.Outgoing, ids); err != nil { + return err + } + ctx.Infof("fetched outgoing configs for %d configs in %s", size, time.Since(start)) + + benchResults = append(benchResults, BenchmarkResult{ + BenchType: BenchTypeFetchOutgoing, + ConfigCount: size, + Duration: time.Since(start), + RLSEnabled: !disableRLS, + }) + } + + for _, size := range []int{50, 100, 250} { + ids := shuffleAndPickNIDs(allConfigIDs, size) + + start := time.Now() + if err := fetchRelatedConfigs(ctx, query.Incoming, ids); err != nil { + return err + } + ctx.Infof("fetched incoming configs for %d configs in %s", size, time.Since(start)) + + benchResults = append(benchResults, BenchmarkResult{ + BenchType: BenchTypeFetchIncoming, + ConfigCount: size, + Duration: time.Since(start), + RLSEnabled: !disableRLS, + }) + } + + for _, size := range []int{50, 100, 250} { + ids := shuffleAndPickNIDs(allConfigIDs, size) + + start := time.Now() + if err := fetchRelatedConfigs(ctx, query.All, ids); err != nil { + return err + } + ctx.Infof("fetched incoming/outgoing configs for %d configs in %s", size, time.Since(start)) + + benchResults = append(benchResults, BenchmarkResult{ + BenchType: BenchTypeFetchBoth, + ConfigCount: size, + Duration: time.Since(start), + RLSEnabled: !disableRLS, + }) + } + + // TODO: add more benchmarks for related changes + jsonData, err := json.MarshalIndent(benchResults, "", " ") + if err != nil { + return err + } + + filename := filepath.Join(".bin", fmt.Sprintf("bench_%s.json", lo.Ternary(disableRLS, "without_rls", "with_rls"))) + if err := os.WriteFile(filename, jsonData, 0644); err != nil { + return err + } + + return nil +} + +func generateConfigItems(ctx context.Context, count int) ([]pkgGenerator.Generated, error) { + var output []pkgGenerator.Generated -func generateConfigItems(ctx context.Context, count int) error { start := time.Now() for { generator := pkgGenerator.ConfigGenerator{ @@ -42,12 +212,13 @@ func generateConfigItems(ctx context.Context, count int) error { generator.GenerateKubernetes() if err := generator.Save(ctx.DB()); err != nil { - return err + return nil, err } + output = append(output, generator.Generated) var totalConfigs int64 if err := ctx.DB().Table("config_items").Count(&totalConfigs).Error; err != nil { - return err + return nil, err } if totalConfigs > int64(count) { @@ -59,46 +230,113 @@ func generateConfigItems(ctx context.Context, count int) error { var configs int64 if err := ctx.DB().Table("config_items").Count(&configs).Error; err != nil { - return err + return nil, err } var changes int64 if err := ctx.DB().Table("config_changes").Count(&changes).Error; err != nil { - return err + return nil, err } logger.Infof("configs %d, changes: %d in %s", configs, changes, time.Since(start)) + return output, nil +} + +func fetchRelatedConfigs(ctx context.Context, relation query.RelationDirection, ids []uuid.UUID) error { + start := time.Now() + + var total int + for i, id := range ids { + res, err := query.GetRelatedConfigs(ctx, query.RelationQuery{ + ID: id, + Relation: relation, + Incoming: query.Both, + Outgoing: query.Both, + MaxDepth: lo.ToPtr(5), + }) + if err != nil { + return fmt.Errorf("failed to fetch relationships for config %s: %w", id, err) + } + total += len(res) + + if (i+1)%50 == 0 { + ctx.Infof("progress:: fetched %d %s configs for %d configs in %s", total, relation, i+1, time.Since(start)) + } + } + return nil } -func main() { - shutdown.WaitForSignal() - flag.IntVar(&count, "count", count, "generates at least these number of configs") - flag.Parse() +func fetchConfigs(ctx context.Context, ids []uuid.UUID) error { + for _, id := range ids { + var config models.ConfigItem + if err := ctx.DB().Find(&config, "id = ?", id).Error; err != nil { + return fmt.Errorf("failed to fetch config %s: %w", id, err) + } + } - // start a postgres db with RLS disabled - if err := run(); err != nil { - shutdown.ShutdownAndExit(1, err.Error()) + return nil +} + +func getAllConfigIDs(generatedList []pkgGenerator.Generated) []uuid.UUID { + var allIDs []uuid.UUID + idMap := make(map[uuid.UUID]struct{}) + + for _, gen := range generatedList { + for _, config := range gen.Configs { + if _, exists := idMap[config.ID]; !exists { + idMap[config.ID] = struct{}{} + allIDs = append(allIDs, config.ID) + } + } } - // TODO: run benchmark on another database RLS enabled - // can't use the same database to avoid caches from the previous benchmark. + return allIDs +} - shutdown.ShutdownAndExit(0, "exiting ...") +func shuffleAndPickNIDs(ids []uuid.UUID, size int) []uuid.UUID { + if size > len(ids) { + size = len(ids) + } + + result := make([]uuid.UUID, len(ids)) + copy(result, ids) + rng := rand.New(rand.NewSource(time.Now().UnixNano())) + + for i := len(result) - 1; i > 0; i-- { + j := rng.Intn(i + 1) + result[i], result[j] = result[j], result[i] + } + + return result[:size] } -func run() error { - // setup a db with RLS disabled - ctx, err := setup.SetupDB("test", setup.WithoutRLS) - if err != nil { +func setupRLSPayload(ctx context.Context) error { + if err := ctx.DB().Exec(`SET request.jwt.claims = '{"tags": [{"cluster": "homelab"}]}'`).Error; err != nil { + return err + } + + var jwt string + if err := ctx.DB().Raw(`SELECT current_setting('request.jwt.claims', TRUE)`).Scan(&jwt).Error; err != nil { return err } - if err := generateConfigItems(ctx, count); err != nil { + if jwt == "" { + return errors.New("jwt parameter not set") + } + + if err := ctx.DB().Exec(`SET role = 'postgrest_api'`).Error; err != nil { return err } - // Run fetch queries + var role string + if err := ctx.DB().Raw(`SELECT CURRENT_USER`).Scan(&role).Error; err != nil { + return err + } + + if role != "postgrest_api" { + return errors.New("role is not set") + } return nil } diff --git a/tests/setup/common.go b/tests/setup/common.go index ef47fb71..650f295c 100644 --- a/tests/setup/common.go +++ b/tests/setup/common.go @@ -35,6 +35,14 @@ import ( _ "github.com/spf13/cobra" ) +// Env vars for embedded db +const ( + DUTY_DB_CREATE = "DUTY_DB_CREATE" + DUTY_DB_DATA_DIR = "DUTY_DB_DATA_DIR" + DUTY_DB_URL = "DUTY_DB_URL" + TEST_DB_PORT = "TEST_DB_PORT" +) + var DefaultContext context.Context var postgresServer *embeddedPG.EmbeddedPostgres @@ -82,7 +90,7 @@ func MustDB() *sql.DB { var WithoutRLS = "rsl_disabled" var WithoutDummyData = "without_dummy_data" var WithExistingDatabase = "with_existing_database" -var recreateDatabase = os.Getenv("DUTY_DB_CREATE") != "false" +var recreateDatabase = os.Getenv(DUTY_DB_CREATE) != "false" func findFileInPath(filename string, depth int) string { if !path.IsAbs(filename) { @@ -140,7 +148,7 @@ func SetupDB(dbName string, args ...interface{}) (context.Context, error) { } var port int - if val, ok := os.LookupEnv("TEST_DB_PORT"); ok { + if val, ok := os.LookupEnv(TEST_DB_PORT); ok { parsed, err := strconv.ParseInt(val, 10, 32) if err != nil { return context.Context{}, err @@ -152,7 +160,7 @@ func SetupDB(dbName string, args ...interface{}) (context.Context, error) { } PgUrl = fmt.Sprintf("postgres://postgres:postgres@localhost:%d/%s?sslmode=disable", port, dbName) - url := os.Getenv("DUTY_DB_URL") + url := os.Getenv(DUTY_DB_URL) if url != "" && !recreateDatabase { PgUrl = url } else if url != "" && recreateDatabase { @@ -172,6 +180,12 @@ func SetupDB(dbName string, args ...interface{}) (context.Context, error) { } else if url == "" && postgresServer == nil { config, _ := GetEmbeddedPGConfig(dbName, port) + + // allow data dir override + if v, ok := os.LookupEnv(DUTY_DB_DATA_DIR); ok { + config = config.DataPath(v) + } + postgresServer = embeddedPG.NewDatabase(config) logger.Infof("starting embedded postgres on port %d", port) if err := postgresServer.Start(); err != nil {