From c8b96318d8264a98b0811322f4ee9acff3ce8ae8 Mon Sep 17 00:00:00 2001
From: r3inbowari <r3inbowari@gmail.com>
Date: Thu, 9 May 2024 15:14:02 +0800
Subject: [PATCH] feat(api): add RunMulti and RunMultiWithContext (#204)

* chore: typo fix

* feat(api): add RunMulti and RunMultiWithContext

* fix: lint error
---
 example/packet_loss/main.go | 13 ++++++++++--
 speedtest.go                |  5 ++---
 speedtest/loss.go           | 41 +++++++++++++++++++++++++++++++++++--
 speedtest/transport/tcp.go  |  3 +++
 4 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/example/packet_loss/main.go b/example/packet_loss/main.go
index 0f66183..c6fae4a 100644
--- a/example/packet_loss/main.go
+++ b/example/packet_loss/main.go
@@ -6,6 +6,7 @@ import (
 	"github.com/showwin/speedtest-go/speedtest/transport"
 	"log"
 	"sync"
+	"time"
 )
 
 // Note: The current packet loss analyzer does not support udp over http.
@@ -19,12 +20,15 @@ func main() {
 	targets := serverList.Available()
 
 	// 2. Create a packet loss analyzer, use default options
-	analyzer, err := speedtest.NewPacketLossAnalyzer(nil)
-	checkError(err)
+	analyzer := speedtest.NewPacketLossAnalyzer(&speedtest.PacketLossAnalyzerOptions{
+		PacketSendingInterval: time.Millisecond * 100,
+	})
 
 	wg := &sync.WaitGroup{}
 	// 3. Perform packet loss analysis on all available servers
+	var hosts []string
 	for _, server := range *targets {
+		hosts = append(hosts, server.Host)
 		wg.Add(1)
 		//ctx, cancel := context.WithTimeout(context.Background(), time.Second*20)
 		//go func(server *speedtest.Server, analyzer *speedtest.PacketLossAnalyzer, ctx context.Context, cancel context.CancelFunc) {
@@ -46,6 +50,11 @@ func main() {
 		}(server, analyzer)
 	}
 	wg.Wait()
+
+	// use mixed PacketLoss
+	mixed, err := analyzer.RunMulti(hosts)
+	checkError(err)
+	fmt.Printf("Mixed packets lossed: %.2f\n", mixed)
 }
 
 func checkError(err error) {
diff --git a/speedtest.go b/speedtest.go
index a2f922e..78e7790 100644
--- a/speedtest.go
+++ b/speedtest.go
@@ -136,8 +136,7 @@ func main() {
 		})
 
 		// 3.0 create a packet loss analyzer, use default options
-		var analyzer *speedtest.PacketLossAnalyzer
-		analyzer, err = speedtest.NewPacketLossAnalyzer(&speedtest.PacketLossAnalyzerOptions{
+		var analyzer = speedtest.NewPacketLossAnalyzer(&speedtest.PacketLossAnalyzerOptions{
 			SourceInterface: *source,
 		})
 
@@ -151,7 +150,7 @@ func main() {
 					packetLossAnalyzerCancel() // cancel early
 				}
 			}()
-			task.Println("Packet Loss Analyzer: Running in background (<= 30 Sec)")
+			task.Println("Packet Loss Analyzer: Running in background (<= 30 Secs)")
 			task.Complete()
 		})
 
diff --git a/speedtest/loss.go b/speedtest/loss.go
index 781fa60..a7cc3d1 100644
--- a/speedtest/loss.go
+++ b/speedtest/loss.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"github.com/showwin/speedtest-go/speedtest/transport"
 	"net"
+	"sync"
 	"time"
 )
 
@@ -22,7 +23,7 @@ type PacketLossAnalyzer struct {
 	options *PacketLossAnalyzerOptions
 }
 
-func NewPacketLossAnalyzer(options *PacketLossAnalyzerOptions) (*PacketLossAnalyzer, error) {
+func NewPacketLossAnalyzer(options *PacketLossAnalyzerOptions) *PacketLossAnalyzer {
 	if options == nil {
 		options = &PacketLossAnalyzerOptions{}
 	}
@@ -56,7 +57,43 @@ func NewPacketLossAnalyzer(options *PacketLossAnalyzerOptions) (*PacketLossAnaly
 	}
 	return &PacketLossAnalyzer{
 		options: options,
-	}, nil
+	}
+}
+
+// RunMulti Mix all servers to get the average packet loss.
+func (pla *PacketLossAnalyzer) RunMulti(hosts []string) (float64, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), pla.options.SamplingDuration)
+	defer cancel()
+	return pla.RunMultiWithContext(ctx, hosts)
+}
+
+func (pla *PacketLossAnalyzer) RunMultiWithContext(ctx context.Context, hosts []string) (float64, error) {
+	results := make(map[string]float64)
+	mutex := &sync.Mutex{}
+	wg := &sync.WaitGroup{}
+	for _, host := range hosts {
+		wg.Add(1)
+		go func(h string) {
+			defer wg.Done()
+			_ = pla.RunWithContext(ctx, h, func(packetLoss *transport.PLoss) {
+				loss := packetLoss.Loss()
+				if loss != -1 {
+					mutex.Lock()
+					results[h] = loss
+					mutex.Unlock()
+				}
+			})
+		}(host)
+	}
+	wg.Wait()
+	if len(results) == 0 {
+		return -1, transport.ErrUnsupported
+	}
+	packetLossAvg := 0.0
+	for _, hostPacketLoss := range results {
+		packetLossAvg += hostPacketLoss
+	}
+	return packetLossAvg / float64(len(results)), nil
 }
 
 func (pla *PacketLossAnalyzer) Run(host string, callback func(packetLoss *transport.PLoss)) error {
diff --git a/speedtest/transport/tcp.go b/speedtest/transport/tcp.go
index 2dcbd27..87b2236 100644
--- a/speedtest/transport/tcp.go
+++ b/speedtest/transport/tcp.go
@@ -191,6 +191,9 @@ func (p PLoss) String() string {
 }
 
 func (p PLoss) Loss() float64 {
+	if p.Sent == 0 {
+		return -1
+	}
 	return 1 - (float64(p.Sent-p.Dup))/float64(p.Max+1)
 }