Agent-System-Dev/tests/load/collaboration/sse_probe/main.go

// 协作存储 U7：SSE 持久事件可见延迟探针（负载合同：P95≤1s）。
// k6 核心不支持 SSE 流式读取，此项由本 Go 探针承担；与 k6 结果记录在
// 同一目录（run_load.sh 编排）。
//
// 行为：并发 C 个会话各挂一个 SSE 订阅（共享登录 cookie）；写者按固定
// 间隔轮转 POST 消息（mock 即时回合）；统计「POST 前基准 seq → 会话内
// 新帧到达」延迟分布与 seq 单调性，输出 JSON 报告并以退出码执行
// P95≤1s、零乱序断言（数据不变量零违反，方案 §3）。
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/http/cookiejar"
	"os"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

type sample struct {
	Conversation string  `json:"conversation"`
	SendSeq      int64   `json:"sendSeq"`
	LatencyMS    float64 `json:"latencyMs"`
}

type report struct {
	Target      string   `json:"target"`
	Duration    string   `json:"duration"`
	Subscribers int      `json:"subscribers"`
	Messages    int      `json:"messages"`
	Samples     int      `json:"samples"`
	P50MS       float64  `json:"p50Ms"`
	P95MS       float64  `json:"p95Ms"`
	P99MS       float64  `json:"p99Ms"`
	MaxMS       float64  `json:"maxMs"`
	OutOfOrder  int      `json:"outOfOrder"` // seq 乱序/回退（不变量违反）
	ThresholdMS float64  `json:"thresholdMs"`
	Pass        bool     `json:"pass"`
	Detail      []sample `json:"samplesDetail"`
}

var (
	jar, _       = cookiejar.New(nil)
	client       = &http.Client{Timeout: 10 * time.Second, Jar: jar}
	streamClient = &http.Client{Jar: jar} // SSE：无整体超时，stop 控制
	pendingMu    sync.Mutex
	pending      = map[string]pendingSend{}
	outOfOrderMu sync.Mutex
	outOfOrderN  int
	samplesMu    sync.Mutex
	samples      []sample
)

type pendingSend struct {
	seq  int64
	sent time.Time
}

func main() {
	base := flag.String("base", "http://127.0.0.1:18980", "目标实例基址")
	owner := flag.String("owner", "probe_user", "探针用户")
	password := flag.String("password", "probe-user-pass-Aa9", "探针用户口令")
	convs := flag.Int("convs", 10, "并发订阅会话数")
	every := flag.Duration("every", 500*time.Millisecond, "写者间隔（轮转各会话）")
	dur := flag.Duration("dur", 60*time.Second, "采样时长")
	out := flag.String("out", "", "JSON 报告输出路径（空 = stdout）")
	flag.Parse()

	login(*base, *owner, *password)

	projectID := apiJSON("POST", *base, "/api/projects", map[string]any{"name": "sse-probe"})["id"].(string)
	convIDs := make([]string, *convs)
	for i := range convIDs {
		convIDs[i] = apiJSON("POST", *base, "/api/projects/"+projectID+"/conversations", map[string]any{"title": fmt.Sprintf("probe-%d", i)})["id"].(string)
	}

	stop := make(chan struct{})
	var wg sync.WaitGroup
	for _, cid := range convIDs {
		wg.Add(1)
		go func(cid string) {
			defer wg.Done()
			subscribe(*base, cid, stop)
		}(cid)
	}
	time.Sleep(500 * time.Millisecond) // 订阅就绪

	deadline := time.Now().Add(*dur)
	sent := 0
	for i := 0; time.Now().Before(deadline); i++ {
		cid := convIDs[i%len(convIDs)]
		detail := apiJSON("GET", *base, "/api/conversations/"+cid, nil)
		seq, _ := detail["lastSeq"].(float64)
		pendingMu.Lock()
		pending[cid] = pendingSend{seq: int64(seq), sent: time.Now()}
		pendingMu.Unlock()
		apiJSON("POST", *base, "/api/conversations/"+cid+"/messages", map[string]any{"text": fmt.Sprintf("probe %d", i)})
		sent++
		time.Sleep(*every)
	}

	time.Sleep(3 * time.Second) // 最后一批帧在途余量
	close(stop)
	wg.Wait()

	r := buildReport(*base, *dur, *convs, sent)
	data, _ := json.MarshalIndent(r, "", "  ")
	if *out != "" {
		if err := os.WriteFile(*out, data, 0o644); err != nil {
			fmt.Fprintln(os.Stderr, "write report:", err)
			os.Exit(2)
		}
		fmt.Println("report:", *out)
	} else {
		fmt.Println(string(data))
	}
	if !r.Pass {
		fmt.Fprintf(os.Stderr, "sse probe FAIL: p95=%.1fms (≤%.0fms) samples=%d outOfOrder=%d\n", r.P95MS, r.ThresholdMS, r.Samples, r.OutOfOrder)
		os.Exit(1)
	}
	fmt.Printf("sse probe PASS: p95=%.1fms samples=%d\n", r.P95MS, r.Samples)
}

