lpra/internal/agent/speed.go

package agent

import (
	"context"
	"crypto/rand"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"time"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

const (
	// speedWindow is how long each direction is measured for, unless the
	// server asks for another window; speedWindowMax bounds what it may ask.
	speedWindow    = 3 * time.Second
	speedWindowMax = 30 * time.Second
	// speedWait is how long the download may take to start: in a relayed test
	// its bytes come from the browser.
	speedWait = 10 * time.Second
	// speedUpCap bounds the upload, whatever the link does in the window.
	speedUpCap = 48 << 20
)

// doSpeed measures the agent's HTTP link to the server: round-trip time of an
// empty request, then how much flows each way in the window. In a relayed
// test the server pipes both streams to the browser, so they run end to end.
func (a *Agent) doSpeed(ctx context.Context, env proto.TaskEnvelope) proto.TaskResult {
	var args proto.SysSpeedArgs
	if len(env.Args) > 0 {
		_ = json.Unmarshal(env.Args, &args)
	}
	window := speedWindow
	if args.WindowMs > 0 {
		window = min(time.Duration(args.WindowMs)*time.Millisecond, speedWindowMax)
	}
	res, err := a.measureSpeed(ctx, args.Relay, window)
	if err != nil {
		return proto.TaskResult{Error: err.Error()}
	}
	data, _ := json.Marshal(res)
	return proto.TaskResult{OK: true, Data: data}
}

func (a *Agent) measureSpeed(ctx context.Context, relay bool, window time.Duration) (proto.SysSpeedResult, error) {
	res := proto.SysSpeedResult{Relay: relay}
	streams := "/api/agent/speed"
	if relay {
		streams += "/relay"
	}
	// Ping: the first request may pay for a new connection, so keep the best.
	for i := 0; i < 4; i++ {
		start := time.Now()
		resp, err := a.speedRequest(ctx, http.MethodGet, "/api/agent/speed/down?bytes=0", nil)
		if err != nil {
			return res, err
		}
		io.Copy(io.Discard, resp.Body)
		resp.Body.Close()
		ms := float64(time.Since(start).Microseconds()) / 1000
		if i == 0 || ms < res.PingMs {
			res.PingMs = ms
		}
	}

	down, err := a.speedDown(ctx, streams+"/down", window)
	if err != nil {
		return res, err
	}
	res.DownBps = down

	// Up: send until the window closes or the cap is hit.
	body := &speedSource{deadline: time.Now().Add(window), left: speedUpCap}
	start := time.Now()
	resp, err := a.speedRequest(ctx, http.MethodPost, streams+"/up", body)
	if err != nil {
		return res, err
	}
	io.Copy(io.Discard, resp.Body)
	resp.Body.Close()
	res.UpBps = bps(int64(speedUpCap-body.left), time.Since(start))
	return res, nil
}

// speedDown reads what the server streams for the window, counted from the
// first byte, and returns the rate.
func (a *Agent) speedDown(ctx context.Context, path string, window time.Duration) (float64, error) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	t := time.AfterFunc(speedWait, cancel)
	defer t.Stop()
	resp, err := a.speedRequest(ctx, http.MethodGet, path, nil)
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()
	var (
		n     int64
		start time.Time
		buf   = make([]byte, 64<<10)
	)
	for {
		k, err := resp.Body.Read(buf)
		if k > 0 && start.IsZero() {
			start = time.Now()
			t.Reset(window)
		}
		n += int64(k)
		if err != nil {
			break
		}
	}
	if start.IsZero() {
		return 0, errors.New("нет данных для замера")
	}
	return bps(n, time.Since(start)), nil
}

func (a *Agent) speedRequest(ctx context.Context, method, path string, body io.Reader) (*http.Response, error) {
	req, err := http.NewRequestWithContext(ctx, method, a.cfg.Host+path, body)
	if err != nil {
		return nil, err
	}
	a.authHeaders(req)
	resp, err := a.http.Do(req)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent {
		resp.Body.Close()
		return nil, errStatus(resp.StatusCode)
	}
	return resp, nil
}

func bps(n int64, d time.Duration) float64 {
	if d <= 0 {
		return 0
	}
	return float64(n) * 8 / d.Seconds()
}

// speedSource is an upload body of incompressible bytes that ends at the
// deadline or after `left` bytes.
type speedSource struct {
	deadline time.Time
	left     int
	block    []byte
}

func (s *speedSource) Read(p []byte) (int, error) {
	if s.left <= 0 || !time.Now().Before(s.deadline) {
		return 0, io.EOF
	}
	if s.block == nil {
		s.block = make([]byte, 64<<10)
		rand.Read(s.block)
	}
	n := copy(p, s.block)
	if n > s.left {
		n = s.left
	}
	s.left -= n
	return n, nil
}