@@ -1962,3 +1962,134 @@ func assertReadMessage(ctx context.Context, c *websocket.Conn, typ websocket.Mes
1962
1962
}
1963
1963
return assertEqualf (p , actP , "unexpected frame %v payload" , actTyp )
1964
1964
}
1965
+
1966
+ func BenchmarkConn (b * testing.B ) {
1967
+ sizes := []int {
1968
+ 2 ,
1969
+ 16 ,
1970
+ 32 ,
1971
+ 512 ,
1972
+ 4096 ,
1973
+ 16384 ,
1974
+ }
1975
+
1976
+ b .Run ("write" , func (b * testing.B ) {
1977
+ for _ , size := range sizes {
1978
+ b .Run (strconv .Itoa (size ), func (b * testing.B ) {
1979
+ b .Run ("stream" , func (b * testing.B ) {
1980
+ benchConn (b , false , true , size )
1981
+ })
1982
+ b .Run ("buffer" , func (b * testing.B ) {
1983
+ benchConn (b , false , false , size )
1984
+ })
1985
+ })
1986
+ }
1987
+ })
1988
+
1989
+ b .Run ("echo" , func (b * testing.B ) {
1990
+ for _ , size := range sizes {
1991
+ b .Run (strconv .Itoa (size ), func (b * testing.B ) {
1992
+ benchConn (b , true , true , size )
1993
+ })
1994
+ }
1995
+ })
1996
+ }
1997
+
1998
+ func benchConn (b * testing.B , echo , stream bool , size int ) {
1999
+ s , closeFn := testServer (b , func (w http.ResponseWriter , r * http.Request ) error {
2000
+ c , err := websocket .Accept (w , r , nil )
2001
+ if err != nil {
2002
+ return err
2003
+ }
2004
+ if echo {
2005
+ wsecho .Loop (r .Context (), c )
2006
+ } else {
2007
+ discardLoop (r .Context (), c )
2008
+ }
2009
+ return nil
2010
+ }, false )
2011
+ defer closeFn ()
2012
+
2013
+ wsURL := strings .Replace (s .URL , "http" , "ws" , 1 )
2014
+
2015
+ ctx , cancel := context .WithTimeout (context .Background (), time .Minute * 5 )
2016
+ defer cancel ()
2017
+
2018
+ c , _ , err := websocket .Dial (ctx , wsURL , nil )
2019
+ if err != nil {
2020
+ b .Fatal (err )
2021
+ }
2022
+ defer c .Close (websocket .StatusInternalError , "" )
2023
+
2024
+ msg := []byte (strings .Repeat ("2" , size ))
2025
+ readBuf := make ([]byte , len (msg ))
2026
+ b .SetBytes (int64 (len (msg )))
2027
+ b .ReportAllocs ()
2028
+ b .ResetTimer ()
2029
+ for i := 0 ; i < b .N ; i ++ {
2030
+ if stream {
2031
+ w , err := c .Writer (ctx , websocket .MessageText )
2032
+ if err != nil {
2033
+ b .Fatal (err )
2034
+ }
2035
+
2036
+ _ , err = w .Write (msg )
2037
+ if err != nil {
2038
+ b .Fatal (err )
2039
+ }
2040
+
2041
+ err = w .Close ()
2042
+ if err != nil {
2043
+ b .Fatal (err )
2044
+ }
2045
+ } else {
2046
+ err = c .Write (ctx , websocket .MessageText , msg )
2047
+ if err != nil {
2048
+ b .Fatal (err )
2049
+ }
2050
+ }
2051
+
2052
+ if echo {
2053
+ _ , r , err := c .Reader (ctx )
2054
+ if err != nil {
2055
+ b .Fatal (err )
2056
+ }
2057
+
2058
+ _ , err = io .ReadFull (r , readBuf )
2059
+ if err != nil {
2060
+ b .Fatal (err )
2061
+ }
2062
+ }
2063
+ }
2064
+ b .StopTimer ()
2065
+
2066
+ c .Close (websocket .StatusNormalClosure , "" )
2067
+ }
2068
+
2069
+ func discardLoop (ctx context.Context , c * websocket.Conn ) {
2070
+ defer c .Close (websocket .StatusInternalError , "" )
2071
+
2072
+ ctx , cancel := context .WithTimeout (ctx , time .Minute )
2073
+ defer cancel ()
2074
+
2075
+ b := make ([]byte , 32768 )
2076
+ echo := func () error {
2077
+ _ , r , err := c .Reader (ctx )
2078
+ if err != nil {
2079
+ return err
2080
+ }
2081
+
2082
+ _ , err = io .CopyBuffer (ioutil .Discard , r , b )
2083
+ if err != nil {
2084
+ return err
2085
+ }
2086
+ return nil
2087
+ }
2088
+
2089
+ for {
2090
+ err := echo ()
2091
+ if err != nil {
2092
+ return
2093
+ }
2094
+ }
2095
+ }
0 commit comments