Merge pull request #27 from servak/probe-wait

 Wait for Probe results before ending the command

README.md

@@ -11,10 +11,13 @@ mping is a program to send ICMP echo.
 
 ## Install
 
+https://github.com/servak/mping/releases
+
+## Build
+
 ```
 go install github.com/servak/mping/cmd/mping@latest
 ```
-
 ## Permission
 
 ```
@@ -31,17 +34,27 @@ sudo chmod u+s mping
 ## Usage
 
 ```
-Usage: mping [OPTIONS] [TARGET...]
-Options:
-  -c, --config string      config path (default "~/.mping.yml")
-  -f, --fiilename string   use contents of file
-  -h, --help               Display help and exit
-  -i, --interval int       interval(ms) (default 1000)
-  -t, --timeout int        timeout(ms) (default 1000)
-  -n, --title string       print title
-  -v, --version            print version
+Usage:
+  mping [IP or HOSTNAME]... [flags]
+  mping [command]
+
 Examples:
-  mping localhost google.com 8.8.8.8 192.168.1.0/24
-  mping google.com icmpv6:google.com
-  mping -f hostslist
+mping 1.1.1.1 8.8.8.8
+mping icmpv6:google.com
+mping http://google.com
+
+Available Commands:
+  batch       Disables TUI and performs probing for a set number of iterations
+  help        Help about any command
+
+Flags:
+  -c, --config string     config path (default "~/.mping.yml")
+  -f, --filename string   use contents of file
+  -h, --help              help for mping
+  -i, --interval int      interval(ms) (default 1000)
+  -t, --timeout int       timeout(ms) (default 1000)
+  -n, --title string      print title
+  -v, --version           Display version
+
+Use "mping [command] --help" for more information about a command.
 ```

internal/command/batch.go

@@ -2,8 +2,7 @@ package command
 
 import (
 	"errors"
-	"fmt"
-	"path/filepath"
+	"sync"
 	"time"
 
 	"github.com/jedib0t/go-pretty/v6/table"
@@ -44,10 +43,6 @@ mping batch http://google.com`,
 			} else if timeout == 0 {
 				return errors.New("timeout can't be zero")
 			}
-			title, err := flags.GetString("title")
-			if err != nil {
-				return err
-			}
 			path, err := flags.GetString("config")
 			if err != nil {
 				return err
@@ -64,65 +59,50 @@ mping batch http://google.com`,
 				return nil
 			}
 
-			cfgPath, _ := filepath.Abs(path)
-			cfg, _ := config.LoadFile(cfgPath)
-			cfg.SetTitle(title)
+			cfg, _ := config.LoadFile(path)
 			_interval := time.Duration(interval) * time.Millisecond
 			_timeout := time.Duration(timeout) * time.Millisecond
 
 			res := make(chan *prober.Event)
 			probeTargets := splitProber(addDefaultProbeType(hosts), cfg)
 			manager := stats.NewMetricsManager()
-			startProbers(probeTargets, res, _interval, _timeout, manager)
 			manager.Subscribe(res)
-			ticker := time.NewTicker(_interval)
+			probers := setupProbers(probeTargets, res, manager)
+			var wg sync.WaitGroup
+			for _, p := range probers {
+				wg.Add(1)
+				go func(p prober.Prober) {
+					p.Start(res, _interval, _timeout)
+					wg.Done()
+				}(p)
+			}
 			cmd.Print("probe")
-			for range ticker.C {
-				counter--
-				if 0 < counter {
-					cmd.Print(".")
-					continue
+			for {
+				if 0 >= counter {
+					break
 				}
-				cmd.Println("")
-				cmd.Println(render(manager))
-				break
+				counter--
+				cmd.Print(".")
+				time.Sleep(_interval)
+			}
+			for _, p := range probers {
+				p.Stop()
 			}
+			wg.Wait()
+			cmd.Print("\r")
+			t := ui.TableRender(manager, stats.Success)
+			t.SetStyle(table.StyleLight)
+			cmd.Println(t.Render())
 			return nil
 		},
 	}
 
 	flags := cmd.Flags()
 	flags.StringP("filename", "f", "", "use contents of file")
