test(bench): improve server and client

Author: ernado

internal/cmd/ch-bench-numbers/main.go

@@ -9,6 +9,7 @@ import (
 	"runtime/trace"
 	"time"
 
+	"github.com/dustin/go-humanize"
 	"github.com/go-faster/errors"
 	"go.uber.org/multierr"
 
@@ -25,7 +26,7 @@ func run(ctx context.Context) (re error) {
 		BlockSize int
 		Address   string
 	}
-	flag.IntVar(&arg.Count, "n", 20, "count")
+	flag.IntVar(&arg.Count, "n", 1, "count")
 	flag.IntVar(&arg.Numbers, "numbers", 500_000_000, "numbers count")
 	flag.IntVar(&arg.BlockSize, "block-size", 65_536, "maximum row count in block")
 	flag.StringVar(&arg.Profile, "profile", "", "cpu profile")
@@ -87,6 +88,12 @@ func run(ctx context.Context) (re error) {
 		}); err != nil {
 			return errors.Wrap(err, "query")
 		}
+		duration := time.Since(start)
+		gotBytes := uint64(arg.Numbers * 8)
+		fmt.Println(duration.Round(time.Millisecond), arg.Numbers, "rows",
+			humanize.Bytes(gotBytes),
+			humanize.Bytes(uint64(float64(gotBytes)/duration.Seconds()))+"/s",
+		)
 		fmt.Println(time.Since(start))
 	}
 

internal/cmd/ch-bench-server/main.go

@@ -9,11 +9,12 @@ import (
 
 	"github.com/dustin/go-humanize"
 	"github.com/go-faster/errors"
+	"golang.org/x/sync/errgroup"
 
 	"github.com/go-faster/ch/proto"
 )
 
-func run(_ context.Context) (re error) {
+func run(ctx context.Context) error {
 	ln, err := net.Listen("tcp4", "127.0.0.1:9001")
 	if err != nil {
 		return errors.Wrap(err, "listen")
@@ -54,45 +55,61 @@ func run(_ context.Context) (re error) {
 
 	fmt.Println("starting", "with chunk of", humanize.Bytes(uint64(len(raw))))
 
-	for {
-		conn, err := ln.Accept()
-		if err != nil {
-			return errors.Wrap(err, "accept")
-		}
-
+	process := func(conn net.Conn) error {
 		go func() {
+			// Drain input.
 			_, _ = io.Copy(io.Discard, conn)
 		}()
-		go func() {
-			defer func() { _ = conn.Close() }()
-			b := new(proto.Buffer)
-			b.EncodeAware(&proto.ServerHello{
-				Name:        "ch-bench-server",
-				Major:       0,
-				Minor:       12,
-				Revision:    proto.Version,
-				Timezone:    "UTC",
-				DisplayName: "Bench",
-				Patch:       1,
-			}, proto.Version)
-			if _, err := conn.Write(b.Buf); err != nil {
-				return
-			}
+		defer func() { _ = conn.Close() }()
+		b := new(proto.Buffer)
+		b.EncodeAware(&proto.ServerHello{
+			Name:        "ch-bench-server",
+			Major:       0,
+			Minor:       12,
+			Revision:    proto.Version,
+			Timezone:    "UTC",
+			DisplayName: "Bench",
+			Patch:       1,
+		}, proto.Version)
+		if _, err := conn.Write(b.Buf); err != nil {
+			return errors.Wrap(err, "write server hello")
+		}
 
-			for i := 0; i < chunks; i++ {
-				if _, err := conn.Write(raw); err != nil {
-					return
-				}
+		for i := 0; i < chunks; i++ {
+			if _, err := conn.Write(raw); err != nil {
+				return errors.Wrap(err, "write chunk")
 			}
+		}
+
+		// End of data.
+		b.Reset()
+		proto.ServerCodeEndOfStream.Encode(b)
+		if _, err := conn.Write(b.Buf); err != nil {
+			return errors.Wrap(err, "write end of stream")
+		}
+
+		return nil
+	}
 
-			// End of data.
-			b.Reset()
-			proto.ServerCodeEndOfStream.Encode(b)
-			if _, err := conn.Write(b.Buf); err != nil {
-				return
+	g, ctx := errgroup.WithContext(ctx)
+	for i := 0; i < 10; i++ {
+		g.Go(func() error {
+			for {
+				if ctx.Err() != nil {
+					return ctx.Err()
+				}
+				conn, err := ln.Accept()
+				if err != nil {
+					return errors.Wrap(err, "accept")
+				}
+				if err := process(conn); err != nil {
+					return errors.Wrap(err, "process")
+				}
 			}
-		}()
+		})
 	}
+
+	return g.Wait()
 }
 
 func main() {