feat: add dual-host IM pressure test tool
This commit is contained in:
richjohnrambo
@@ -0,0 +1,323 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type StartInfo struct {
|
||||
RunID string `json:"run_id"`
|
||||
Scenario string `json:"scenario"`
|
||||
StartAt time.Time `json:"start_at"`
|
||||
}
|
||||
|
||||
type AgentReady struct {
|
||||
Machine string `json:"machine"`
|
||||
ReadyConnections int `json:"ready_connections"`
|
||||
}
|
||||
|
||||
type coordinatorState struct {
|
||||
mu sync.Mutex
|
||||
start StartInfo
|
||||
startDelay time.Duration
|
||||
expectedAgents int
|
||||
ready map[string]AgentReady
|
||||
ledgers map[string][]AcceptedMessage
|
||||
summaries map[string]AgentSummary
|
||||
complete chan struct{}
|
||||
completeOnce sync.Once
|
||||
}
|
||||
|
||||
func runCoordinator(ctx context.Context, listen, runID, scenario, output string, expectedAgents int, startDelay time.Duration) error {
|
||||
if !runIDPattern.MatchString(runID) {
|
||||
return fmt.Errorf("run-id must match %s", runIDPattern.String())
|
||||
}
|
||||
if !validScenario(scenario) {
|
||||
return fmt.Errorf("unsupported scenario %q", scenario)
|
||||
}
|
||||
if expectedAgents <= 0 {
|
||||
expectedAgents = 2
|
||||
}
|
||||
state := &coordinatorState{
|
||||
start: StartInfo{RunID: runID, Scenario: scenario}, startDelay: startDelay, expectedAgents: expectedAgents,
|
||||
ready: make(map[string]AgentReady), ledgers: make(map[string][]AcceptedMessage),
|
||||
summaries: make(map[string]AgentSummary), complete: make(chan struct{}),
|
||||
}
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("/v1/ready", state.handleReady)
|
||||
mux.HandleFunc("/v1/start", state.handleStart)
|
||||
mux.HandleFunc("/v1/ledger", state.handleLedger)
|
||||
mux.HandleFunc("/v1/summary", state.handleSummary)
|
||||
mux.HandleFunc("/v1/status", state.handleStatus)
|
||||
server := &http.Server{Addr: listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second}
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
log.Printf("coordinator listening=%s run_id=%s scenario=%s expected_agents=%d", listen, runID, scenario, expectedAgents)
|
||||
if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
errCh <- err
|
||||
}
|
||||
}()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case err := <-errCh:
|
||||
return err
|
||||
case <-state.complete:
|
||||
log.Printf("all agent summaries received")
|
||||
}
|
||||
shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_ = server.Shutdown(shutdownCtx)
|
||||
if output != "" {
|
||||
if err := state.writeOutput(output); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleReady(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
var ready AgentReady
|
||||
if !decodeLimitedJSON(w, r, &ready) || (ready.Machine != "a" && ready.Machine != "b") || ready.ReadyConnections <= 0 {
|
||||
if ready.Machine != "a" && ready.Machine != "b" {
|
||||
http.Error(w, "invalid machine", http.StatusBadRequest)
|
||||
}
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.ready[ready.Machine] = ready
|
||||
if len(s.ready) == s.expectedAgents && s.start.StartAt.IsZero() {
|
||||
s.start.StartAt = time.Now().UTC().Add(s.startDelay)
|
||||
log.Printf("all agents ready; start_at=%s", s.start.StartAt.Format(time.RFC3339Nano))
|
||||
}
|
||||
start := s.start
|
||||
s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, start)
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleStart(w http.ResponseWriter, r *http.Request) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.start.StartAt.IsZero() {
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{"ready_agents": len(s.ready), "expected_agents": s.expectedAgents})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, s.start)
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleLedger(w http.ResponseWriter, r *http.Request) {
|
||||
machine := strings.ToLower(strings.TrimSpace(r.URL.Query().Get("machine")))
|
||||
switch r.Method {
|
||||
case http.MethodPost:
|
||||
if machine != "a" && machine != "b" {
|
||||
http.Error(w, "invalid machine", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var messages []AcceptedMessage
|
||||
if !decodeLimitedJSON(w, r, &messages) {
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.ledgers[machine] = messages
|
||||
count := len(s.ledgers)
|
||||
s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]int{"submitted_agents": count})
|
||||
case http.MethodGet:
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if len(s.ledgers) < s.expectedAgents {
|
||||
writeJSON(w, http.StatusAccepted, map[string]int{"submitted_agents": len(s.ledgers), "expected_agents": s.expectedAgents})
|
||||
return
|
||||
}
|
||||
var merged []AcceptedMessage
|
||||
for _, messages := range s.ledgers {
|
||||
merged = append(merged, messages...)
|
||||
}
|
||||
sort.Slice(merged, func(i, j int) bool { return merged[i].ClientMid < merged[j].ClientMid })
|
||||
writeJSON(w, http.StatusOK, merged)
|
||||
default:
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleSummary(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
||||
return
|
||||
}
|
||||
var summary AgentSummary
|
||||
if !decodeLimitedJSON(w, r, &summary) {
|
||||
return
|
||||
}
|
||||
if summary.Machine != "a" && summary.Machine != "b" {
|
||||
http.Error(w, "invalid machine", http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.summaries[summary.Machine] = summary
|
||||
complete := len(s.summaries) == s.expectedAgents
|
||||
s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]any{"complete": complete})
|
||||
if complete {
|
||||
s.completeOnce.Do(func() {
|
||||
go func() {
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
close(s.complete)
|
||||
}()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *coordinatorState) handleStatus(w http.ResponseWriter, _ *http.Request) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"start": s.start, "ready": s.ready, "ledger_agents": len(s.ledgers), "summary_agents": len(s.summaries),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *coordinatorState) writeOutput(path string) error {
|
||||
s.mu.Lock()
|
||||
data := struct {
|
||||
Start StartInfo `json:"start"`
|
||||
Ready map[string]AgentReady `json:"ready"`
|
||||
Ledgers map[string][]AcceptedMessage `json:"ledgers"`
|
||||
Summaries map[string]AgentSummary `json:"summaries"`
|
||||
}{s.start, s.ready, s.ledgers, s.summaries}
|
||||
s.mu.Unlock()
|
||||
raw, err := json.MarshalIndent(data, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err = os.MkdirAll(filepath.Dir(path), 0o750); err != nil && filepath.Dir(path) != "." {
|
||||
return err
|
||||
}
|
||||
return os.WriteFile(path, append(raw, '\n'), 0o600)
|
||||
}
|
||||
|
||||
func decodeLimitedJSON(w http.ResponseWriter, r *http.Request, output any) bool {
|
||||
dec := json.NewDecoder(io.LimitReader(r.Body, 32<<20))
|
||||
if err := dec.Decode(output); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, value any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(value)
|
||||
}
|
||||
|
||||
type coordinatorClient struct {
|
||||
base string
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func newCoordinatorClient(base string) *coordinatorClient {
|
||||
return &coordinatorClient{base: strings.TrimRight(base, "/"), client: &http.Client{Timeout: 35 * time.Second}}
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) ready(ctx context.Context, ready AgentReady) (StartInfo, error) {
|
||||
var start StartInfo
|
||||
if _, err := c.request(ctx, http.MethodPost, "/v1/ready", ready, &start); err != nil {
|
||||
return start, err
|
||||
}
|
||||
for start.StartAt.IsZero() {
|
||||
if err := waitContext(ctx, time.Second); err != nil {
|
||||
return start, err
|
||||
}
|
||||
status, err := c.request(ctx, http.MethodGet, "/v1/start", nil, &start)
|
||||
if err != nil {
|
||||
return start, err
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
start = StartInfo{}
|
||||
}
|
||||
}
|
||||
return start, nil
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) submitLedger(ctx context.Context, machine string, messages []AcceptedMessage) error {
|
||||
_, err := c.request(ctx, http.MethodPost, "/v1/ledger?machine="+machine, messages, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) ledger(ctx context.Context) ([]AcceptedMessage, error) {
|
||||
for {
|
||||
var messages []AcceptedMessage
|
||||
status, err := c.request(ctx, http.MethodGet, "/v1/ledger", nil, &messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status == http.StatusOK {
|
||||
return messages, nil
|
||||
}
|
||||
if err = waitContext(ctx, time.Second); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) submitSummary(ctx context.Context, summary AgentSummary) error {
|
||||
_, err := c.request(ctx, http.MethodPost, "/v1/summary", summary, nil)
|
||||
return err
|
||||
}
|
||||
|
||||
func (c *coordinatorClient) request(ctx context.Context, method, path string, input, output any) (int, error) {
|
||||
var body io.Reader
|
||||
if input != nil {
|
||||
raw, err := json.Marshal(input)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
body = strings.NewReader(string(raw))
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, method, c.base+path, body)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if input != nil {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
resp, err := c.client.Do(req)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode >= 400 {
|
||||
raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
|
||||
return resp.StatusCode, fmt.Errorf("coordinator HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
|
||||
}
|
||||
if output != nil && resp.StatusCode == http.StatusOK {
|
||||
if err = json.NewDecoder(resp.Body).Decode(output); err != nil {
|
||||
return resp.StatusCode, err
|
||||
}
|
||||
}
|
||||
return resp.StatusCode, nil
|
||||
}
|
||||
|
||||
func waitContext(ctx context.Context, delay time.Duration) error {
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user