feat(dstp): implement concurrency for running tests concurrently

Signed-off-by: Yagiz Degirmenci <yagizcanilbey1903@gmail.com>
ycd · Nov 22, 2021 · 4a52d9f · 4a52d9f
1 parent 95a1d7a
commit 4a52d9f
Show file tree

Hide file tree

Showing 8 changed files with 109 additions and 100 deletions.
diff --git a/README.md b/README.md
@@ -17,7 +17,6 @@ Usage: dstp [OPTIONS] [ARGS]
 Options:
 	-a, --addr   <string>  The URL or the IP address to run tests against      [REQUIRED]
 	-o, --out    <string>  The type of the output, either json or plaintext    [Default: plaintext] 
-	-c           <bool>    Run all the tests concurrently.                     [Default: false]
 	-p           <int>     Number of ping packets                              [Default: 3]
 	-t           <int>     Give up on ping after this many seconds             [Default: 2s per ping packet]
 	-h, --help             Show this message and exit.
@@ -128,7 +127,6 @@ for 64-bit Windows, macOS, and Linux targets. They contain the compiled executab
    Options:
    -a, --addr   <string>  The URL or the IP address to run tests against      [REQUIRED]
    -o, --out    <string>  The type of the output, either json or plaintext    [Default: plaintext]
-   -c           <bool>    Run all the tests concurrently.                     [Default: false]
    -p           <int>     Number of ping packets                              [Default: 3]
    -t           <int>     Give up on ping after this many seconds             [Default: 2s per ping packet]
    -h, --help             Show this message and exit.

