Skip to content

Commit

Permalink
feat: refactor runner
Browse files Browse the repository at this point in the history
  • Loading branch information
b4nst committed Sep 4, 2024
1 parent 4341eab commit 57394dc
Show file tree
Hide file tree
Showing 6 changed files with 107 additions and 144 deletions.
64 changes: 40 additions & 24 deletions cmd/icmperf/main.go
Original file line number Diff line number Diff line change
@@ -1,42 +1,58 @@
package main

import (
"context"
"net"
"time"
"fmt"

"github.com/alecthomas/kong"
tea "github.com/charmbracelet/bubbletea"
"golang.org/x/net/icmp"
"github.com/prometheus-community/pro-bing"

"github.com/b4nst/icmperf/pkg/cli"
"github.com/b4nst/icmperf/pkg/model"
"github.com/b4nst/icmperf/pkg/pinger"
"github.com/b4nst/icmperf/pkg/recorder"
"github.com/b4nst/icmperf/pkg/session"
)

func main() {
cli := cli.CLI{}
ktx := kong.Parse(&cli)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()

pinger := pinger.NewPinger(cli.BindAddr, cli.MTU, cli.Timeout)
record := recorder.NewRecord()
pinger.OnRecv(func(m *icmp.Message, t time.Time) error {
body := m.Body.(*icmp.Echo)
id := uint64(body.ID)<<32 | uint64(body.Seq)
record.PacketReceived(id, len(body.Data), t)
return nil
})

peer := net.UDPAddr(cli.Target)
m := model.NewModel(pinger, record, &peer, cli.MTU, cli.Duration)
if err := pinger.Start(ctx); err != nil {
ktx.FatalIfErrorf(err)

pingers := make([]*probing.Pinger, 0, len(cli.DataSizes)+1)

// Latenncy pinger
p, err := probing.NewPinger(cli.Target)
ktx.FatalIfErrorf(err)
if cli.Duration > 0 {
p.Timeout = cli.Duration
} else {
p.Count = cli.Count
}
p.SetPrivileged(cli.Privileged)
p.Size = 24 // Minimum size
pingers = append(pingers, p)

if _, err := tea.NewProgram(m).Run(); err != nil {
// Data pingers
for _, ds := range cli.DataSizes {
p, err := probing.NewPinger(cli.Target)
ktx.FatalIfErrorf(err)
if cli.Duration > 0 {
p.Timeout = cli.Duration
} else {
p.Count = cli.Count
}
p.Size = ds
p.SetPrivileged(cli.Privileged)
pingers = append(pingers, p)
}

s := session.NewSession(pingers)
ktx.FatalIfErrorf(s.Run())

stats := s.Statistics()
stat, err := recorder.ProcessStats(stats)
ktx.FatalIfErrorf(err)

for _, s := range stats {
fmt.Printf("%s\n", s)
}
fmt.Println("- - - - - - - - - - - - - - - - - - - - - - - - -")
fmt.Printf("%s\n", stat)
}
10 changes: 8 additions & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,11 @@ require (
github.com/charmbracelet/bubbles v0.19.0
github.com/charmbracelet/bubbletea v1.1.0
github.com/dustin/go-humanize v1.0.1
github.com/emirpasic/gods v1.18.1
github.com/montanaflynn/stats v0.7.1
github.com/prometheus-community/pro-bing v0.4.1
github.com/stretchr/testify v1.9.0
golang.org/x/net v0.28.0
golang.org/x/sync v0.8.0
)

require (
Expand All @@ -17,16 +20,19 @@ require (
github.com/charmbracelet/lipgloss v0.13.0 // indirect
github.com/charmbracelet/x/ansi v0.2.3 // indirect
github.com/charmbracelet/x/term v0.2.0 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
github.com/mattn/go-isatty v0.0.20 // indirect
github.com/mattn/go-localereader v0.0.1 // indirect
github.com/mattn/go-runewidth v0.0.16 // indirect
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect
github.com/muesli/cancelreader v0.2.2 // indirect
github.com/muesli/termenv v0.15.2 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/rivo/uniseg v0.4.7 // indirect
golang.org/x/sync v0.8.0 // indirect
golang.org/x/sys v0.24.0 // indirect
golang.org/x/text v0.17.0 // indirect
gopkg.in/yaml.v3 v3.0.1 // indirect
)
18 changes: 16 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,14 @@ github.com/charmbracelet/x/ansi v0.2.3 h1:VfFN0NUpcjBRd4DnKfRaIRo53KRgey/nhOoEqo
github.com/charmbracelet/x/ansi v0.2.3/go.mod h1:dk73KoMTT5AX5BsX0KrqhsTqAnhZZoCBjs7dGWp4Ktw=
github.com/charmbracelet/x/term v0.2.0 h1:cNB9Ot9q8I711MyZ7myUR5HFWL/lc3OpU8jZ4hwm0x0=
github.com/charmbracelet/x/term v0.2.0/go.mod h1:GVxgxAbjUrmpvIINHIQnJJKpMlHiZ4cktEQCN6GWyF0=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
github.com/emirpasic/gods v1.18.1 h1:FXtiHYKDGKCW2KzwZKx0iC0PQmdlorYgdFG9jPXJ1Bc=
github.com/emirpasic/gods v1.18.1/go.mod h1:8tpGGwCnJ5H4r6BWwaV6OrWmMoPhUl5jm/FMNAnJvWQ=
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f/go.mod h1:vw97MGsxSvLiUE2X8qFplwetxpGLQrlU1Q9AUEIzCaM=
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/hexops/gotextdiff v1.0.3 h1:gitA9+qJrrTCsiCl7+kh75nPqQt1cx4ZkudSTLoUqJM=
github.com/hexops/gotextdiff v1.0.3/go.mod h1:pSWU5MAI3yDq+fZBTazCSJysOMbxWL1BSow5/V2vxeg=
github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
Expand All @@ -34,15 +36,23 @@ github.com/mattn/go-localereader v0.0.1 h1:ygSAOl7ZXTx4RdPYinUpg6W99U8jWvWi9Ye2J
github.com/mattn/go-localereader v0.0.1/go.mod h1:8fBrzywKY7BI3czFoHkuzRoWE9C+EiG4R1k4Cjx5p88=
github.com/mattn/go-runewidth v0.0.16 h1:E5ScNMtiwvlvB5paMFdw9p4kSQzbXFikJ5SQO6TULQc=
github.com/mattn/go-runewidth v0.0.16/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
github.com/montanaflynn/stats v0.7.1 h1:etflOAAHORrCC44V+aR6Ftzort912ZU+YLiSTuV8eaE=
github.com/montanaflynn/stats v0.7.1/go.mod h1:etXPPgVO6n31NxCd9KQUMvCM+ve0ruNzt6R8Bnaayow=
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 h1:ZK8zHtRHOkbHy6Mmr5D264iyp3TiX5OmNcI5cIARiQI=
github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6/go.mod h1:CJlz5H+gyd6CUWT45Oy4q24RdLyn7Md9Vj2/ldJBSIo=
github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELUXHmA=
github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo=
github.com/muesli/termenv v0.15.2 h1:GohcuySI0QmI3wN8Ok9PtKGkgkFIk7y6Vpb5PvrY+Wo=
github.com/muesli/termenv v0.15.2/go.mod h1:Epx+iuz8sNs7mNKhxzH4fWXGNpZwUaJKRS1noLXviQ8=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
github.com/prometheus-community/pro-bing v0.4.1 h1:aMaJwyifHZO0y+h8+icUz0xbToHbia0wdmzdVZ+Kl3w=
github.com/prometheus-community/pro-bing v0.4.1/go.mod h1:aLsw+zqCaDoa2RLVVSX3+UiCkBBXTMtZC3c7EkfWnAE=
github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
github.com/rivo/uniseg v0.4.7 h1:WUdvkW8uEhrYfLC4ZzdpI2ztxP1I582+49Oc5Mq64VQ=
github.com/rivo/uniseg v0.4.7/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88=
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE=
golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg=
golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ=
Expand All @@ -53,3 +63,7 @@ golang.org/x/sys v0.24.0 h1:Twjiwq9dn6R1fQcyiK+wQyHWfaz/BJB+YIpzU/Cv3Xg=
golang.org/x/sys v0.24.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc=
golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
29 changes: 17 additions & 12 deletions pkg/cli/cli.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,25 +2,30 @@ package cli

import (
"fmt"
"net"
"time"
)

type Target net.UDPAddr
type DataSizes []int

func (t *Target) UnmarshalText(text []byte) error {
ip, err := net.LookupIP(string(text))
if err != nil {
return fmt.Errorf("failed to lookup target: %w", err)
func (ds *DataSizes) Validate() error {
if len(*ds) == 0 {
return fmt.Errorf("data sizes must have at least one element")
}

for _, size := range *ds {
if size < 24 {
return fmt.Errorf("data size must be greater than 0")
}
}
t.IP = ip[0]
return nil
}

type CLI struct {
Target Target `arg:"" help:"The target host to ping."`
MTU int `help:"The maximum transmission unit of your interface." short:"m" default:"1500"`
Timeout time.Duration `help:"The timeout for each ping." short:"t" default:"5s"`
Duration time.Duration `help:"The duration of the test." short:"d" default:"30s"`
BindAddr string `help:"The address to bind the ICMP listener to." default:"0.0.0.0" short:"l"`
Timeout time.Duration `help:"The timeout for each ping." short:"t" default:"5s"`
Duration time.Duration `help:"The duration of the test." short:"d" xor:"duration"`
Count int `help:"The number of pings to send." short:"c" default:"10" xor:"duration"`
Privileged bool `help:"Use privileged ICMP sockets." default:"false" short:"p"`
DataSizes []int `help:"The size of the data to send." short:"s" default:"1024"`

Target string `arg:"" help:"The target host to ping."`
}
75 changes: 4 additions & 71 deletions pkg/model/model.go
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
package model

import (
"fmt"
"net"
"strconv"
"strings"
"time"

Expand All @@ -20,11 +20,8 @@ var (
)

type Model struct {
pinger *pinger.Pinger
record *recorder.Record
peer *net.UDPAddr
payload []byte
stats *recorder.Stats

duration time.Duration
timer timer.Model
Expand All @@ -33,21 +30,16 @@ type Model struct {

func NewModel(pinger *pinger.Pinger, record *recorder.Record, peer *net.UDPAddr, mtu int, duration time.Duration) *Model {
return &Model{
pinger: pinger,
record: record,
peer: peer,
payload: make([]byte, mtu-28),
stats: nil,

duration: duration,
timer: timer.NewWithInterval(duration, 200*time.Millisecond),
progress: progress.New(progress.WithFillCharacters('▱', ' '), progress.WithDefaultGradient()),
}
}

type latencyTick time.Time
type pingTick time.Time
type statsMsg *recorder.Stats

func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
cmds := []tea.Cmd{}
Expand All @@ -62,23 +54,11 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.timer, tc = m.timer.Update(msg)
pc = m.progress.SetPercent(1.0 - m.timer.Timeout.Seconds()/m.duration.Seconds())
cmds = append(cmds, tc, pc)
case timer.TimeoutMsg:
cmds = append(cmds, m.progress.SetPercent(1.0))
cmds = append(cmds, m.statsCmd())
case tea.WindowSizeMsg:
m.progress.Width = min(msg.Width-4, maxWidth)
case latencyTick:
if !m.timer.Timedout() {
cmds = append(cmds, m.pingLatency())
}
case pingTick:
if !m.timer.Timedout() {
cmds = append(cmds, m.pingBandwidth())
}
case statsMsg:
m.stats = msg
case error:
fmt.Println(msg)
cmds = append(cmds, tea.Quit)

case progress.FrameMsg:
progressModel, cmd := m.progress.Update(msg)
m.progress = progressModel.(progress.Model)
Expand All @@ -91,10 +71,6 @@ func (m *Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
func (m *Model) View() string {
view := strings.Builder{}

view.WriteString("Probing peer at ")
view.WriteString(m.peer.IP.String())
view.WriteString("\n\n")

view.WriteString("Probe size: ")
view.WriteString(humanize.Bytes(uint64(len(m.payload))))
view.WriteString(" ")
Expand All @@ -105,52 +81,9 @@ func (m *Model) View() string {
view.WriteString(m.progress.View())
view.WriteString("\n\n")

if m.stats != nil {
view.WriteString("Bandwidth: ")
view.WriteString(humanize.Bytes(uint64(m.stats.Bandwidth())))
view.WriteString("/s ")

view.WriteString("Latency: ")
view.WriteString(m.stats.Latency().Round(time.Millisecond).String())
view.WriteString(" ")

view.WriteString("Loss rate: ")
view.WriteString(strconv.FormatFloat(m.stats.PacketLoss()*100, 'f', 2, 64))
view.WriteString("%\n")
}

return view.String()
}

func (m *Model) Init() tea.Cmd {
return tea.Batch(m.pingLatency(), m.pingBandwidth(), m.timer.Init())
}

func (m *Model) pingLatency() tea.Cmd {
return tea.Tick(1*time.Second, func(t time.Time) tea.Msg {
sendAndRecord(m.pinger, m.peer, []byte{}, m.record)
return latencyTick(t)
})
}

func (m *Model) pingBandwidth() tea.Cmd {
return tea.Tick(500*time.Millisecond, func(t time.Time) tea.Msg {
sendAndRecord(m.pinger, m.peer, m.payload, m.record)
return pingTick(t)
})
}

func (m *Model) statsCmd() tea.Cmd {
return func() tea.Msg {
return statsMsg(m.record.Stats())
}
}

func sendAndRecord(p *pinger.Pinger, peer *net.UDPAddr, payload []byte, r *recorder.Record) error {
id, t, err := p.Send(peer, payload)
if err != nil {
return err
}
r.PacketSent(id, len(payload), t)
return nil
return tea.Batch(m.timer.Init())
}
Loading

0 comments on commit 57394dc

Please sign in to comment.