bench: querying config items

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 --disable-rls
 
 fmt:
 	go fmt ./...

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        |

cmd/bench/bench.go

@@ -1,19 +1,121 @@
 package main
 
 import (
+	"encoding/json"
+	"errors"
 	"flag"
+	"fmt"
+	"math/rand"
+	"os"
 	"time"
 
+	"github.com/google/uuid"
+
 	"github.com/flanksource/commons/logger"
 	"github.com/flanksource/duty/context"
+	"github.com/flanksource/duty/models"
 	"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 BenchmarkResult struct {
+	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
+
+	querySize := []int{10_000, 25_000, 50_000, 100_000}
+	for _, size := range querySize {
+		ids := shuffleAndPickNIDs(allConfigIDs, size)
+
+		start := time.Now()
+		if err := fetchConfigs(ctx, ids); err != nil {
+			return err
+		}
+
+		benchResults = append(benchResults, BenchmarkResult{
+			ConfigCount: size,
+			Duration:    time.Since(start),
+			RLSEnabled:  !disableRLS,
+		})
+	}
+
+	// Todo: add more benchmarks for changes and relationships queries
+
+	jsonData, err := json.MarshalIndent(benchResults, "", "  ")
+	if err != nil {
+		return err
+	}
+
+	fmt.Println(string(jsonData))
+	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 +144,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 +162,106 @@ 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 fetchConfigs(ctx context.Context, ids []uuid.UUID) error {
+	var found int
+	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)
+		} else if config.ID == id {
+			found++
+		}
+
+		// var changes []models.ConfigChange
+		// if err := ctx.DB().Where("config_id = ?", id).Find(&changes).Error; err != nil {
+		// 	return fmt.Errorf("failed to fetch changes for config %s: %w", id, err)
+		// }
+
+		// var analysis []models.ConfigAnalysis
+		// if err := ctx.DB().Where("config_id = ?", id).Find(&analysis).Error; err != nil {
+		// 	return fmt.Errorf("failed to fetch analysis for config %s: %w", id, err)
+		// }
+
+		// var relationships []models.ConfigRelationship
+		// if err := ctx.DB().Where("config_id = ? OR related_id = ?", id, id).Find(&relationships).Error; err != nil {
+		// 	return fmt.Errorf("failed to fetch relationships for config %s: %w", id, err)
+		// }
+	}
+
 	return nil
 }
 
-func main() {
-	shutdown.WaitForSignal()
-	flag.IntVar(&count, "count", count, "generates at least these number of configs")
-	flag.Parse()
+func getAllConfigIDs(generatedList []pkgGenerator.Generated) []uuid.UUID {
+	var allIDs []uuid.UUID
+	idMap := make(map[uuid.UUID]struct{})
 
-	// start a postgres db with RLS disabled
-	if err := run(); err != nil {
-		shutdown.ShutdownAndExit(1, err.Error())
+	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 jwt == "" {
+		return errors.New("jwt parameter not set")
+	}
+
+	if err := ctx.DB().Exec(`SET role = 'postgrest_api'`).Error; err != nil {
 		return err
 	}
 
-	if err := generateConfigItems(ctx, count); err != nil {
+	var role string
+	if err := ctx.DB().Raw(`SELECT CURRENT_USER`).Scan(&role).Error; err != nil {
 		return err
 	}
 
-	// Run fetch queries
+	if role != "postgrest_api" {
+		return errors.New("role is not set")
+	}
 
 	return nil
 }