lingniu-vehicle-ingest/go/vehicle-gateway/internal/loadsim/runner.go

package loadsim

import (
	"context"
	"errors"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"
)

type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

type Runner struct {
	Dial DialFunc
}

type Stats struct {
	ConnectionsOpened int64
	ConnectionsFailed int64
	FramesWritten     int64
	WriteErrors       int64
	ResponseBytes     int64
	ReadErrors        int64
}

func (r Runner) Run(ctx context.Context, cfg Config) (Stats, error) {
	phoneBase := cfg.JT808PhoneBase
	if phoneBase <= 0 {
		phoneBase = DefaultJT808PhoneBase
	}
	if _, err := (&FlagConfig{
		protocol:       string(cfg.Protocol),
		addr:           cfg.Addr,
		connections:    cfg.Connections,
		connectRate:    cfg.ConnectRatePerSecond,
		sendInterval:   cfg.SendInterval,
		duration:       cfg.Duration,
		template:       cfg.Template,
		jt808PhoneBase: phoneBase,
		drainResponses: cfg.DrainResponses,
	}).Build(); err != nil {
		return Stats{}, err
	}
	var factory *FrameFactory
	if cfg.SendFrames {
		var err error
		factory, err = NewFrameFactoryWithJT808PhoneBase(cfg.Protocol, cfg.Template, phoneBase)
		if err != nil {
			return Stats{}, err
		}
	}
	dial := r.Dial
	if dial == nil {
		dialer := &net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}
		dial = dialer.DialContext
	}

	runCtx, cancel := context.WithTimeout(ctx, cfg.Duration)
	defer cancel()

	var stats Stats
	var wg sync.WaitGroup
	connectionIndex := 0
	for i, batch := range ConnectionBatches(cfg.Connections, cfg.ConnectRatePerSecond) {
		if i > 0 {
			if !sleepOrDone(runCtx, time.Second) {
				break
			}
		}
		for range batch {
			if runCtx.Err() != nil {
				break
			}
			conn, err := dial(runCtx, "tcp", cfg.Addr)
			if err != nil {
				atomic.AddInt64(&stats.ConnectionsFailed, 1)
				continue
			}
			atomic.AddInt64(&stats.ConnectionsOpened, 1)
			wg.Add(1)
			connIndex := connectionIndex
			connectionIndex++
			go func() {
				defer wg.Done()
				if cfg.SendFrames {
					var readDone chan struct{}
					if cfg.DrainResponses {
						readDone = make(chan struct{})
						go func() {
							drainResponses(runCtx, conn, &stats)
							close(readDone)
						}()
					}
					writeLoop(runCtx, conn, factory, connIndex, cfg.SendInterval, &stats)
					_ = conn.Close()
					if readDone != nil {
						<-readDone
					}
					return
				}
				<-runCtx.Done()
				_ = conn.Close()
			}()
		}
	}
	<-runCtx.Done()
	wg.Wait()
	return stats, nil
}

func drainResponses(ctx context.Context, conn net.Conn, stats *Stats) {
	buffer := make([]byte, 4096)
	for {
		_ = conn.SetReadDeadline(time.Now().Add(time.Second))
		read, err := conn.Read(buffer)
		if read > 0 {
			atomic.AddInt64(&stats.ResponseBytes, int64(read))
		}
		if err == nil {
			continue
		}
		if ctx.Err() != nil || errors.Is(err, net.ErrClosed) || errors.Is(err, io.EOF) {
			return
		}
		if timeout, ok := err.(net.Error); ok && timeout.Timeout() {
			continue
		}
		atomic.AddInt64(&stats.ReadErrors, 1)
		return
	}
}

func writeLoop(ctx context.Context, conn net.Conn, factory *FrameFactory, connectionIndex int, interval time.Duration, stats *Stats) {
	for frameIndex := int64(0); ; frameIndex++ {
		payload, err := factory.Frame(connectionIndex, frameIndex)
		if err != nil {
			atomic.AddInt64(&stats.WriteErrors, 1)
			return
		}
		_ = conn.SetWriteDeadline(time.Now().Add(3 * time.Second))
		if _, err := conn.Write(payload); err != nil {
			atomic.AddInt64(&stats.WriteErrors, 1)
			return
		}
		atomic.AddInt64(&stats.FramesWritten, 1)
		if !sleepOrDone(ctx, interval) {
			return
		}
	}
}

func sleepOrDone(ctx context.Context, duration time.Duration) bool {
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-timer.C:
		return true
	}
}