im_server/loadtest/im-pressure/coordinator_test.go

package main

import (
	"bytes"
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"
)

func TestCoordinatorResetsBarrierWhenConnectionCountDrops(t *testing.T) {
	state := &coordinatorState{
		start: StartInfo{RunID: "run-reset", Scenario: "group"}, startDelay: time.Minute,
		commitDelay: time.Second, expectedAgents: 2, ready: make(map[string]AgentReady),
		armed: make(map[string]bool), ledgers: make(map[string][]AcceptedMessage),
		started:   make(map[string]time.Time),
		summaries: make(map[string]AgentSummary), complete: make(chan struct{}),
	}
	postReady := func(ready AgentReady) StartInfo {
		raw, err := json.Marshal(ready)
		if err != nil {
			t.Fatal(err)
		}
		req := httptest.NewRequest(http.MethodPost, "/v1/ready", bytes.NewReader(raw))
		w := httptest.NewRecorder()
		state.handleReady(w, req)
		if w.Code != http.StatusOK {
			t.Fatalf("ready status=%d body=%s", w.Code, w.Body.String())
		}
		var start StartInfo
		if err = json.Unmarshal(w.Body.Bytes(), &start); err != nil {
			t.Fatal(err)
		}
		return start
	}
	fullA := AgentReady{Machine: "a", ReadyConnections: 1000, ExpectedConnections: 1000}
	fullB := AgentReady{Machine: "b", ReadyConnections: 1000, ExpectedConnections: 1000}
	postReady(fullA)
	if start := postReady(fullB); start.StartAt.IsZero() {
		t.Fatal("full agents must create a validation barrier")
	}
	fullA.ReadyConnections = 999
	if start := postReady(fullA); !start.StartAt.IsZero() || start.Committed {
		t.Fatalf("connection drop must reset barrier: %+v", start)
	}
}

func TestCoordinatorClientWaitsForBothAgentsAndMergesLedger(t *testing.T) {
	state := &coordinatorState{
		start: StartInfo{RunID: "run1", Scenario: "p2p"}, startDelay: 50 * time.Millisecond,
		commitDelay: 50 * time.Millisecond, expectedAgents: 2,
		ready: make(map[string]AgentReady), ledgers: make(map[string][]AcceptedMessage),
		armed: make(map[string]bool), summaries: make(map[string]AgentSummary), complete: make(chan struct{}),
		started: make(map[string]time.Time),
	}
	server := httptest.NewServer(testCoordinatorMux(state))
	defer server.Close()
	client := newCoordinatorClient(server.URL)

	aDone := make(chan StartInfo, 1)
	go func() {
		start, _ := client.ready(context.Background(), func() AgentReady {
			return AgentReady{Machine: "a", ReadyConnections: 1000, ExpectedConnections: 1000}
		})
		aDone <- start
	}()
	time.Sleep(50 * time.Millisecond)
	startB, err := client.ready(context.Background(), func() AgentReady {
		return AgentReady{Machine: "b", ReadyConnections: 1000, ExpectedConnections: 1000}
	})
	if err != nil || startB.StartAt.IsZero() || !startB.Committed {
		t.Fatalf("machine B ready: start=%+v err=%v", startB, err)
	}
	select {
	case startA := <-aDone:
		if !startA.StartAt.Equal(startB.StartAt) {
			t.Fatalf("different start times A=%s B=%s", startA.StartAt, startB.StartAt)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("machine A did not receive start barrier")
	}

	if err = client.submitLedger(context.Background(), "a", []AcceptedMessage{{ClientMid: "b"}}); err != nil {
		t.Fatal(err)
	}
	if err = client.submitLedger(context.Background(), "b", []AcceptedMessage{{ClientMid: "a"}}); err != nil {
		t.Fatal(err)
	}
	ledger, err := client.ledger(context.Background())
	if err != nil || len(ledger) != 2 || ledger[0].ClientMid != "a" || ledger[1].ClientMid != "b" {
		t.Fatalf("ledger=%+v err=%v", ledger, err)
	}
	if err = client.submitObserved(context.Background(), "a", []AcceptedMessage{{ClientMid: "unknown", Seq: 10}}); err != nil {
		t.Fatal(err)
	}
	if err = client.submitObserved(context.Background(), "b", []AcceptedMessage{{ClientMid: "unknown", Seq: 10}}); err != nil {
		t.Fatal(err)
	}
	observed, err := client.observed(context.Background())
	if err != nil || len(observed) != 1 || observed[0].ClientMid != "unknown" {
		t.Fatalf("observed=%+v err=%v", observed, err)
	}
	if err = client.submitHistory(context.Background(), "a", []string{"mid-2", "mid-1"}); err != nil {
		t.Fatal(err)
	}
	if err = client.submitHistory(context.Background(), "b", []string{"mid-2", "mid-3"}); err != nil {
		t.Fatal(err)
	}
	history, err := client.history(context.Background())
	if err != nil || len(history) != 3 || history[0] != "mid-1" || history[1] != "mid-2" || history[2] != "mid-3" {
		t.Fatalf("history=%+v err=%v", history, err)
	}
}

func testCoordinatorMux(state *coordinatorState) *http.ServeMux {
	mux := http.NewServeMux()
	mux.HandleFunc("/v1/ready", state.handleReady)
	mux.HandleFunc("/v1/start", state.handleStart)
	mux.HandleFunc("/v1/started", state.handleStarted)
	mux.HandleFunc("/v1/ledger", state.handleLedger)
	mux.HandleFunc("/v1/observed", state.handleObserved)
	mux.HandleFunc("/v1/history", state.handleHistory)
	mux.HandleFunc("/v1/summary", state.handleSummary)
	mux.HandleFunc("/v1/status", state.handleStatus)
	return mux
}

func TestCoordinatorRejectsLateAgentForBothMachines(t *testing.T) {
	startAt := time.Now().Add(-maxAgentStartSkew - time.Second).UTC()
	state := &coordinatorState{
		start:          StartInfo{RunID: "run-late", Scenario: "group", StartAt: startAt, Committed: true},
		expectedAgents: 2, started: make(map[string]time.Time),
	}
	server := httptest.NewServer(testCoordinatorMux(state))
	defer server.Close()
	client := newCoordinatorClient(server.URL)
	if err := client.started(context.Background(), "b", startAt); err == nil {
		t.Fatal("late agent must be rejected")
	}
	if err := client.started(context.Background(), "a", startAt); err == nil {
		t.Fatal("the other agent must observe the same rejected run")
	}
}

func TestCoordinatorReleasesAgentsOnlyAfterBothReachStart(t *testing.T) {
	startAt := time.Now().UTC()
	state := &coordinatorState{
		start:          StartInfo{RunID: "run-synced", Scenario: "group", StartAt: startAt, Committed: true},
		expectedAgents: 2, started: make(map[string]time.Time),
	}
	server := httptest.NewServer(testCoordinatorMux(state))
	defer server.Close()
	client := newCoordinatorClient(server.URL)
	aDone := make(chan error, 1)
	go func() { aDone <- client.started(context.Background(), "a", startAt) }()
	time.Sleep(50 * time.Millisecond)
	select {
	case err := <-aDone:
		t.Fatalf("machine A passed before B arrived: %v", err)
	default:
	}
	if err := client.started(context.Background(), "b", startAt); err != nil {
		t.Fatalf("machine B start: %v", err)
	}
	select {
	case err := <-aDone:
		if err != nil {
			t.Fatalf("machine A start: %v", err)
		}
	case <-time.After(time.Second):
		t.Fatal("machine A was not released")
	}
}