stemedb/cmd/load-test/main.go

// Package main provides load testing for StemeDB.
//
// This tool benchmarks StemeDB performance with three scenarios:
//   - Baseline: 10K assertions to establish latency baseline
//   - Sustained: 1K writes/sec for configurable duration
//   - Concurrent: 100 concurrent readers with background writes
//
// Usage:
//
//	go run ./cmd/load-test --api-url http://localhost:18180
//	go run ./cmd/load-test --scenario sustained --duration 1h
//	go run ./cmd/load-test --scenario concurrent --readers 100
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/orchard9/stemedb-go/steme"
)

const (
	defaultAPIURL      = "http://127.0.0.1:18180"
	defaultKeysFile    = "demo/keys/agents.json"
	defaultOutputDir   = "uat/production-readiness/results"
	defaultDuration    = 5 * time.Minute
	defaultTargetRPS   = 1000
	defaultReaderCount = 100
)

// DemoAgent represents a demo agent from agents.json.
type DemoAgent struct {
	Seed        string `json:"seed"`
	PublicKey   string `json:"public_key"`
	Tier        int    `json:"tier"`
	Description string `json:"description"`
}

func main() {
	// Parse flags
	apiURL := flag.String("api-url", "", "StemeDB API URL (default: $STEMEDB_API_URL or http://127.0.0.1:18180)")
	keysFile := flag.String("keys-file", "", "Path to agents.json (default: demo/keys/agents.json)")
	scenario := flag.String("scenario", "all", "Test scenario: baseline|sustained|concurrent|all")
	duration := flag.Duration("duration", defaultDuration, "Duration for sustained write test")
	targetRPS := flag.Int("target-rps", defaultTargetRPS, "Target requests per second for sustained test")
	readers := flag.Int("readers", defaultReaderCount, "Number of concurrent readers")
	outputDir := flag.String("output", "", "Output directory for reports (default: uat/production-readiness/results)")
	verbose := flag.Bool("verbose", false, "Enable verbose output")
	flag.Parse()

	// Resolve API URL
	url := *apiURL
	if url == "" {
		url = os.Getenv("STEMEDB_API_URL")
	}
	if url == "" {
		url = defaultAPIURL
	}

	// Resolve keys file
	keysPath := *keysFile
	if keysPath == "" {
		keysPath = findKeysFile()
	}

	// Resolve output directory
	outDir := *outputDir
	if outDir == "" {
		outDir = findOutputDir()
	}

	// Check if meter is disabled
	meterDisabled := os.Getenv("STEMEDB_METER_ENABLED") == "false"

	fmt.Println("=== StemeDB Load Tester ===")
	fmt.Println()
	fmt.Printf("API URL:     %s\n", url)
	fmt.Printf("Keys File:   %s\n", keysPath)
	fmt.Printf("Scenario:    %s\n", *scenario)
	fmt.Printf("Output Dir:  %s\n", outDir)
	if *scenario == "sustained" || *scenario == "all" {
		fmt.Printf("Duration:    %v\n", *duration)
		fmt.Printf("Target RPS:  %d\n", *targetRPS)
		if !meterDisabled {
			fmt.Println("WARNING: Meter is enabled. Set STEMEDB_METER_ENABLED=false for accurate sustained test results.")
		}
	}
	if *scenario == "concurrent" || *scenario == "all" {
		fmt.Printf("Readers:     %d\n", *readers)
	}
	fmt.Println()

	// Load agents
	clients, signers, err := loadAgents(keysPath, url)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to load agents: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("Loaded %d demo agents\n\n", len(clients))

	// Verify server is reachable
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// Handle interrupt
	sigCh := make(chan os.Signal, 1)
	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-sigCh
		fmt.Println("\nInterrupted. Stopping tests...")
		cancel()
	}()

	health, err := clients[0].Health(ctx)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Server not reachable at %s: %v\n", url, err)
		fmt.Fprintln(os.Stderr, "Start the server with: STEMEDB_METER_ENABLED=false cargo run --release --bin stemedb-api")
		os.Exit(1)
	}
	fmt.Printf("Server Status: %s (v%s, %d assertions)\n\n", health.Status, health.Version, health.AssertionsCount)

	// Prepare scenario config
	cfg := ScenarioConfig{
		Clients:   clients,
		Signers:   signers,
		Duration:  *duration,
		TargetRPS: *targetRPS,
		Readers:   *readers,
		Verbose:   *verbose,
	}

	// Run scenarios
	results := TestResults{
		APIURL:        url,
		Duration:      *duration,
		TargetRPS:     *targetRPS,
		ReaderCount:   *readers,
		MeterDisabled: meterDisabled,
	}

	switch *scenario {
	case "baseline":
		results.BaselineMetrics, err = RunBaselineLatency(ctx, cfg)
	case "sustained":
		results.SustainedMetrics, err = RunSustainedWrites(ctx, cfg)
	case "concurrent":
		results.ConcurrentMetrics, err = RunConcurrentReaders(ctx, cfg)
	case "all":
		results.BaselineMetrics, err = RunBaselineLatency(ctx, cfg)
		if err == nil {
			results.SustainedMetrics, err = RunSustainedWrites(ctx, cfg)
		}
		if err == nil {
			results.ConcurrentMetrics, err = RunConcurrentReaders(ctx, cfg)
		}
	default:
		fmt.Fprintf(os.Stderr, "Unknown scenario: %s\n", *scenario)
		os.Exit(1)
	}

	if err != nil && err != context.Canceled {
		fmt.Fprintf(os.Stderr, "Test failed: %v\n", err)
		os.Exit(1)
	}

	// Generate report
	fmt.Println("=== Generating Report ===")
	reportPath, err := GenerateReport(results, outDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Failed to generate report: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("Report saved to: %s\n", reportPath)

	// Print summary
	fmt.Println()
	fmt.Println("=== Summary ===")
	printSummary(results)
}