-	flags.StringP("title", "n", "", "print title")
 	flags.StringP("config", "c", "~/.mping.yml", "config path")
 	flags.IntP("interval", "i", 1000, "interval(ms)")
 	flags.IntP("timeout", "t", 1000, "timeout(ms)")
 	flags.IntP("count", "", 10, "repeat count")
 
 	return cmd
 }
-
-func render(mm *stats.MetricsManager) string {
-	t := table.NewWriter()
-	t.AppendHeader(table.Row{stats.Host, stats.Sent, stats.Success, stats.Fail, stats.Loss, stats.Last, stats.Avg, stats.Best, stats.Worst, stats.LastSuccTime, stats.LastFailTime, "FAIL Reason"})
-	df := ui.DurationFormater
-	tf := ui.TimeFormater
-	for _, m := range mm.GetSortedMetricsByKey(stats.Host) {
-		t.AppendRow(table.Row{
-			m.Name,
-			m.Total,
-			m.Successful,
-			m.Failed,
-			fmt.Sprintf("%5.1f%%", m.Loss),
-			df(m.LastRTT),
-			df(m.AverageRTT),
-			df(m.MinimumRTT),
-			df(m.MaximumRTT),
-			tf(m.LastSuccTime),
-			tf(m.LastFailTime),
-			m.LastFailDetail,
-		})
-	}
-	t.SetStyle(table.StyleLight)
-	return t.Render()
-}

internal/command/mping.go

@@ -6,8 +6,10 @@ import (
 	"net"
 	"os"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/jedib0t/go-pretty/v6/table"
 	"github.com/spf13/cobra"
 
 	"github.com/servak/mping/internal/config"
@@ -71,9 +73,29 @@ mping http://google.com`,
 			res := make(chan *prober.Event)
 			probeTargets := splitProber(addDefaultProbeType(hosts), cfg)
 			manager := stats.NewMetricsManager()
-			startProbers(probeTargets, res, _interval, _timeout, manager)
+			probers := setupProbers(probeTargets, res, manager)
 			manager.Subscribe(res)
-			startCUI(manager, cfg.UI.CUI, _interval)
+			var wg sync.WaitGroup
+			for _, p := range probers {
+				wg.Add(1)
+				go func(p prober.Prober) {
+					p.Start(res, _interval, _timeout)
+					wg.Done()
+				}(p)
+			}
+			go func() {
+				startCUI(manager, cfg.UI.CUI, _interval)
+				for _, p := range probers {
+					p.Stop()
+				}
+			}()
+
+			cmd.Print("Waiting for the results of the probe. Please stand by.")
+			wg.Wait()
+			cmd.Print("\r")
+			t := ui.TableRender(manager, stats.Success)
+			t.SetStyle(table.StyleLight)
+			cmd.Println(t.Render())
 			return nil
 		},
 	}
@@ -88,15 +110,17 @@ mping http://google.com`,
 	return cmd
 }
 
