github.com/metaworking/channeld@v0.7.3/pkg/channeld/connection_websocket_test.go (about)

     1  package channeld
     2  
     3  import (
     4  	"net/http"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/gorilla/websocket"
    10  	"github.com/metaworking/channeld/internal/testpb"
    11  	"google.golang.org/protobuf/proto"
    12  )
    13  
    14  func getBenchmarkBytes() []byte {
    15  	testMsg := &testpb.TestMapMessage{
    16  		Kv:  make(map[uint32]string),
    17  		Kv2: make(map[uint32]*testpb.TestMapMessage_StringWrapper),
    18  	}
    19  	testMsg.Kv[1] = "a"
    20  	testMsg.Kv[2] = "b"
    21  	testMsg.Kv[3] = "c"
    22  	testMsg.Kv[4] = "d"
    23  
    24  	testMsg.Kv2[1] = &testpb.TestMapMessage_StringWrapper{Content: "a"}
    25  	testMsg.Kv2[2] = &testpb.TestMapMessage_StringWrapper{Content: "b", Num: 2}
    26  
    27  	bytes, _ := proto.Marshal(testMsg)
    28  	return bytes
    29  }
    30  
    31  const writeTimes = 100000
    32  
    33  var wg sync.WaitGroup
    34  
    35  func TestGorillaWebSocket(t *testing.T) {
    36  	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    37  		conn, err := upgrader.Upgrade(w, r, nil)
    38  		if err == nil {
    39  			data := getBenchmarkBytes()
    40  
    41  			startTime := time.Now()
    42  			w, _ := conn.NextWriter(websocket.BinaryMessage)
    43  			for i := 0; i < writeTimes; i++ {
    44  				w.Write(data)
    45  			}
    46  			w.Close()
    47  			t.Logf("Write %d times: %dms", writeTimes, time.Since(startTime).Milliseconds())
    48  
    49  			conn.Close()
    50  		} else {
    51  			t.Errorf("failed to upgrade websocket connection: %v", err)
    52  		}
    53  		wg.Done()
    54  	})
    55  	wg.Add(1)
    56  	go http.ListenAndServe("localhost:8081", nil)
    57  	time.Sleep(time.Second)
    58  
    59  	conn, _, err := websocket.DefaultDialer.Dial("ws://localhost:8081", nil)
    60  	if err == nil {
    61  		startTime := time.Now()
    62  		_, bytes, err := conn.ReadMessage()
    63  		if err == nil {
    64  			t.Logf("read in %dms (%d bytes)\n", time.Since(startTime).Milliseconds(), len(bytes))
    65  		} else {
    66  			t.Errorf("error reading message: %v", err)
    67  		}
    68  	} else {
    69  		t.Fatalf("failed dialing: %v", err)
    70  	}
    71  
    72  	wg.Wait()
    73  
    74  	// Result: 10x FASTER than nhooyr/websocket!
    75  	// write 10,000 times: 2ms; read in 3ms
    76  	// write 100,000 times: 24-25ms; read in 24-25ms
    77  }
    78  
    79  /*
    80  func TestNhooyrWebSocket(t *testing.T) {
    81  	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
    82  		conn, err := nws.Accept(w, r, nil)
    83  		if err == nil {
    84  			data := getBenchmarkBytes()
    85  			ctx := r.Context()
    86  
    87  			startTime := time.Now()
    88  			w, _ := conn.Writer(ctx, nws.MessageBinary)
    89  			for i := 0; i < writeTimes; i++ {
    90  				w.Write(data)
    91  			}
    92  			w.Close()
    93  			t.Logf("write %d times: %dms", writeTimes, time.Since(startTime).Milliseconds())
    94  
    95  			conn.Close(nws.StatusNormalClosure, "finished benchmark")
    96  		} else {
    97  			t.Errorf("failed to upgrade websocket connection: %v", err)
    98  		}
    99  		wg.Done()
   100  	})
   101  	wg.Add(1)
   102  	go http.ListenAndServe("localhost:8081", nil)
   103  	time.Sleep(time.Second)
   104  
   105  	clientCtx := context.Background()
   106  	conn, _, err := nws.Dial(clientCtx, "ws://localhost:8081", nil)
   107  	if err == nil {
   108  		conn.SetReadLimit(0xffffff)
   109  		startTime := time.Now()
   110  		_, bytes, err := conn.Read(clientCtx)
   111  		if err == nil {
   112  			t.Logf("read in %dms (%d bytes)\n", time.Since(startTime).Milliseconds(), len(bytes))
   113  		} else {
   114  			t.Errorf("error reading message: %v", err)
   115  		}
   116  	} else {
   117  		t.Fatalf("failed dialing: %v", err)
   118  	}
   119  
   120  	wg.Wait()
   121  
   122  	// Result:
   123  	// write 10,000 times: 19-22ms; read in 28-32ms
   124  	// write 100,000 times: 197-215ms; read in 290-317ms
   125  }
   126  */