// loadAgents loads demo agents and creates clients.
func loadAgents(path, baseURL string) ([]*steme.Client, []*steme.Signer, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, nil, fmt.Errorf("failed to read %s: %w", path, err)
	}

	var agents map[string]DemoAgent
	if err := json.Unmarshal(data, &agents); err != nil {
		return nil, nil, fmt.Errorf("failed to parse %s: %w", path, err)
	}

	var clients []*steme.Client
	var signers []*steme.Signer

	for name, agent := range agents {
		signer, err := steme.NewSignerFromHex(agent.Seed)
		if err != nil {
			return nil, nil, fmt.Errorf("failed to create signer for %s: %w", name, err)
		}
		signers = append(signers, signer)
		clients = append(clients, steme.NewClient(baseURL, signer))
	}

	return clients, signers, nil
}

// findKeysFile locates the agents.json file.
func findKeysFile() string {
	// Try current directory
	if _, err := os.Stat(defaultKeysFile); err == nil {
		return defaultKeysFile
	}

	// Try one level up (if running from cmd/load-test)
	upPath := filepath.Join("..", "..", defaultKeysFile)
	if _, err := os.Stat(upPath); err == nil {
		return upPath
	}

	return defaultKeysFile
}

// findOutputDir locates the output directory.
func findOutputDir() string {
	// Try current directory
	if _, err := os.Stat("uat/production-readiness"); err == nil {
		return defaultOutputDir
	}

	// Try one level up
	upPath := filepath.Join("..", "..", defaultOutputDir)
	if _, err := os.Stat(filepath.Dir(upPath)); err == nil {
		return upPath
	}

	return defaultOutputDir
}

func printSummary(results TestResults) {
	allPass := true

	if results.BaselineMetrics != nil {
		m := results.BaselineMetrics
		pass := m.WriteP99() <= Targets.BaselineP99 && m.WriteErrorRate() == 0
		status := "PASS"
		if !pass {
			status = "FAIL"
			allPass = false
		}
		fmt.Printf("Baseline:   %s (p99: %v, errors: %d)\n", status, m.WriteP99(), m.WriteErrors())
	}

	if results.SustainedMetrics != nil {
		m := results.SustainedMetrics
		pass := m.WriteP99() <= Targets.SustainedP99 && m.WriteErrorRate() <= Targets.SustainedErrRate
		status := "PASS"
		if !pass {
			status = "FAIL"
			allPass = false
		}
		fmt.Printf("Sustained:  %s (p99: %v, %.0f/sec, errors: %d)\n",
			status, m.WriteP99(), m.WriteThroughput(), m.WriteErrors())
	}

	if results.ConcurrentMetrics != nil {
		m := results.ConcurrentMetrics
		pass := m.DegradationPct() <= Targets.DegradationPct
		status := "PASS"
		if !pass {
			status = "FAIL"
			allPass = false
		}
		fmt.Printf("Concurrent: %s (degradation: %.1f%%, readers: %d)\n",
			status, m.DegradationPct(), m.ReaderCount())
	}

	fmt.Println()
	if allPass {
		fmt.Println("Overall: PASS - System ready for production load")
	} else {
		fmt.Println("Overall: FAIL - See report for recommendations")
	}
}