im_server/loadtest/im-pressure/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"
)

func main() {
	log.SetFlags(log.LstdFlags | log.Lmicroseconds | log.LUTC)
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	if len(os.Args) < 2 {
		usage()
		os.Exit(2)
	}
	var err error
	switch os.Args[1] {
	case "prepare":
		err = prepareCommand(ctx, os.Args[2:])
	case "verify":
		err = verifyCommand(ctx, os.Args[2:])
	case "coordinator":
		err = coordinatorCommand(ctx, os.Args[2:])
	case "agent":
		err = agentCommand(ctx, os.Args[2:])
	case "merge":
		err = mergeCommand(os.Args[2:])
	case "collect":
		err = collectCommand(ctx, os.Args[2:])
	case "verify-db":
		err = verifyDBCommand(ctx, os.Args[2:])
	default:
		usage()
		os.Exit(2)
	}
	if err != nil {
		log.Fatal(err)
	}
}

func verifyDBCommand(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("verify-db", flag.ContinueOnError)
	coordinator := fs.String("coordinator", "loadtest/im-pressure/results/coordinator.json", "coordinator ledger path")
	output := fs.String("output", "-", "verification JSON output path, or - for stdout")
	tenant := fs.String("tenant", "", "tenant code")
	dsnEnv := fs.String("dsn-env", "MYSQL_DSN", "environment variable containing the business database DSN")
	registryDB := fs.String("registry-db", "", "tenant registry database; defaults to TENANT_REGISTRY_DB")
	if err := fs.Parse(args); err != nil {
		return err
	}
	if strings.TrimSpace(*tenant) == "" {
		return fmt.Errorf("--tenant is required")
	}
	return runDBVerify(ctx, *coordinator, *output, *tenant, *dsnEnv, *registryDB)
}

func prepareCommand(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("prepare", flag.ContinueOnError)
	configPath := fs.String("config", "loadtest/im-pressure/config.example.json", "JSON configuration path")
	manifestPath := fs.String("manifest", "loadtest/im-pressure/manifest.json", "manifest output path")
	if err := fs.Parse(args); err != nil {
		return err
	}
	cfg, err := loadConfig(*configPath)
	if err != nil {
		return err
	}
	return runPrepare(ctx, cfg, *manifestPath)
}

func verifyCommand(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("verify", flag.ContinueOnError)
	configPath := fs.String("config", "loadtest/im-pressure/config.example.json", "JSON configuration path")
	manifestPath := fs.String("manifest", "loadtest/im-pressure/manifest.json", "manifest path")
	if err := fs.Parse(args); err != nil {
		return err
	}
	cfg, err := loadConfig(*configPath)
	if err != nil {
		return err
	}
	manifest, err := loadManifest(*manifestPath)
	if err != nil {
		return err
	}
	apiKey, err := cfg.apiKey()
	if err != nil {
		return err
	}
	return runVerify(ctx, cfg, manifest, newAPIClient(cfg, apiKey))
}

func coordinatorCommand(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("coordinator", flag.ContinueOnError)
	listen := fs.String("listen", ":8090", "coordinator listen address")
	runID := fs.String("run-id", "", "unique run ID, max 24 characters")
	scenario := fs.String("scenario", "", "p2p, group, mixed, or idle")
	output := fs.String("output", "loadtest/im-pressure/results/coordinator.json", "coordinator result path")
	expected := fs.Int("expected-agents", 2, "number of agents")
	delay := fs.Duration("start-delay", 30*time.Second, "delay after all agents are ready")
	if err := fs.Parse(args); err != nil {
		return err
	}
	return runCoordinator(ctx, *listen, *runID, *scenario, *output, *expected, *delay)
}

func agentCommand(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("agent", flag.ContinueOnError)
	configPath := fs.String("config", "loadtest/im-pressure/config.example.json", "JSON configuration path")
	manifestPath := fs.String("manifest", "loadtest/im-pressure/manifest.json", "manifest path")
	machine := fs.String("machine", "", "a or b")
	coordinatorURL := fs.String("coordinator", "", "coordinator base URL")
	output := fs.String("output", "loadtest/im-pressure/results/agent-summary.json", "agent result path")
	metricsListen := fs.String("metrics-listen", ":9105", "Prometheus listen address; empty disables it")
	if err := fs.Parse(args); err != nil {
		return err
	}
	if *coordinatorURL == "" {
		return fmt.Errorf("--coordinator is required")
	}
	cfg, err := loadConfig(*configPath)
	if err != nil {
		return err
	}
	manifest, err := loadManifest(*manifestPath)
	if err != nil {
		return err
	}
	return runAgent(ctx, cfg, manifest, *machine, *coordinatorURL, *output, *metricsListen)
}

func mergeCommand(args []string) error {
	fs := flag.NewFlagSet("merge", flag.ContinueOnError)
	a := fs.String("a", "loadtest/im-pressure/results/agent-a-summary.json", "machine A summary")
	b := fs.String("b", "loadtest/im-pressure/results/agent-b-summary.json", "machine B summary")
	output := fs.String("output", "loadtest/im-pressure/results/report.json", "combined report")
	if err := fs.Parse(args); err != nil {
		return err
	}
	return runMerge(*a, *b, *output)
}

func collectCommand(ctx context.Context, args []string) error {
	fs := flag.NewFlagSet("collect", flag.ContinueOnError)
	var urls namedURLs
	fs.Var(&urls, "url", "repeatable name=https://host/metrics source")
	output := fs.String("output", "loadtest/im-pressure/results/metrics.jsonl", "snapshot output")
	interval := fs.Duration("interval", 5*time.Second, "collection interval")
	duration := fs.Duration("duration", 10*time.Minute, "collection duration")
	coordinatorURL := fs.String("coordinator", "", "optional coordinator base URL used to label before/during phases")
	coordinatorOutput := fs.String("coordinator-output", "", "coordinator JSON path used to detect the after phase")
	if err := fs.Parse(args); err != nil {
		return err
	}
	return runCollector(ctx, urls, *output, *interval, *duration, *coordinatorURL, *coordinatorOutput)
}

func usage() {
	fmt.Fprintln(os.Stderr, "usage: im-pressure <prepare|verify|coordinator|agent|merge|collect|verify-db> [flags]")
}