From d685f35111ce8800a82e19eed6673723d9c7b498 Mon Sep 17 00:00:00 2001
From: Andrey Petrov <andrey.petrov@shazow.net>
Date: Mon, 21 Oct 2019 14:19:18 +0200
Subject: [PATCH] wip: Duration

---
 main.go    | 86 +++++++++++++++++++-----------------------------------
 request.go | 59 +++++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 56 deletions(-)
 create mode 100644 request.go

diff --git a/main.go b/main.go
index 0107103..bb0d8fa 100644
--- a/main.go
+++ b/main.go
@@ -2,12 +2,11 @@ package main // import "github.com/INFURA/versus"
 
 import (
 	"bufio"
-	"bytes"
 	"context"
 	"fmt"
 	"log"
 	"os"
-	"strings"
+	"os/signal"
 	"time"
 
 	flags "github.com/jessevdk/go-flags"
@@ -18,6 +17,7 @@ const chanBuffer = 20
 // Options contains the flag options
 type Options struct {
 	Endpoints []string `positional-args:"yes"`
+	Duration  string   `long:"duration" description:"Stop after some duration."`
 }
 
 func exit(code int, format string, args ...interface{}) {
@@ -35,17 +35,42 @@ func main() {
 		return
 	}
 
+	// Setup signals
+	ctx, abort := context.WithCancel(context.Background())
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt)
+	go func(abort context.CancelFunc) {
+		<-sigCh
+		abort()
+		<-sigCh
+		panic("aborted")
+	}(abort)
+
+	if options.Duration != "" {
+		d, err := time.ParseDuration(options.Duration)
+		if err != nil {
+			exit(1, "failed to parse duration: %s", err)
+		}
+		timeoutCtx, cancel := context.WithTimeout(ctx, d)
+		defer cancel()
+		ctx = timeoutCtx
+	}
+
+	// Launch clients
 	clients := make(Clients, 0, len(options.Endpoints))
 	for _, endpoint := range options.Endpoints {
 		c, err := NewClient(endpoint)
 		if err != nil {
-			exit(1, "failed to create client: %s", err)
+			exit(2, "failed to create client: %s", err)
 		}
+		go func() {
+			if err := c.Serve(ctx); err != nil {
+				exit(2, "failed to start client: %s", err)
+			}
+		}()
 		clients = append(clients, *c)
 	}
 
-	ctx := context.Background()
-
 	scanner := bufio.NewScanner(os.Stdin)
 	for scanner.Scan() {
 		select {
@@ -65,54 +90,3 @@ func main() {
 		return
 	}
 }
-
-type Request struct {
-	client *Client
-	id     int
-
-	Line      []byte
-	Timestamp time.Time
-}
-
-func (req *Request) Do() {
-	body, err := req.client.t.Send(req.Line)
-	resp := Response{
-		req.client, body, err, time.Now().Sub(req.Timestamp),
-	}
-	req.client.out <- resp
-}
-
-type Response struct {
-	client *Client
-
-	Body []byte
-	Err  error
-
-	Elapsed time.Duration
-}
-
-func (r *Response) Equal(other Response) bool {
-	return r.Err.Error() == other.Err.Error() && bytes.Equal(r.Body, other.Body)
-}
-
-type Responses []Response
-
-func (r Responses) String() string {
-	var buf strings.Builder
-
-	// TODO: Sort before printing
-	last := r[0]
-	for i, resp := range r {
-		fmt.Fprintf(&buf, "\t%s", resp.Elapsed)
-
-		if resp.Err.Error() != last.Err.Error() {
-			fmt.Fprintf(&buf, "[%d: error mismatch: %s != %s]", i, resp.Err, last.Err)
-		}
-
-		if !bytes.Equal(resp.Body, last.Body) {
-			fmt.Fprintf(&buf, "[%d: body mismatch: %s != %s]", i, resp.Body[:20], last.Body[:20])
-		}
-	}
-
-	return buf.String()
-}
diff --git a/request.go b/request.go
new file mode 100644
index 0000000..2c75fe3
--- /dev/null
+++ b/request.go
@@ -0,0 +1,59 @@
+package main
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"time"
+)
+
+type Request struct {
+	client *Client
+	id     int
+
+	Line      []byte
+	Timestamp time.Time
+}
+
+func (req *Request) Do() {
+	body, err := req.client.t.Send(req.Line)
+	resp := Response{
+		req.client, body, err, time.Now().Sub(req.Timestamp),
+	}
+	req.client.out <- resp
+}
+
+type Response struct {
+	client *Client
+
+	Body []byte
+	Err  error
+
+	Elapsed time.Duration
+}
+
+func (r *Response) Equal(other Response) bool {
+	return r.Err.Error() == other.Err.Error() && bytes.Equal(r.Body, other.Body)
+}
+
+type Responses []Response
+
+func (r Responses) String() string {
+	var buf strings.Builder
+
+	// TODO: Sort before printing
+	last := r[0]
+	for i, resp := range r {
+		fmt.Fprintf(&buf, "\t%s", resp.Elapsed)
+
+		if resp.Err.Error() != last.Err.Error() {
+			fmt.Fprintf(&buf, "[%d: error mismatch: %s != %s]", i, resp.Err, last.Err)
+		}
+
+		if !bytes.Equal(resp.Body, last.Body) {
+			fmt.Fprintf(&buf, "[%d: body mismatch: %s != %s]", i, resp.Body[:20], last.Body[:20])
+		}
+	}
+
+	return buf.String()
+}