udp throughput test update

examples/throughput/main.go

@@ -1,19 +1,21 @@
 package main
 
 import (
+	"crypto/rand"
 	"encoding/binary"
 	"encoding/hex"
 	"flag"
 	"fmt"
+	"github.com/nknorg/tuna/udp"
 	"log"
 	"math"
 	"net"
 	"os"
 	"strings"
 	"time"
 
-	ncp "github.com/nknorg/ncp-go"
-	nkn "github.com/nknorg/nkn-sdk-go"
+	"github.com/nknorg/ncp-go"
+	"github.com/nknorg/nkn-sdk-go"
 	ts "github.com/nknorg/nkn-tuna-session"
 	"github.com/nknorg/tuna/geo"
 )
@@ -23,6 +25,9 @@ const (
 	listenID = "bob"
 )
 
+var udpBytesReceived = 0
+var udpBytesSend = 0
+
 func read(sess net.Conn) error {
 	timeStart := time.Now()
 
@@ -95,6 +100,47 @@ func write(sess net.Conn, numBytes int) error {
 	return nil
 }
 
+func readUDP(conn *udp.EncryptUDPConn) error {
+	defer conn.Close()
+	buffer := make([]byte, 1024)
+	for {
+		n, _, err := conn.ReadFromUDP(buffer)
+		if err != nil {
+			return err
+		}
+		udpBytesReceived += n
+	}
+}
+
+func writeUDP(conn *udp.EncryptUDPConn, numBytes int) error {
+	defer conn.Close()
+	buffer := make([]byte, 1024)
+	timeStart := time.Now()
+	for {
+		rand.Read(buffer)
+		n, _, err := conn.WriteMsgUDP(buffer, nil, nil)
+		if err != nil {
+			return err
+		}
+		time.Sleep(100 * time.Nanosecond)
+		udpBytesSend += n
+		if udpBytesSend > numBytes {
+			break
+		}
+		if ((udpBytesSend - n) * 10 / numBytes) != (udpBytesSend * 10 / numBytes) {
+			mbTobytes := math.Pow(2, 20)
+			sent := float64(udpBytesSend) / mbTobytes
+			received := float64(udpBytesReceived) / mbTobytes
+			timeEnd := time.Now()
+			elapsed := timeEnd.Sub(timeStart)
+			speed := sent / elapsed.Seconds()
+			log.Printf("UDP: Sent %.2f MB bytes, received %.2f MB bytes, speed: %.2f MB/s, package loss:  %.2f%% \n", sent, received, speed, 100*(1-received/sent))
+		}
+	}
+
+	return nil
+}
+
 func main() {
 	numTunaListeners := flag.Int("n", 1, "number of tuna listeners")
 	numBytes := flag.Int("m", 1, "data to send (MB)")
@@ -109,6 +155,7 @@ func main() {
 	tunaMeasureBandwidth := flag.Bool("tmb", false, "tuna measure bandwidth")
 	tunaMaxPrice := flag.String("price", "0.01", "tuna reverse service max price in unit of NKN/MB")
 	mtu := flag.Int("mtu", 0, "ncp session mtu")
+	u := flag.Bool("u", false, "test udp")
 
 	flag.Parse()
 
@@ -188,6 +235,19 @@ func main() {
 				}(s)
 			}
 		}()
+
+		if *u {
+			uConn, err := c.ListenUDP(nil)
+			if err != nil {
+				log.Fatal(err)
+			}
+			go func() {
+				err = readUDP(uConn)
+				if err != nil {
+					log.Fatal(err)
+				}
+			}()
+		}
 	}
 
 	if *dial {
@@ -226,6 +286,19 @@ func main() {
 				time.Sleep(time.Millisecond * 100)
 			}
 		}()
+
+		if *u {
+			udpConn, err := c.DialUDPWithConfig(*dialAddr, dialConfig)
+			if err != nil {
+				log.Fatal(err)
+			}
+			go func() {
+				err := writeUDP(udpConn, *numBytes)
+				if err != nil {
+					log.Fatal(err)
+				}
+			}()
+		}
 	}
 
 	select {}

tests/session_test.go

@@ -3,6 +3,7 @@ package tests
 import (
 	"bytes"
 	"crypto/rand"
+	"encoding/hex"
 	"errors"
 	"github.com/nknorg/ncp-go"
 	"github.com/nknorg/nkn-sdk-go"
@@ -14,6 +15,7 @@ import (
 	"net"
 	"os"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 )
@@ -102,20 +104,6 @@ func TestTunaSession(t *testing.T) {
 		}
 	}()
 
-	go func() {
-		buffer := make([]byte, 65536)
-		for {
-			n, addr, err := uConn.ReadFromUDP(buffer)
-			if err != nil {
-				t.Fatal(err)
-			}
-			n, _, err = uConn.WriteMsgUDP(buffer[:n], nil, addr)
-			if err != nil {
-				t.Fatal(err)
-			}
-		}
-	}()
-
 	mm, err := nkn.NewMultiClient(account2, dialID, 4, false, clientConfig)
 	if err != nil {
 		t.Fatal(err)
@@ -148,7 +136,7 @@ func TestTunaSession(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	err = testUDP(udpConn)
+	err = testUDP(udpConn, uConn)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -169,16 +157,56 @@ func testTCP(conn net.Conn) error {
 	return nil
 }
 
-func testUDP(conn *udp.EncryptUDPConn) error {
-	send := make([]byte, 4096)
-	receive := make([]byte, 4096)
-	for i := 0; i < 10; i++ {
-		rand.Read(send)
-		conn.WriteMsgUDP(send, nil, nil)
-		conn.ReadFromUDP(receive)
-		if !bytes.Equal(send, receive) {
-			return errors.New("bytes not equal")
+func testUDP(from, to *udp.EncryptUDPConn) error {
+	count := 1000
+	sendList := make([]string, count)
+	recvList := make([]string, count)
+	sendNum := 0
+	recvNum := 0
+	var wg sync.WaitGroup
+	var e error
+	go func() {
+		wg.Add(1)
+		receive := make([]byte, 1024)
+		for i := 0; i < count; i++ {
+			_, _, err := to.ReadFromUDP(receive)
+			if err != nil {
+				e = err
+				return
+			}
+			recvNum++
+			recvList = append(recvList, hex.EncodeToString(receive))
+		}
+		wg.Done()
+	}()
+
+	go func() {
+		time.Sleep(1 * time.Second)
+		send := make([]byte, 1024)
+		wg.Add(1)
+		for i := 0; i < count; i++ {
+			rand.Read(send)
+			_, _, err := from.WriteMsgUDP(send, nil, nil)
+			if err != nil {
+				e = err
+				return
+			}
+			sendNum++
+			sendList = append(sendList, hex.EncodeToString(send))
 		}
+		wg.Done()
+	}()
+
+	wg.Wait()
+	if sendNum != recvNum {
+		return errors.New("package lost")
 	}
-	return nil
+
+	for i := 0; i < sendNum; i++ {
+		if sendList[i] != recvList[i] {
+			return errors.New("data mismatch")
+		}
+	}
+
+	return e
 }