From dd9da23a595474054d65dcf2cfeb76b02ee56b9d Mon Sep 17 00:00:00 2001
From: yuhan6665 <1588741+yuhan6665@users.noreply.github.com>
Date: Sat, 4 Dec 2021 17:45:01 -0500
Subject: [PATCH] Add statistics in conn testing

Occasional error is observed when we execute long test
Print time and memory for better troubleshooting in the future

Co-authored-by: Jebbs <qjebbs@gmail.com>
---
 common/units/bytesize.go      | 100 ++++++++++++++++++++++++++++++++++
 common/units/bytesize_test.go |  66 ++++++++++++++++++++++
 testing/scenarios/common.go   |  14 ++++-
 3 files changed, 179 insertions(+), 1 deletion(-)
 create mode 100644 common/units/bytesize.go
 create mode 100644 common/units/bytesize_test.go

diff --git a/common/units/bytesize.go b/common/units/bytesize.go
new file mode 100644
index 00000000..bf3c84d3
--- /dev/null
+++ b/common/units/bytesize.go
@@ -0,0 +1,100 @@
+package units
+
+import (
+	"errors"
+	"strconv"
+	"strings"
+	"unicode"
+)
+
+var (
+	errInvalidSize = errors.New("invalid size")
+	errInvalidUnit = errors.New("invalid or unsupported unit")
+)
+
+// ByteSize is the size of bytes
+type ByteSize uint64
+
+const (
+	_ = iota
+	// KB = 1KB
+	KB ByteSize = 1 << (10 * iota)
+	// MB = 1MB
+	MB
+	// GB = 1GB
+	GB
+	// TB = 1TB
+	TB
+	// PB = 1PB
+	PB
+	// EB = 1EB
+	EB
+)
+
+func (b ByteSize) String() string {
+	unit := ""
+	value := float64(0)
+	switch {
+	case b == 0:
+		return "0"
+	case b < KB:
+		unit = "B"
+		value = float64(b)
+	case b < MB:
+		unit = "KB"
+		value = float64(b) / float64(KB)
+	case b < GB:
+		unit = "MB"
+		value = float64(b) / float64(MB)
+	case b < TB:
+		unit = "GB"
+		value = float64(b) / float64(GB)
+	case b < PB:
+		unit = "TB"
+		value = float64(b) / float64(TB)
+	case b < EB:
+		unit = "PB"
+		value = float64(b) / float64(PB)
+	default:
+		unit = "EB"
+		value = float64(b) / float64(EB)
+	}
+	result := strconv.FormatFloat(value, 'f', 2, 64)
+	result = strings.TrimSuffix(result, ".0")
+	return result + unit
+}
+
+// Parse parses ByteSize from string
+func (b *ByteSize) Parse(s string) error {
+	s = strings.TrimSpace(s)
+	s = strings.ToUpper(s)
+	i := strings.IndexFunc(s, unicode.IsLetter)
+	if i == -1 {
+		return errInvalidUnit
+	}
+
+	bytesString, multiple := s[:i], s[i:]
+	bytes, err := strconv.ParseFloat(bytesString, 64)
+	if err != nil || bytes <= 0 {
+		return errInvalidSize
+	}
+	switch multiple {
+	case "B":
+		*b = ByteSize(bytes)
+	case "K", "KB", "KIB":
+		*b = ByteSize(bytes * float64(KB))
+	case "M", "MB", "MIB":
+		*b = ByteSize(bytes * float64(MB))
+	case "G", "GB", "GIB":
+		*b = ByteSize(bytes * float64(GB))
+	case "T", "TB", "TIB":
+		*b = ByteSize(bytes * float64(TB))
+	case "P", "PB", "PIB":
+		*b = ByteSize(bytes * float64(PB))
+	case "E", "EB", "EIB":
+		*b = ByteSize(bytes * float64(EB))
+	default:
+		return errInvalidUnit
+	}
+	return nil
+}
diff --git a/common/units/bytesize_test.go b/common/units/bytesize_test.go
new file mode 100644
index 00000000..3b8b9c88
--- /dev/null
+++ b/common/units/bytesize_test.go
@@ -0,0 +1,66 @@
+package units_test
+
+import (
+	"testing"
+
+	"github.com/xtls/xray-core/common/units"
+)
+
+func TestByteSizes(t *testing.T) {
+	size := units.ByteSize(0)
+	assertSizeString(t, size, "0")
+	size++
+	assertSizeValue(t,
+		assertSizeString(t, size, "1.00B"),
+		size,
+	)
+	size <<= 10
+	assertSizeValue(t,
+		assertSizeString(t, size, "1.00KB"),
+		size,
+	)
+	size <<= 10
+	assertSizeValue(t,
+		assertSizeString(t, size, "1.00MB"),
+		size,
+	)
+	size <<= 10
+	assertSizeValue(t,
+		assertSizeString(t, size, "1.00GB"),
+		size,
+	)
+	size <<= 10
+	assertSizeValue(t,
+		assertSizeString(t, size, "1.00TB"),
+		size,
+	)
+	size <<= 10
+	assertSizeValue(t,
+		assertSizeString(t, size, "1.00PB"),
+		size,
+	)
+	size <<= 10
+	assertSizeValue(t,
+		assertSizeString(t, size, "1.00EB"),
+		size,
+	)
+}
+
+func assertSizeValue(t *testing.T, size string, expected units.ByteSize) {
+	actual := units.ByteSize(0)
+	err := actual.Parse(size)
+	if err != nil {
+		t.Error(err)
+	}
+	if actual != expected {
+		t.Errorf("expect %s, but got %s", expected, actual)
+	}
+}
+
+func assertSizeString(t *testing.T, size units.ByteSize, expected string) string {
+	actual := size.String()
+	if actual != expected {
+		t.Errorf("expect %s, but got %s", expected, actual)
+	}
+	return expected
+}
diff --git a/testing/scenarios/common.go b/testing/scenarios/common.go
index 67535dfa..d9f66e64 100644
--- a/testing/scenarios/common.go
+++ b/testing/scenarios/common.go
@@ -23,6 +23,7 @@ import (
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/retry"
 	"github.com/xtls/xray-core/common/serial"
+	"github.com/xtls/xray-core/common/units"
 	core "github.com/xtls/xray-core/core"
 )
 
@@ -198,7 +199,18 @@ func testUDPConn(port net.Port, payloadSize int, timeout time.Duration) func() e
 }
 
 func testTCPConn2(conn net.Conn, payloadSize int, timeout time.Duration) func() error {
-	return func() error {
+	return func() (err1 error) {
+		start := time.Now()
+		defer func() {
+			var m runtime.MemStats
+			runtime.ReadMemStats(&m)
+			// For info on each, see: https://golang.org/pkg/runtime/#MemStats
+			fmt.Println("testConn finishes:", time.Since(start).Milliseconds(), "ms\t",
+				err1, "\tAlloc =", units.ByteSize(m.Alloc).String(),
+				"\tTotalAlloc =", units.ByteSize(m.TotalAlloc).String(),
+				"\tSys =", units.ByteSize(m.Sys).String(),
+				"\tNumGC =", m.NumGC)
+		}()
 		payload := make([]byte, payloadSize)
 		common.Must2(rand.Read(payload))