-func startProbers(probeTargets map[*prober.ProberConfig][]string, res chan *prober.Event, interval, timeout time.Duration, manager *stats.MetricsManager) {
+func setupProbers(probeTargets map[*prober.ProberConfig][]string, res chan *prober.Event, manager *stats.MetricsManager) []prober.Prober {
+	var probers []prober.Prober
 	for cfg, targets := range probeTargets {
-		prober, err := newProber(cfg, manager, targets)
+		p, err := newProber(cfg, manager, targets)
 		if err != nil {
 			fmt.Println(err.Error())
 			os.Exit(1)
 		}
-		go prober.Start(res, interval, timeout)
+		probers = append(probers, p)
 	}
+	return probers
 }
 
 func startCUI(manager *stats.MetricsManager, cui *ui.CUIConfig, interval time.Duration) {

internal/prober/http.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -17,9 +18,11 @@ const (
 
 type (
 	HTTPProber struct {
-		client  *http.Client
-		targets []string
-		config  *HTTPConfig
+		client   *http.Client
+		targets  []string
+		config   *HTTPConfig
+		exitChan chan bool
+		wg       sync.WaitGroup
 	}
 
 	HTTPConfig struct {
@@ -30,9 +33,10 @@ type (
 
 func NewHTTPProber(targets []string, cfg *HTTPConfig) *HTTPProber {
 	return &HTTPProber{
-		client:  &http.Client{},
-		targets: targets,
-		config:  cfg,
+		client:   &http.Client{},
+		targets:  targets,
+		config:   cfg,
+		exitChan: make(chan bool),
 	}
 }
 
@@ -64,6 +68,8 @@ func (p *HTTPProber) failed(r chan *Event, target string, now time.Time, err err
 }
 
 func (p *HTTPProber) probe(r chan *Event, target string) {
+	p.wg.Add(1)
+	defer p.wg.Done()
 	now := time.Now()
 	p.sent(r, target)
 	resp, err := p.client.Get(target)
@@ -99,12 +105,24 @@ func (p *HTTPProber) probe(r chan *Event, target string) {
 func (p *HTTPProber) Start(r chan *Event, interval, timeout time.Duration) error {
 	p.client.Timeout = timeout
 	ticker := time.NewTicker(interval)
-	for {
-		select {
-		case <-ticker.C:
-			for _, target := range p.targets {
-				go p.probe(r, target)
+	p.wg.Add(1)
+	go func() {
+		defer p.wg.Done()
+		for {
+			select {
+			case <-p.exitChan:
+				return
+			case <-ticker.C:
+				for _, target := range p.targets {
+					go p.probe(r, target)
+				}
 			}
 		}
-	}
+	}()
+	p.wg.Wait()
+	return nil
+}
+
+func (p *HTTPProber) Stop() {
+	p.exitChan <- true
 }

internal/prober/icmp.go

@@ -20,15 +20,17 @@ const (
 
 type (
 	ICMPProber struct {
-		version ProbeType
-		c       *icmp.PacketConn
-		body    []byte
-		targets []*net.IPAddr
-		timeout time.Duration
-		runCnt  int
-		runID   int
-		tables  map[runTime]map[string]bool
-		mu      sync.Mutex
+		version  ProbeType
+		c        *icmp.PacketConn
+		body     []byte
+		targets  []*net.IPAddr
+		timeout  time.Duration
+		runCnt   int
+		runID    int
+		tables   map[runTime]map[string]bool
+		mu       sync.Mutex
+		exitChan chan bool
+		wg       sync.WaitGroup
 	}
 
 	ICMPConfig struct {
@@ -71,13 +73,14 @@ func NewICMPProber(t ProbeType, addrs []*net.IPAddr, cfg *ICMPConfig) (*ICMPProb
 		c, err = icmp.ListenPacket("ip6:ipv6-icmp", "::")
 	}
 	return &ICMPProber{
-		version: t,
-		c:       c,
-		tables:  make(map[runTime]map[string]bool),
-		targets: addrs,
-		runID:   os.Getpid() & 0xffff,
-		runCnt:  0,
-		body:    []byte(cfg.Body),
+		version:  t,
+		c:        c,
+		tables:   make(map[runTime]map[string]bool),
+		targets:  addrs,
+		runID:    os.Getpid() & 0xffff,
+		runCnt:   0,
+		body:     []byte(cfg.Body),
+		exitChan: make(chan bool),
 	}, err
 }
 
@@ -253,11 +256,30 @@ func (p *ICMPProber) Start(r chan *Event, interval, timeout time.Duration) error
 	p.timeout = timeout
 	ticker := time.NewTicker(interval)
 	go p.recvPkts(r)
+	p.wg.Add(1)
+	go func() {
+		defer p.wg.Done()
+		for {
+			select {
+			case <-p.exitChan:
+				return
+			case <-ticker.C:
+				p.probe(r)
+				go p.checkTimeout(r)
+			}
+		}
+	}()
+	p.wg.Wait()
 	for {
-		select {
-		case <-ticker.C:
-			p.probe(r)
-			go p.checkTimeout(r)
+		p.checkTimeout(r)
+		if len(p.tables) == 0 {
+			break
 		}
+		time.Sleep(interval)
 	}
+	return nil
+}
+
+func (p *ICMPProber) Stop() {
+	p.exitChan <- true
 }

internal/prober/prober.go

@@ -26,4 +26,5 @@ type Event struct {
 
 type Prober interface {
 	Start(chan *Event, time.Duration, time.Duration) error
+	Stop()
 }