111 lines
2.3 KiB
Go
111 lines
2.3 KiB
Go
package loadsim
|
|
|
|
import (
|
|
"context"
|
|
"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
|
|
}
|
|
|
|
func (r Runner) Run(ctx context.Context, cfg Config) (Stats, error) {
|
|
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,
|
|
}).Build(); err != nil {
|
|
return Stats{}, err
|
|
}
|
|
payload, err := FrameTemplate(cfg.Protocol, cfg.Template)
|
|
if err != nil {
|
|
return Stats{}, err
|
|
}
|
|
if !cfg.SendFrames {
|
|
payload = nil
|
|
}
|
|
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
|
|
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)
|
|
go func() {
|
|
defer wg.Done()
|
|
defer conn.Close()
|
|
if cfg.SendFrames {
|
|
writeLoop(runCtx, conn, payload, cfg.SendInterval, &stats)
|
|
return
|
|
}
|
|
<-runCtx.Done()
|
|
}()
|
|
}
|
|
}
|
|
<-runCtx.Done()
|
|
wg.Wait()
|
|
return stats, nil
|
|
}
|
|
|
|
func writeLoop(ctx context.Context, conn net.Conn, payload []byte, interval time.Duration, stats *Stats) {
|
|
for {
|
|
_ = 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
|
|
}
|
|
}
|