Test User
471 lines
14 KiB
Go
471 lines
14 KiB
Go
package client
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"math"
|
|
"net/http"
|
|
"net/url"
|
|
"strings"
|
|
"time"
|
|
|
|
"speedtest/internal/probe"
|
|
)
|
|
|
|
const (
|
|
DefaultDownloadBytes = 4 * 1024 * 1024
|
|
DefaultUploadBytes = 2 * 1024 * 1024
|
|
)
|
|
|
|
type Options struct {
|
|
ServerURL string
|
|
Target string
|
|
Duration time.Duration
|
|
Timeout time.Duration
|
|
DownloadBytes int
|
|
UploadBytes int
|
|
QueueInterval time.Duration
|
|
OnSample func(probe.Sample)
|
|
OnQueueWaiting func(string)
|
|
}
|
|
|
|
type createResponse struct {
|
|
ID string `json:"id"`
|
|
Status string `json:"status"`
|
|
QueuePosition int `json:"queuePosition"`
|
|
QueueURL string `json:"queueUrl"`
|
|
DownloadURL string `json:"downloadUrl"`
|
|
UploadURL string `json:"uploadUrl"`
|
|
CompleteURL string `json:"completeUrl"`
|
|
CancelURL string `json:"cancelUrl"`
|
|
}
|
|
|
|
type completeResponse struct {
|
|
Summary probe.Summary `json:"summary"`
|
|
}
|
|
|
|
type queueResponse struct {
|
|
ID string `json:"id"`
|
|
Status string `json:"status"`
|
|
QueuePosition int `json:"queuePosition"`
|
|
}
|
|
|
|
type errorResponse struct {
|
|
Error string `json:"error"`
|
|
}
|
|
|
|
func Run(ctx context.Context, options Options) (probe.Summary, error) {
|
|
options, server, err := normalizeOptions(options)
|
|
if err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
httpClient := http.Client{Timeout: transferTimeout(options.Timeout)}
|
|
|
|
session, err := createSession(ctx, &httpClient, server, options)
|
|
if err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
var samples []probe.Sample
|
|
var startedAt time.Time
|
|
complete := false
|
|
defer func() {
|
|
if !complete {
|
|
cancelCtx, cancel := context.WithTimeout(context.Background(), cancelRequestTimeout(options.Timeout))
|
|
defer cancel()
|
|
_ = cancelSession(cancelCtx, &httpClient, server, session.CancelURL)
|
|
}
|
|
}()
|
|
completePartial := func(original error) (probe.Summary, error) {
|
|
if len(samples) == 0 || startedAt.IsZero() {
|
|
return probe.Summary{}, original
|
|
}
|
|
finishedAt := time.Now().UTC()
|
|
completeCtx, cancel := context.WithTimeout(context.Background(), cancelRequestTimeout(options.Timeout))
|
|
defer cancel()
|
|
if _, completeErr := completeSession(completeCtx, &httpClient, server, session.CompleteURL, startedAt, finishedAt, samples); completeErr != nil {
|
|
return probe.Summary{}, fmt.Errorf("%w; also failed to complete partial test session: %v", original, completeErr)
|
|
}
|
|
complete = true
|
|
return probe.Summary{}, original
|
|
}
|
|
|
|
if err := waitForQueue(ctx, &httpClient, server, session, options); err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
|
|
record := func(sample probe.Sample) {
|
|
samples = append(samples, sample)
|
|
if options.OnSample != nil {
|
|
options.OnSample(sample)
|
|
}
|
|
}
|
|
|
|
startedAt = time.Now().UTC()
|
|
record(measureLatency(ctx, &httpClient, server))
|
|
if err := runPhase(ctx, options.Duration, func() probe.Sample {
|
|
return measureDownload(ctx, &httpClient, server, session.DownloadURL, options.DownloadBytes)
|
|
}, record); err != nil {
|
|
return completePartial(err)
|
|
}
|
|
record(measureLatency(ctx, &httpClient, server))
|
|
if err := runPhase(ctx, options.Duration, func() probe.Sample {
|
|
return measureUpload(ctx, &httpClient, server, session.UploadURL, options.UploadBytes)
|
|
}, record); err != nil {
|
|
return completePartial(err)
|
|
}
|
|
record(measureLatency(ctx, &httpClient, server))
|
|
finishedAt := time.Now().UTC()
|
|
|
|
summary, err := completeSession(ctx, &httpClient, server, session.CompleteURL, startedAt, finishedAt, samples)
|
|
if err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
complete = true
|
|
return summary, nil
|
|
}
|
|
|
|
func normalizeOptions(options Options) (Options, *url.URL, error) {
|
|
if strings.TrimSpace(options.ServerURL) == "" {
|
|
return Options{}, nil, errors.New("server URL is required")
|
|
}
|
|
server, err := url.Parse(strings.TrimRight(options.ServerURL, "/"))
|
|
if err != nil || server.Scheme == "" || server.Host == "" || (server.Scheme != "http" && server.Scheme != "https") {
|
|
return Options{}, nil, errors.New("server URL must be a valid http or https URL")
|
|
}
|
|
if strings.TrimSpace(options.Target) == "" {
|
|
options.Target = server.Host
|
|
}
|
|
if options.Duration <= 0 {
|
|
options.Duration = 30 * time.Second
|
|
}
|
|
if options.Timeout <= 0 {
|
|
options.Timeout = 3 * time.Second
|
|
}
|
|
if options.DownloadBytes <= 0 {
|
|
options.DownloadBytes = DefaultDownloadBytes
|
|
}
|
|
if options.UploadBytes <= 0 {
|
|
options.UploadBytes = DefaultUploadBytes
|
|
}
|
|
if options.QueueInterval <= 0 {
|
|
options.QueueInterval = 5 * time.Second
|
|
}
|
|
return options, server, nil
|
|
}
|
|
|
|
func createSession(ctx context.Context, httpClient *http.Client, server *url.URL, options Options) (createResponse, error) {
|
|
payload := map[string]any{
|
|
"target": options.Target,
|
|
"httpUrl": resolve(server, "/api/config"),
|
|
"durationSeconds": int(math.Max(1, math.Ceil(options.Duration.Seconds()))),
|
|
"timeoutMillis": int(math.Max(100, float64(options.Timeout.Milliseconds()))),
|
|
}
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return createResponse{}, err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, resolve(server, "/api/tests"), bytes.NewReader(body))
|
|
if err != nil {
|
|
return createResponse{}, err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := httpClient.Do(request)
|
|
if err != nil {
|
|
return createResponse{}, err
|
|
}
|
|
defer response.Body.Close()
|
|
|
|
var session createResponse
|
|
if response.StatusCode >= 200 && response.StatusCode < 300 {
|
|
if err := json.NewDecoder(response.Body).Decode(&session); err != nil {
|
|
return createResponse{}, err
|
|
}
|
|
return session, nil
|
|
}
|
|
message := decodeError(response.Body)
|
|
if message == "" {
|
|
message = response.Status
|
|
}
|
|
return createResponse{}, fmt.Errorf("create test session: %s", message)
|
|
}
|
|
|
|
func waitForQueue(ctx context.Context, httpClient *http.Client, server *url.URL, session createResponse, options Options) error {
|
|
if session.QueuePosition <= 0 && session.Status != "queued" {
|
|
return nil
|
|
}
|
|
queueURL := session.QueueURL
|
|
if strings.TrimSpace(queueURL) == "" {
|
|
return errors.New("server queued the test without a queue status URL")
|
|
}
|
|
|
|
position := session.QueuePosition
|
|
for {
|
|
if options.OnQueueWaiting != nil {
|
|
options.OnQueueWaiting(queueMessage(position))
|
|
}
|
|
|
|
timer := time.NewTimer(options.QueueInterval)
|
|
select {
|
|
case <-ctx.Done():
|
|
timer.Stop()
|
|
return ctx.Err()
|
|
case <-timer.C:
|
|
}
|
|
|
|
status, err := fetchQueueStatus(ctx, httpClient, server, queueURL)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if status.Status == "ready" || status.QueuePosition <= 0 {
|
|
return nil
|
|
}
|
|
position = status.QueuePosition
|
|
}
|
|
}
|
|
|
|
func fetchQueueStatus(ctx context.Context, httpClient *http.Client, server *url.URL, path string) (queueResponse, error) {
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodGet, resolve(server, path), nil)
|
|
if err != nil {
|
|
return queueResponse{}, err
|
|
}
|
|
response, err := httpClient.Do(request)
|
|
if err != nil {
|
|
return queueResponse{}, err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
message := decodeError(response.Body)
|
|
if message == "" {
|
|
message = response.Status
|
|
}
|
|
return queueResponse{}, fmt.Errorf("queue status: %s", message)
|
|
}
|
|
var status queueResponse
|
|
if err := json.NewDecoder(response.Body).Decode(&status); err != nil {
|
|
return queueResponse{}, err
|
|
}
|
|
return status, nil
|
|
}
|
|
|
|
func queueMessage(position int) string {
|
|
if position <= 0 {
|
|
return "正在排队，等待其他用户完成..."
|
|
}
|
|
return fmt.Sprintf("正在排队，当前排第 %d 位，等待其他用户完成...", position)
|
|
}
|
|
|
|
func runPhase(ctx context.Context, duration time.Duration, measure func() probe.Sample, record func(probe.Sample)) error {
|
|
deadline := time.Now().Add(duration)
|
|
samples := 0
|
|
for time.Now().Before(deadline) || samples == 0 {
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
record(measure())
|
|
samples++
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func measureLatency(ctx context.Context, httpClient *http.Client, server *url.URL) probe.Sample {
|
|
startedAt := time.Now()
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodGet, resolve(server, "/api/config?cache="+cacheToken()), nil)
|
|
if err != nil {
|
|
return failedSample("latency", startedAt, 0, 0, err)
|
|
}
|
|
response, err := httpClient.Do(request)
|
|
if err != nil {
|
|
return failedSample("latency", startedAt, 0, 0, err)
|
|
}
|
|
defer response.Body.Close()
|
|
_, _ = io.Copy(io.Discard, response.Body)
|
|
elapsed := time.Since(startedAt)
|
|
return probe.Sample{
|
|
At: time.Now().UTC(),
|
|
Kind: "latency",
|
|
Success: response.StatusCode >= 200 && response.StatusCode < 300,
|
|
Latency: elapsed,
|
|
LatencyMS: round2(float64(elapsed.Microseconds()) / 1000),
|
|
Error: statusError(response),
|
|
}
|
|
}
|
|
|
|
func measureDownload(ctx context.Context, httpClient *http.Client, server *url.URL, path string, size int) probe.Sample {
|
|
startedAt := time.Now()
|
|
requestURL := resolve(server, path+"&bytes="+fmt.Sprint(size)+"&cache="+cacheToken())
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil)
|
|
if err != nil {
|
|
return failedSample("download", startedAt, 0, 0, err)
|
|
}
|
|
response, err := httpClient.Do(request)
|
|
if err != nil {
|
|
return failedSample("download", startedAt, 0, 0, err)
|
|
}
|
|
defer response.Body.Close()
|
|
received, _ := io.Copy(io.Discard, response.Body)
|
|
elapsed := time.Since(startedAt)
|
|
return transferSample("download", response, startedAt, elapsed, received)
|
|
}
|
|
|
|
func measureUpload(ctx context.Context, httpClient *http.Client, server *url.URL, path string, size int) probe.Sample {
|
|
startedAt := time.Now()
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, resolve(server, path+"&cache="+cacheToken()), bytes.NewReader(make([]byte, size)))
|
|
if err != nil {
|
|
return failedSample("upload", startedAt, 0, 0, err)
|
|
}
|
|
response, err := httpClient.Do(request)
|
|
if err != nil {
|
|
return failedSample("upload", startedAt, 0, 0, err)
|
|
}
|
|
defer response.Body.Close()
|
|
_, _ = io.Copy(io.Discard, response.Body)
|
|
elapsed := time.Since(startedAt)
|
|
return transferSample("upload", response, startedAt, elapsed, int64(size))
|
|
}
|
|
|
|
func transferSample(kind string, response *http.Response, startedAt time.Time, elapsed time.Duration, bytes int64) probe.Sample {
|
|
return probe.Sample{
|
|
At: time.Now().UTC(),
|
|
Kind: kind,
|
|
Success: response.StatusCode >= 200 && response.StatusCode < 300,
|
|
Latency: elapsed,
|
|
LatencyMS: round2(float64(elapsed.Microseconds()) / 1000),
|
|
Bytes: bytes,
|
|
Mbps: mbps(bytes, elapsed),
|
|
Error: statusError(response),
|
|
}
|
|
}
|
|
|
|
func completeSession(ctx context.Context, httpClient *http.Client, server *url.URL, path string, startedAt time.Time, finishedAt time.Time, samples []probe.Sample) (probe.Summary, error) {
|
|
payload := map[string]any{
|
|
"startedAt": startedAt,
|
|
"finishedAt": finishedAt,
|
|
"samples": samples,
|
|
}
|
|
body, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, resolve(server, path), bytes.NewReader(body))
|
|
if err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
request.Header.Set("Content-Type", "application/json")
|
|
response, err := httpClient.Do(request)
|
|
if err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
message := decodeError(response.Body)
|
|
if message == "" {
|
|
message = response.Status
|
|
}
|
|
return probe.Summary{}, fmt.Errorf("complete test session: %s", message)
|
|
}
|
|
var completed completeResponse
|
|
if err := json.NewDecoder(response.Body).Decode(&completed); err != nil {
|
|
return probe.Summary{}, err
|
|
}
|
|
return completed.Summary, nil
|
|
}
|
|
|
|
func cancelSession(ctx context.Context, httpClient *http.Client, server *url.URL, path string) error {
|
|
if strings.TrimSpace(path) == "" {
|
|
return nil
|
|
}
|
|
request, err := http.NewRequestWithContext(ctx, http.MethodPost, resolve(server, path), nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
response, err := httpClient.Do(request)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer response.Body.Close()
|
|
if response.StatusCode == http.StatusNotFound {
|
|
return nil
|
|
}
|
|
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
|
message := decodeError(response.Body)
|
|
if message == "" {
|
|
message = response.Status
|
|
}
|
|
return fmt.Errorf("cancel test session: %s", message)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func resolve(server *url.URL, path string) string {
|
|
relative, err := url.Parse(path)
|
|
if err != nil {
|
|
return server.String() + path
|
|
}
|
|
return server.ResolveReference(relative).String()
|
|
}
|
|
|
|
func decodeError(reader io.Reader) string {
|
|
var payload errorResponse
|
|
if err := json.NewDecoder(reader).Decode(&payload); err != nil {
|
|
return ""
|
|
}
|
|
return payload.Error
|
|
}
|
|
|
|
func failedSample(kind string, startedAt time.Time, bytes int64, mbps float64, err error) probe.Sample {
|
|
elapsed := time.Since(startedAt)
|
|
return probe.Sample{
|
|
At: time.Now().UTC(),
|
|
Kind: kind,
|
|
Success: false,
|
|
Latency: elapsed,
|
|
LatencyMS: round2(float64(elapsed.Microseconds()) / 1000),
|
|
Bytes: bytes,
|
|
Mbps: mbps,
|
|
Error: err.Error(),
|
|
}
|
|
}
|
|
|
|
func statusError(response *http.Response) string {
|
|
if response.StatusCode >= 200 && response.StatusCode < 300 {
|
|
return ""
|
|
}
|
|
return response.Status
|
|
}
|
|
|
|
func mbps(bytes int64, elapsed time.Duration) float64 {
|
|
if bytes <= 0 || elapsed <= 0 {
|
|
return 0
|
|
}
|
|
return round2(float64(bytes) * 8 / float64(elapsed.Microseconds()))
|
|
}
|
|
|
|
func round2(value float64) float64 {
|
|
return math.Round(value*100) / 100
|
|
}
|
|
|
|
func transferTimeout(timeout time.Duration) time.Duration {
|
|
if timeout < 30*time.Second {
|
|
return 30 * time.Second
|
|
}
|
|
return timeout
|
|
}
|
|
|
|
func cancelRequestTimeout(timeout time.Duration) time.Duration {
|
|
if timeout <= 0 {
|
|
return 3 * time.Second
|
|
}
|
|
if timeout > 5*time.Second {
|
|
return 5 * time.Second
|
|
}
|
|
return timeout
|
|
}
|
|
|
|
func cacheToken() string {
|
|
return fmt.Sprint(time.Now().UnixNano())
|
|
}
|