diff --git a/config/config.go b/config/config.go
@@ -9,20 +9,18 @@ import (
 )
 
 type Config struct {
-	Addr       string
-	Output     string
-	PingCount  int
-	Timeout    int
-	ShowHelp   bool
-	Concurrent bool
+	Addr      string
+	Output    string
+	PingCount int
+	Timeout   int
+	ShowHelp  bool
 }
 
 var usageStr = `
 Usage: dstp [OPTIONS] [ARGS]
 Options:
 	-a, --addr   <string>  The URL or the IP address to run tests against      [REQUIRED]
 	-o, --out    <string>  The type of the output, either json or plaintext    [Default: plaintext] 
-	-c           <bool>    Run all the tests concurrently.                     [Default: false]
 	-p           <int>     Number of ping packets                              [Default: 3]
 	-t           <int>     Give up on ping after this many seconds             [Default: 2s per ping packet]
 	-h, --help             Show this message and exit.
@@ -52,7 +50,6 @@ func ConfigureOptions(fs *flag.FlagSet, args []string) (*Config, error) {
 	fs.StringVar(&opts.Output, "out", "plaintext", "The type of the output")
 	fs.IntVar(&opts.PingCount, "p", 3, "Number of ping packets")
 	fs.IntVar(&opts.Timeout, "t", -1, "Give up on ping after this many seconds")
-	fs.BoolVar(&opts.Concurrent, "c", false, "Run all the tests concurrently")
 	fs.BoolVar(&opts.ShowHelp, "h", false, "Show help message")
 	fs.BoolVar(&opts.ShowHelp, "help", false, "Show help message")
 

diff --git a/pkg/common/types.go b/pkg/common/types.go
@@ -1,5 +1,12 @@
 package common
 
+import (
+	"encoding/json"
+	"fmt"
+	"reflect"
+	"sync"
+)
+
 type Output string
 
 type Args []string
@@ -13,3 +20,35 @@ func (o Output) String() string {
 func (a Address) String() string {
 	return string(a)
 }
+
+type Result struct {
+	Ping      string `json:"ping"`
+	DNS       string `json:"dns"`
+	SystemDNS string `json:"system_dns"`
+	TLS       string `json:"tls"`
+	HTTPS     string `json:"https"`
+
+	Mu sync.Mutex `json:"-"`
+}
+
+func (r Result) Output(outputType string) string {
+	var output string
+
+	switch outputType {
+	case "plaintext":
+		v := reflect.ValueOf(r)
+		for i := 0; i < v.NumField(); i++ {
+			if v.Type().Field(i).Name == "Mu" {
+				continue
+			}
+
+			output += fmt.Sprintf("%s: %v\n", White(v.Type().Field(i).Name), Green(v.Field(i).Interface()))
+		}
+	case "json":
+		// SAFETY: we are sure that this never fails
+		byt, _ := json.MarshalIndent(r, "", "  ")
+		output += string(byt)
+	}
+
+	return output
+}
diff --git a/pkg/dstp/dstp.go b/pkg/dstp/dstp.go
@@ -3,11 +3,10 @@ package dstp
 import (
 	"context"
 	"crypto/tls"
-	"encoding/json"
 	"fmt"
 	"math"
 	"net/http"
-	"reflect"
+	"sync"
 	"time"
 
 	"github.com/ycd/dstp/config"
@@ -16,70 +15,28 @@ import (
 	"github.com/ycd/dstp/pkg/ping"
 )
 
-type Result struct {
-	Ping      string `json:"ping"`
-	DNS       string `json:"dns"`
-	SystemDNS string `json:"system_dns"`
-	TLS       string `json:"tls"`
-	HTTPS     string `json:"https"`
-}
-
-func (r Result) Output(outputType string) string {
-	var output string
-
-	switch outputType {
-	case "plaintext":
-		v := reflect.ValueOf(r)
-		for i := 0; i < v.NumField(); i++ {
-			output += fmt.Sprintf("%s: %v\n", common.White(v.Type().Field(i).Name), common.Green(v.Field(i).Interface()))
-		}
-	case "json":
-		// SAFETY: we are sure that this never fails
-		byt, _ := json.MarshalIndent(r, "", "  ")
-		output += string(byt)
-	}
-
-	return output
-}
-
 // RunAllTests executes all the tests against the given domain, IP or DNS server.
 func RunAllTests(ctx context.Context, config config.Config) error {
-	var result Result
+	var result common.Result
 
 	addr, err := getAddr(config.Addr)
 	if err != nil {
 		return err
 	}
 
-	if out, err := ping.RunTest(ctx, common.Address(addr), config.PingCount, config.Timeout); err != nil {
-		result.Ping = err.Error()
-	} else {
-		result.Ping = out.String()
-	}
+	var wg sync.WaitGroup
+	wg.Add(5)
 
-	if out, err := ping.RunDNSTest(ctx, common.Address(addr), config.PingCount, config.Timeout); err != nil {
-		result.DNS = err.Error()
-	} else {
-		result.DNS = out.String()
-	}
+	go ping.RunTest(ctx, &wg, common.Address(addr), config.PingCount, config.Timeout, &result)
 
-	if out, err := lookup.Host(ctx, common.Address(addr)); err != nil {
-		result.SystemDNS = err.Error()
-	} else {
-		result.SystemDNS = out.String()
-	}
+	go ping.RunDNSTest(ctx, &wg, common.Address(addr), config.PingCount, config.Timeout, &result)
 
-	if out, err := testTLS(ctx, common.Address(addr)); err != nil {
-		result.TLS = err.Error()
-	} else {
-		result.TLS = out.String()
-	}
+	go lookup.Host(ctx, &wg, common.Address(addr), &result)
 
-	if out, err := testHTTPS(ctx, common.Address(addr), config.Timeout); err != nil {
-		result.HTTPS = err.Error()
-	} else {
-		result.HTTPS = out.String()
-	}
+	go testTLS(ctx, &wg, common.Address(addr), &result)
+
+	go testHTTPS(ctx, &wg, common.Address(addr), config.Timeout, &result)
+	wg.Wait()
 
 	s := result.Output(config.Output)
 	s += "\n"
@@ -89,16 +46,17 @@ func RunAllTests(ctx context.Context, config config.Config) error {
 	return nil
 }
 
-func testTLS(ctx context.Context, address common.Address) (common.Output, error) {
+func testTLS(ctx context.Context, wg *sync.WaitGroup, address common.Address, result *common.Result) error {
 	var output string
+	defer wg.Done()
 
 	conn, err := tls.Dial("tcp", fmt.Sprintf("%s:443", string(address)), nil)
 	if err != nil {
-		return "", err
+		return err
 	}
 	err = conn.VerifyHostname(string(address))
 	if err != nil {
-		return "", err
+		return err
 	}
 
 	notAfter := conn.ConnectionState().PeerCertificates[0].NotAfter
@@ -110,15 +68,19 @@ func testTLS(ctx context.Context, address common.Address) (common.Output, error)
 		output += fmt.Sprintf("the certificate expired %v days ago", -expiry)
 	}
 
-	return common.Output(output), nil
+	result.Mu.Lock()
+	result.TLS = output
+	result.Mu.Unlock()
+
+	return nil
 }
 
-func testHTTPS(ctx context.Context, address common.Address, t int) (common.Output, error) {
-	var output string
+func testHTTPS(ctx context.Context, wg *sync.WaitGroup, address common.Address, t int, result *common.Result) error {
+	defer wg.Done()
 
 	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf("https://%s", address), nil)
 	if err != nil {
-		return "", err
+		return err
 	}
 
 	client := http.Client{
@@ -127,9 +89,12 @@ func testHTTPS(ctx context.Context, address common.Address, t int) (common.Outpu
 
 	resp, err := client.Do(req)
 	if err != nil {
-		return "", err
+		return err
 	}
 
-	output += fmt.Sprintf("got %s", resp.Status)
-	return common.Output(output), nil
+	result.Mu.Lock()
+	result.HTTPS = fmt.Sprintf("got %s", resp.Status)
+	result.Mu.Unlock()
+
+	return nil
 }
diff --git a/pkg/dstp/dstp_test.go b/pkg/dstp/dstp_test.go
@@ -15,7 +15,6 @@ func TestRunAllTests(t *testing.T) {
 		Addr:       "https://jvns.ca",
 		Output:     "plaintext",
 		ShowHelp:   false,
-		Concurrent: false,
 		Timeout:    3,
 		PingCount:  3,
 	}
@@ -24,7 +23,6 @@ func TestRunAllTests(t *testing.T) {
 		Addr:       "8.8.8.8",
 		Output:     "plaintext",
 		ShowHelp:   false,
-		Concurrent: false,
 		Timeout:    3,
 		PingCount:  3,
 	}
@@ -33,7 +31,6 @@ func TestRunAllTests(t *testing.T) {
 		Addr:       "facebook.com",
 		Output:     "plaintext",
 		ShowHelp:   false,
-		Concurrent: false,
 		Timeout:    3,
 		PingCount:  3,
 	}
@@ -42,7 +39,6 @@ func TestRunAllTests(t *testing.T) {
 		Addr:       "https://meta.stackoverflow.com/",
 		Output:     "plaintext",
 		ShowHelp:   false,
-		Concurrent: false,
 		Timeout:    3,
 		PingCount:  3,
 	}
@@ -51,7 +47,6 @@ func TestRunAllTests(t *testing.T) {
 		Addr:       "facebook.com:80",
 		Output:     "plaintext",
 		ShowHelp:   false,
-		Concurrent: false,
 		Timeout:    3,
 		PingCount:  3,
 	}

diff --git a/pkg/lookup/host.go b/pkg/lookup/host.go
@@ -5,15 +5,18 @@ import (
 	"github.com/ycd/dstp/pkg/common"
 	"net"
 	"strings"
+	"sync"
 )
 
-func Host(ctx context.Context, addr common.Address) (common.Output, error) {
+func Host(ctx context.Context, wg *sync.WaitGroup, addr common.Address, result *common.Result) error {
+	defer wg.Done()
+
 	addrs, err := net.LookupHost(addr.String())
 	if err != nil {
-		return "", err
+		return err
 	}
 
-	output := "resolving " + strings.Join(addrs, ", ")
+	result.SystemDNS = "resolving " + strings.Join(addrs, ", ")
 
-	return common.Output(output), nil
+	return nil
 }
diff --git a/pkg/ping/dns.go b/pkg/ping/dns.go
@@ -4,15 +4,16 @@ import (
 	"context"
 	"fmt"
 	"github.com/ycd/dstp/pkg/common"
+	"sync"
 	"time"
 )
 
-func RunDNSTest(ctx context.Context, addr common.Address, count int, timeout int) (common.Output, error) {
-	var output string
+func RunDNSTest(ctx context.Context, wg *sync.WaitGroup, addr common.Address, count int, timeout int, result *common.Result) error {
+	defer wg.Done()
 
 	pinger, err := createPinger(addr.String())
 	if err != nil {
-		return "", err
+		return err
 	}
 
 	pinger.Count = count
@@ -23,9 +24,11 @@ func RunDNSTest(ctx context.Context, addr common.Address, count int, timeout int
 	}
 	err = pinger.Run()
 	if err != nil {
-		return "", fmt.Errorf("failed to run ping: %v", err.Error())
+		return fmt.Errorf("failed to run ping: %v", err.Error())
 	}
 
-	output += joinS("resolving", pinger.IPAddr().String())
-	return common.Output(output), nil
+	result.Mu.Lock()
+	result.DNS = joinS("resolving", pinger.IPAddr().String())
+	result.Mu.Unlock()
+	return nil
 }