feat: add dual-host IM pressure test tool

loadtest/im-pressure/collector.go

@@ -0,0 +1,114 @@
+package main
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+)
+
+type namedURLs []string
+
+func (n *namedURLs) String() string { return strings.Join(*n, ",") }
+func (n *namedURLs) Set(value string) error {
+	if !strings.Contains(value, "=") {
+		return fmt.Errorf("URL must use name=https://host/path")
+	}
+	*n = append(*n, value)
+	return nil
+}
+
+type metricSnapshot struct {
+	Timestamp time.Time `json:"timestamp"`
+	Name      string    `json:"name"`
+	URL       string    `json:"url"`
+	Status    int       `json:"status"`
+	Duration  int64     `json:"duration_ms"`
+	Body      string    `json:"body,omitempty"`
+	Error     string    `json:"error,omitempty"`
+}
+
+func runCollector(ctx context.Context, sources namedURLs, output string, interval, duration time.Duration) error {
+	if len(sources) == 0 {
+		return fmt.Errorf("at least one --url name=url is required")
+	}
+	if interval <= 0 || duration <= 0 {
+		return fmt.Errorf("interval and duration must be positive")
+	}
+	if dir := filepath.Dir(output); dir != "." {
+		if err := os.MkdirAll(dir, 0o750); err != nil {
+			return err
+		}
+	}
+	file, err := os.OpenFile(output, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+	writer := bufio.NewWriterSize(file, 256<<10)
+	defer writer.Flush()
+	client := &http.Client{Timeout: 10 * time.Second}
+	deadline := time.Now().Add(duration)
+	for {
+		snapshots := collectOnce(ctx, client, sources)
+		for _, snapshot := range snapshots {
+			raw, _ := json.Marshal(snapshot)
+			if _, err = writer.Write(append(raw, '\n')); err != nil {
+				return err
+			}
+		}
+		if err = writer.Flush(); err != nil {
+			return err
+		}
+		if time.Now().Add(interval).After(deadline) {
+			return nil
+		}
+		if err = waitContext(ctx, interval); err != nil {
+			return err
+		}
+	}
+}
+
+func collectOnce(ctx context.Context, client *http.Client, sources namedURLs) []metricSnapshot {
+	out := make([]metricSnapshot, len(sources))
+	var wg sync.WaitGroup
+	for i, source := range sources {
+		i, source := i, source
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			parts := strings.SplitN(source, "=", 2)
+			snapshot := metricSnapshot{Timestamp: time.Now().UTC(), Name: parts[0], URL: parts[1]}
+			started := time.Now()
+			req, err := http.NewRequestWithContext(ctx, http.MethodGet, parts[1], nil)
+			if err == nil {
+				var response *http.Response
+				response, err = client.Do(req)
+				if err == nil {
+					snapshot.Status = response.StatusCode
+					raw, readErr := io.ReadAll(io.LimitReader(response.Body, 16<<20))
+					_ = response.Body.Close()
+					if readErr != nil {
+						err = readErr
+					} else {
+						snapshot.Body = string(raw)
+					}
+				}
+			}
+			snapshot.Duration = time.Since(started).Milliseconds()
+			if err != nil {
+				snapshot.Error = err.Error()
+			}
+			out[i] = snapshot
+		}()
+	}
+	wg.Wait()
+	return out
+}