func login(base, owner, pass string) {
	body, _ := json.Marshal(map[string]string{"owner": owner, "password": pass})
	resp, err := client.Post(base+"/api/auth/login", "application/json", bytes.NewReader(body))
	if err != nil || resp.StatusCode != 200 {
		fmt.Fprintf(os.Stderr, "probe login: err=%v status=%v\n", err, statusOf(resp))
		os.Exit(2)
	}
	resp.Body.Close()
}

func statusOf(resp *http.Response) string {
	if resp == nil {
		return "<nil>"
	}
	return strconv.Itoa(resp.StatusCode)
}

func apiJSON(method, base, path string, body any) map[string]any {
	var rd io.Reader
	_ = rd
	if body != nil {
		buf, _ := json.Marshal(body)
		rd = bytes.NewReader(buf)
	}
	req, _ := http.NewRequest(method, base+path, rd)
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := client.Do(req)
	if err != nil {
		fmt.Fprintf(os.Stderr, "%s %s: %v\n", method, path, err)
		os.Exit(2)
	}
	defer resp.Body.Close()
	out := map[string]any{}
	if resp.StatusCode != 200 {
		fmt.Fprintf(os.Stderr, "%s %s: status %d\n", method, path, resp.StatusCode)
		os.Exit(2)
	}
	if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
		fmt.Fprintf(os.Stderr, "%s %s: decode %v\n", method, path, err)
		os.Exit(2)
	}
	return out
}

// subscribe 读会话 SSE：seq 严格递增校验；帧到达时与写者基准配对成延迟样本。
func subscribe(base, cid string, stop <-chan struct{}) {
	req, _ := http.NewRequest(http.MethodGet, base+"/api/events?conversationId="+cid, nil)
	resp, err := streamClient.Do(req)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	sc := bufio.NewScanner(resp.Body)
	sc.Buffer(make([]byte, 0, 64*1024), 4<<20)
	last := int64(-1)
	for sc.Scan() {
		select {
		case <-stop:
			return
		default:
		}
		line := sc.Text()
		if !strings.HasPrefix(line, "id: ") {
			continue
		}
		seq, _ := strconv.ParseInt(strings.TrimPrefix(line, "id: "), 10, 64)
		arrived := time.Now()
		if last != -1 && seq <= last {
			outOfOrderMu.Lock()
			outOfOrderN++
			outOfOrderMu.Unlock()
		}
		last = seq
		pendingMu.Lock()
		p, ok := pending[cid]
		pendingMu.Unlock()
		if ok && seq > p.seq {
			pendingMu.Lock()
			delete(pending, cid)
			pendingMu.Unlock()
			samplesMu.Lock()
			samples = append(samples, sample{Conversation: cid, SendSeq: seq, LatencyMS: float64(arrived.Sub(p.sent)) / float64(time.Millisecond)})
			samplesMu.Unlock()
		}
	}
}

func buildReport(base string, dur time.Duration, subs, msgs int) report {
	samplesMu.Lock()
	detail := append([]sample(nil), samples...)
	samplesMu.Unlock()
	outOfOrderMu.Lock()
	ooo := outOfOrderN
	outOfOrderMu.Unlock()

	r := report{Target: base, Duration: dur.String(), Subscribers: subs, Messages: msgs, Samples: len(detail), OutOfOrder: ooo, ThresholdMS: 1000, Detail: detail}
	if len(detail) > 0 {
		vals := make([]float64, len(detail))
		for i, s := range detail {
			vals[i] = s.LatencyMS
		}
		sort.Float64s(vals)
		pct := func(p float64) float64 { return vals[int(float64(len(vals)-1)*p)] }
		r.P50MS, r.P95MS, r.P99MS, r.MaxMS = pct(0.50), pct(0.95), pct(0.99), vals[len(vals)-1]
	}
	r.Pass = len(detail) > 0 && r.P95MS <= r.ThresholdMS && ooo == 0
	return r
}