512 lines
15 KiB
Go
512 lines
15 KiB
Go
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"`
|
|
Committed bool `json:"committed"`
|
|
}
|
|
|
|
type AgentReady struct {
|
|
Machine string `json:"machine"`
|
|
ReadyConnections int `json:"ready_connections"`
|
|
ExpectedConnections int `json:"expected_connections"`
|
|
Armed bool `json:"armed"`
|
|
}
|
|
|
|
type coordinatorState struct {
|
|
mu sync.Mutex
|
|
start StartInfo
|
|
startDelay time.Duration
|
|
commitDelay time.Duration
|
|
expectedAgents int
|
|
ready map[string]AgentReady
|
|
armed map[string]bool
|
|
ledgers map[string][]AcceptedMessage
|
|
observed map[string][]AcceptedMessage
|
|
history map[string][]string
|
|
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, commitDelay: 5 * time.Second,
|
|
expectedAgents: expectedAgents,
|
|
ready: make(map[string]AgentReady), ledgers: make(map[string][]AcceptedMessage),
|
|
observed: make(map[string][]AcceptedMessage), history: make(map[string][]string),
|
|
armed: make(map[string]bool), 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/observed", state.handleObserved)
|
|
mux.HandleFunc("/v1/history", state.handleHistory)
|
|
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) handleHistory(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 messageIDs []string
|
|
if !decodeLimitedJSON(w, r, &messageIDs) {
|
|
return
|
|
}
|
|
s.mu.Lock()
|
|
if s.history == nil {
|
|
s.history = make(map[string][]string)
|
|
}
|
|
s.history[machine] = messageIDs
|
|
count := len(s.history)
|
|
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.history) < s.expectedAgents {
|
|
writeJSON(w, http.StatusAccepted, map[string]int{"submitted_agents": len(s.history), "expected_agents": s.expectedAgents})
|
|
return
|
|
}
|
|
unique := make(map[string]bool)
|
|
for _, messageIDs := range s.history {
|
|
for _, mid := range messageIDs {
|
|
if mid != "" {
|
|
unique[mid] = true
|
|
}
|
|
}
|
|
}
|
|
merged := make([]string, 0, len(unique))
|
|
for mid := range unique {
|
|
merged = append(merged, mid)
|
|
}
|
|
sort.Strings(merged)
|
|
writeJSON(w, http.StatusOK, merged)
|
|
default:
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
}
|
|
}
|
|
|
|
func (s *coordinatorState) handleObserved(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()
|
|
if s.observed == nil {
|
|
s.observed = make(map[string][]AcceptedMessage)
|
|
}
|
|
s.observed[machine] = messages
|
|
count := len(s.observed)
|
|
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.observed) < s.expectedAgents {
|
|
writeJSON(w, http.StatusAccepted, map[string]int{"submitted_agents": len(s.observed), "expected_agents": s.expectedAgents})
|
|
return
|
|
}
|
|
var merged []AcceptedMessage
|
|
for _, messages := range s.observed {
|
|
merged = append(merged, messages...)
|
|
}
|
|
merged = mergeMessageLedgers(nil, merged)
|
|
writeJSON(w, http.StatusOK, merged)
|
|
default:
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
}
|
|
}
|
|
|
|
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.ExpectedConnections <= 0 || ready.ReadyConnections < 0 || ready.ReadyConnections > ready.ExpectedConnections {
|
|
if ready.Machine != "a" && ready.Machine != "b" {
|
|
http.Error(w, "invalid machine", http.StatusBadRequest)
|
|
} else {
|
|
http.Error(w, "invalid connection counts", http.StatusBadRequest)
|
|
}
|
|
return
|
|
}
|
|
s.mu.Lock()
|
|
s.ready[ready.Machine] = ready
|
|
now := time.Now().UTC()
|
|
allStable := len(s.ready) == s.expectedAgents
|
|
for _, agent := range s.ready {
|
|
if agent.ReadyConnections != agent.ExpectedConnections {
|
|
allStable = false
|
|
break
|
|
}
|
|
}
|
|
if !allStable && (!s.start.Committed || now.Before(s.start.StartAt)) {
|
|
log.Printf("agent readiness dropped before start; resetting start barrier")
|
|
s.start.StartAt = time.Time{}
|
|
s.start.Committed = false
|
|
clear(s.armed)
|
|
}
|
|
if allStable && s.start.StartAt.IsZero() {
|
|
s.start.StartAt = now.Add(s.startDelay)
|
|
s.start.Committed = false
|
|
clear(s.armed)
|
|
log.Printf("all agents ready; validation_at=%s", s.start.StartAt.Format(time.RFC3339Nano))
|
|
}
|
|
if allStable && !s.start.Committed && ready.Armed && !now.Before(s.start.StartAt) {
|
|
s.armed[ready.Machine] = true
|
|
}
|
|
if allStable && !s.start.Committed && len(s.armed) == s.expectedAgents {
|
|
delay := s.commitDelay
|
|
if delay <= 0 {
|
|
delay = 5 * time.Second
|
|
}
|
|
s.start.StartAt = now.Add(delay)
|
|
s.start.Committed = true
|
|
log.Printf("all agents validated; committed 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, "armed": s.armed,
|
|
"ledger_agents": len(s.ledgers), "observed_agents": len(s.observed),
|
|
"history_agents": len(s.history), "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"`
|
|
Observed map[string][]AcceptedMessage `json:"observed"`
|
|
History map[string][]string `json:"history"`
|
|
Summaries map[string]AgentSummary `json:"summaries"`
|
|
}{
|
|
Start: s.start, Ready: s.ready, Ledgers: s.ledgers, Observed: s.observed,
|
|
History: s.history, Summaries: 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, current func() AgentReady) (StartInfo, error) {
|
|
var start StartInfo
|
|
armed := false
|
|
for {
|
|
state := current()
|
|
state.Armed = armed
|
|
if _, err := c.request(ctx, http.MethodPost, "/v1/ready", state, &start); err != nil {
|
|
return start, err
|
|
}
|
|
if start.StartAt.IsZero() {
|
|
armed = false
|
|
} else {
|
|
remaining := time.Until(start.StartAt)
|
|
if start.Committed && remaining <= 0 {
|
|
return start, nil
|
|
}
|
|
if !start.Committed && remaining <= 0 {
|
|
armed = true
|
|
continue
|
|
}
|
|
if remaining < time.Second {
|
|
if err := waitContext(ctx, remaining); err != nil {
|
|
return start, err
|
|
}
|
|
continue
|
|
}
|
|
}
|
|
if err := waitContext(ctx, time.Second); err != nil {
|
|
return start, err
|
|
}
|
|
}
|
|
}
|
|
|
|
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) submitObserved(ctx context.Context, machine string, messages []AcceptedMessage) error {
|
|
_, err := c.request(ctx, http.MethodPost, "/v1/observed?machine="+machine, messages, nil)
|
|
return err
|
|
}
|
|
|
|
func (c *coordinatorClient) observed(ctx context.Context) ([]AcceptedMessage, error) {
|
|
for {
|
|
var messages []AcceptedMessage
|
|
status, err := c.request(ctx, http.MethodGet, "/v1/observed", 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) submitHistory(ctx context.Context, machine string, messageIDs []string) error {
|
|
_, err := c.request(ctx, http.MethodPost, "/v1/history?machine="+machine, messageIDs, nil)
|
|
return err
|
|
}
|
|
|
|
func (c *coordinatorClient) history(ctx context.Context) ([]string, error) {
|
|
for {
|
|
var messageIDs []string
|
|
status, err := c.request(ctx, http.MethodGet, "/v1/history", nil, &messageIDs)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if status == http.StatusOK {
|
|
return messageIDs, 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
|
|
}
|
|
}
|