324 lines
9.2 KiB
Go
324 lines
9.2 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"`
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|