├── .gitignore ├── LICENSE ├── README.md ├── go.mod ├── go.sum ├── internal ├── random.go └── random_test.go ├── main.go └── pkg ├── broadcast └── broadcast.go └── echo └── echo.go /.gitignore: -------------------------------------------------------------------------------- 1 | # If you prefer the allow list template instead of the deny list, see community template: 2 | # https://github.com/github/gitignore/blob/main/community/Golang/Go.AllowList.gitignore 3 | # 4 | # Binaries for programs and plugins 5 | *.exe 6 | *.exe~ 7 | *.dll 8 | *.so 9 | *.dylib 10 | 11 | # Test binary, built with `go test -c` 12 | *.test 13 | 14 | # Output of the go coverage tool, specifically when used with LiteIDE 15 | *.out 16 | 17 | # Dependency directories (remove the comment below to include it) 18 | vendor/ 19 | bin/ 20 | assets/ 21 | .idea/ 22 | .vscode/ 23 | 24 | # Go workspace file 25 | go.work 26 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 道友请留步 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # wsbench 2 | 3 | ### install 4 | ```bash 5 | go install github.com/lxzan/wsbench@latest 6 | ``` 7 | 8 | ### help 9 | ``` 10 | NAME: 11 | wsbench - testing websocket server iops and latency 12 | 13 | USAGE: 14 | wsbench [global options] command [command options] [arguments...] 15 | 16 | COMMANDS: 17 | echo 18 | broadcast 19 | version 20 | help, h Shows a list of commands or help for one command 21 | 22 | GLOBAL OPTIONS: 23 | --help, -h show help 24 | ``` 25 | 26 | ### example 27 | 28 | ##### Echo 29 | ```bash 30 | wsbench echo -c 1000 -n 1000 -p 1000 -u 'ws://127.0.0.1:8000/connect,ws://127.0.0.1:8001/connect' 31 | ``` 32 | 33 | ##### Broadcast 34 | ```bash 35 | wsbench broadcast -c 1000 -n 1 -p 1000 -i 3 -u 'ws://127.0.0.1:8000/connect,ws://127.0.0.1:8001/connect' 36 | ``` -------------------------------------------------------------------------------- /go.mod: -------------------------------------------------------------------------------- 1 | module github.com/lxzan/wsbench 2 | 3 | go 1.19 4 | 5 | require ( 6 | github.com/lxzan/concurrency v1.2.0 7 | github.com/lxzan/gws v1.8.0-rc2 8 | github.com/rs/zerolog v1.29.1 9 | github.com/stretchr/testify v1.8.4 10 | github.com/urfave/cli/v2 v2.25.3 11 | ) 12 | 13 | require ( 14 | github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect 15 | github.com/davecgh/go-spew v1.1.1 // indirect 16 | github.com/dolthub/maphash v0.1.0 // indirect 17 | github.com/hashicorp/errwrap v1.0.0 // indirect 18 | github.com/hashicorp/go-multierror v1.1.1 // indirect 19 | github.com/klauspost/compress v1.17.5-0.20240119100516-32312d57f3c7 // indirect 20 | github.com/mattn/go-colorable v0.1.12 // indirect 21 | github.com/mattn/go-isatty v0.0.14 // indirect 22 | github.com/pmezard/go-difflib v1.0.0 // indirect 23 | github.com/russross/blackfriday/v2 v2.1.0 // indirect 24 | github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 // indirect 25 | golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6 // indirect 26 | gopkg.in/yaml.v3 v3.0.1 // indirect 27 | ) 28 | -------------------------------------------------------------------------------- /go.sum: -------------------------------------------------------------------------------- 1 | github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc= 2 | github.com/cpuguy83/go-md2man/v2 v2.0.2 h1:p1EgwI/C7NhT0JmVkwCD2ZBK8j4aeHQX2pMHHBfMQ6w= 3 | github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o= 4 | github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= 5 | github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= 6 | github.com/dolthub/maphash v0.1.0 h1:bsQ7JsF4FkkWyrP3oCnFJgrCUAFbFf3kOl4L/QxPDyQ= 7 | github.com/dolthub/maphash v0.1.0/go.mod h1:gkg4Ch4CdCDu5h6PMriVLawB7koZ+5ijb9puGMV50a4= 8 | github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA= 9 | github.com/hashicorp/errwrap v1.0.0 h1:hLrqtEDnRye3+sgx6z4qVLNuviH3MR5aQ0ykNJa/UYA= 10 | github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4= 11 | github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo= 12 | github.com/hashicorp/go-multierror v1.1.1/go.mod h1:iw975J/qwKPdAO1clOe2L8331t/9/fmwbPZ6JB6eMoM= 13 | github.com/klauspost/compress v1.17.5-0.20240119100516-32312d57f3c7 h1:cBVUjnVhEblq1CXwboJtC3F1lQYutr/4sVsnhLUAgiE= 14 | github.com/klauspost/compress v1.17.5-0.20240119100516-32312d57f3c7/go.mod h1:/dCuZOvVtNoHsyb+cuJD3itjs3NbnF6KH9zAO4BDxPM= 15 | github.com/lxzan/concurrency v1.2.0 h1:BVeP+RY34Ksu5JhQKHmGYGpKvkSIzAWROMc0MJd/1rU= 16 | github.com/lxzan/concurrency v1.2.0/go.mod h1:2z0NFVX9YGfZ+HFGzI2SI8vvQTJK4T/cm52g07MWVYw= 17 | github.com/lxzan/gws v1.8.0-rc2 h1:hRn9j0BKFfNqVIv2XLGH/+ohQvp3mP4L76WeLIz0IqU= 18 | github.com/lxzan/gws v1.8.0-rc2/go.mod h1:B2fVs+A1TECBKMNO5aiON+573ffF77c1bpDRB9UzSIo= 19 | github.com/mattn/go-colorable v0.1.12 h1:jF+Du6AlPIjs2BiUiQlKOX0rt3SujHxPnksPKZbaA40= 20 | github.com/mattn/go-colorable v0.1.12/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4= 21 | github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y= 22 | github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94= 23 | github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= 24 | github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= 25 | github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= 26 | github.com/rs/xid v1.4.0/go.mod h1:trrq9SKmegXys3aeAKXMUTdJsYXVwGY3RLcfgqegfbg= 27 | github.com/rs/zerolog v1.29.1 h1:cO+d60CHkknCbvzEWxP0S9K6KqyTjrCNUy1LdQLCGPc= 28 | github.com/rs/zerolog v1.29.1/go.mod h1:Le6ESbR7hc+DP6Lt1THiV8CQSdkkNrd3R0XbEgp3ZBU= 29 | github.com/russross/blackfriday/v2 v2.1.0 h1:JIOH55/0cWyOuilr9/qlrm0BSXldqnqwMsf35Ld67mk= 30 | github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM= 31 | github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= 32 | github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo= 33 | github.com/urfave/cli/v2 v2.25.3 h1:VJkt6wvEBOoSjPFQvOkv6iWIrsJyCrKGtCtxXWwmGeY= 34 | github.com/urfave/cli/v2 v2.25.3/go.mod h1:GHupkWPMM0M/sj1a2b4wUrWBPzazNrIjouW6fmdJLxc= 35 | github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673 h1:bAn7/zixMGCfxrRTfdpNzjtPYqr8smhKouy9mxVdGPU= 36 | github.com/xrash/smetrics v0.0.0-20201216005158-039620a65673/go.mod h1:N3UwUGtsrSj3ccvlPHLoLsHnpR27oXr4ZE984MbSER8= 37 | golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 38 | golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6 h1:foEbQz/B0Oz6YIqu/69kfXPYeFQAuuMYFkjaqXzl5Wo= 39 | golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= 40 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= 41 | gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= 42 | gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= 43 | gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= 44 | -------------------------------------------------------------------------------- /internal/random.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "math/rand" 5 | "sync" 6 | "time" 7 | ) 8 | 9 | type RandomString struct { 10 | mu sync.Mutex 11 | r *rand.Rand 12 | layout string 13 | } 14 | 15 | var ( 16 | AlphabetNumeric = &RandomString{ 17 | layout: "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ", 18 | r: rand.New(rand.NewSource(time.Now().UnixNano())), 19 | mu: sync.Mutex{}, 20 | } 21 | 22 | Numeric = &RandomString{ 23 | layout: "0123456789", 24 | r: rand.New(rand.NewSource(time.Now().UnixNano())), 25 | mu: sync.Mutex{}, 26 | } 27 | ) 28 | 29 | func (c *RandomString) Generate(n int) []byte { 30 | c.mu.Lock() 31 | var b = make([]byte, n, n) 32 | var length = len(c.layout) 33 | for i := 0; i < n; i++ { 34 | var idx = c.r.Intn(length) 35 | b[i] = c.layout[idx] 36 | } 37 | c.mu.Unlock() 38 | return b 39 | } 40 | 41 | func (c *RandomString) Intn(n int) int { 42 | c.mu.Lock() 43 | x := c.r.Intn(n) 44 | c.mu.Unlock() 45 | return x 46 | } 47 | 48 | func (c *RandomString) Uint32() uint32 { 49 | c.mu.Lock() 50 | x := c.r.Uint32() 51 | c.mu.Unlock() 52 | return x 53 | } 54 | -------------------------------------------------------------------------------- /internal/random_test.go: -------------------------------------------------------------------------------- 1 | package internal 2 | 3 | import ( 4 | "testing" 5 | 6 | "github.com/stretchr/testify/assert" 7 | ) 8 | 9 | func TestAlphabetNumeric(t *testing.T) { 10 | count := 1000 11 | set := make(map[string]struct{}, count) 12 | 13 | for i := 0; i < count; i++ { 14 | bs := AlphabetNumeric.Generate(100) 15 | if len(bs) != 100 { 16 | t.Error("generate error") 17 | return 18 | } 19 | set[string(bs)] = struct{}{} 20 | } 21 | 22 | assert.Equal(t, count, len(set)) 23 | } 24 | 25 | func TestNumeric(t *testing.T) { 26 | count := 1000 27 | for i := 0; i < count; i++ { 28 | num := AlphabetNumeric.Intn(100) 29 | 30 | // n >= 0 31 | assert.GreaterOrEqual(t, num, 0) 32 | // n <= 100 33 | assert.LessOrEqual(t, num, 100) 34 | } 35 | } 36 | -------------------------------------------------------------------------------- /main.go: -------------------------------------------------------------------------------- 1 | package main 2 | 3 | import ( 4 | "github.com/lxzan/wsbench/pkg/broadcast" 5 | "github.com/lxzan/wsbench/pkg/echo" 6 | "github.com/rs/zerolog" 7 | "github.com/rs/zerolog/log" 8 | cli "github.com/urfave/cli/v2" 9 | "os" 10 | ) 11 | 12 | const Version = "v1.1.0" 13 | 14 | func main() { 15 | log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr}) 16 | 17 | app := cli.App{ 18 | Name: "wsbench", 19 | Usage: "testing websocket server iops and latency", 20 | Commands: []*cli.Command{ 21 | echo.NewCommand(), 22 | broadcast.NewCommand(), 23 | { 24 | Name: "version", 25 | Action: func(context *cli.Context) error { 26 | println(Version) 27 | return nil 28 | }, 29 | }, 30 | }, 31 | } 32 | 33 | if err := app.Run(os.Args); err != nil { 34 | log.Fatal().Msg(err.Error()) 35 | } 36 | } 37 | -------------------------------------------------------------------------------- /pkg/broadcast/broadcast.go: -------------------------------------------------------------------------------- 1 | package broadcast 2 | 3 | import ( 4 | "crypto/tls" 5 | "github.com/lxzan/concurrency" 6 | "github.com/lxzan/gws" 7 | "github.com/lxzan/wsbench/internal" 8 | "github.com/rs/zerolog/log" 9 | "github.com/urfave/cli/v2" 10 | "os" 11 | "sync" 12 | "sync/atomic" 13 | "time" 14 | ) 15 | 16 | type Params struct { 17 | Serial int64 18 | Urls []string 19 | Compress bool 20 | Payload []byte 21 | PayloadSize int 22 | NumClient int 23 | NumMessage int 24 | Interval int 25 | } 26 | 27 | func NewCommand() *cli.Command { 28 | return &cli.Command{ 29 | Name: "broadcast", 30 | Flags: []cli.Flag{ 31 | &cli.StringSliceFlag{ 32 | Name: "u", 33 | Aliases: []string{"urls"}, 34 | Usage: "server address", 35 | }, 36 | &cli.IntFlag{ 37 | Name: "c", 38 | Usage: "number of multiple requests to make at a time", 39 | DefaultText: "100", 40 | Value: 100, 41 | Aliases: []string{"connection"}, 42 | }, 43 | &cli.IntFlag{ 44 | Name: "n", 45 | Usage: "number of requests to perform", 46 | DefaultText: "10000", 47 | Value: 10000, 48 | Aliases: []string{"message_num"}, 49 | }, 50 | &cli.IntFlag{ 51 | Name: "p", 52 | Usage: "payload size", 53 | DefaultText: "4000", 54 | Value: 4000, 55 | Aliases: []string{"payload_size"}, 56 | }, 57 | &cli.BoolFlag{ 58 | Name: "compress", 59 | Usage: "whether to turn on compression", 60 | DefaultText: "false", 61 | Value: false, 62 | }, 63 | &cli.IntFlag{ 64 | Name: "i", 65 | Usage: "message delivery interval", 66 | DefaultText: "10s", 67 | Value: 10, 68 | Aliases: []string{"interval"}, 69 | }, 70 | &cli.StringFlag{ 71 | Name: "f", 72 | Usage: "load payload content from file", 73 | DefaultText: "", 74 | Aliases: []string{"file"}, 75 | }, 76 | }, 77 | Action: Run, 78 | } 79 | } 80 | 81 | func Run(ctx *cli.Context) error { 82 | var params = &Params{} 83 | params.Urls = ctx.StringSlice("urls") 84 | params.NumClient = ctx.Int("connection") 85 | params.NumMessage = ctx.Int("message_num") 86 | params.PayloadSize = ctx.Int("payload_size") 87 | params.Payload = internal.AlphabetNumeric.Generate(params.PayloadSize) 88 | params.Compress = ctx.Bool("compress") 89 | params.Interval = ctx.Int("interval") 90 | if s := ctx.String("file"); len(s) > 0 { 91 | content, err := os.ReadFile(s) 92 | if err != nil { 93 | return err 94 | } 95 | params.Payload = content 96 | params.PayloadSize = len(content) 97 | } 98 | 99 | var handler = &Handler{ 100 | params: params, 101 | done: make(chan struct{}), 102 | sessions: &sync.Map{}, 103 | } 104 | 105 | var cc = concurrency.NewWorkerGroup[int]() 106 | for i := 0; i < params.NumClient; i++ { 107 | cc.Push(i) 108 | } 109 | cc.OnMessage = func(args int) error { 110 | socket, _, err := gws.NewClient(handler, &gws.ClientOption{ 111 | ReadBufferSize: 8 * 1024, 112 | Addr: handler.SelectURL(), 113 | TlsConfig: &tls.Config{InsecureSkipVerify: true}, 114 | PermessageDeflate: gws.PermessageDeflate{ 115 | Enabled: params.Compress, 116 | ServerContextTakeover: true, 117 | ClientContextTakeover: true, 118 | }, 119 | }) 120 | if err != nil { 121 | return err 122 | } 123 | handler.sessions.Store(socket, 1) 124 | go socket.ReadLoop() 125 | return nil 126 | } 127 | if err := cc.Start(); err != nil { 128 | return err 129 | } 130 | 131 | go func() { 132 | ticker := time.NewTicker(time.Duration(params.Interval) * time.Second) 133 | defer ticker.Stop() 134 | 135 | broadcaster := gws.NewBroadcaster(gws.OpcodeBinary, params.Payload) 136 | for { 137 | <-ticker.C 138 | 139 | handler.sessions.Range(func(key, value any) bool { 140 | socket := key.(*gws.Conn) 141 | for i := 0; i < params.NumMessage; i++ { 142 | _ = broadcaster.Broadcast(socket) 143 | } 144 | return true 145 | }) 146 | } 147 | }() 148 | 149 | handler.ShowProgress() 150 | return nil 151 | } 152 | 153 | type Handler struct { 154 | params *Params 155 | num int64 156 | sessions *sync.Map 157 | done chan struct{} 158 | } 159 | 160 | func (c *Handler) SelectURL() string { 161 | nextId := atomic.AddInt64(&c.params.Serial, 1) 162 | return c.params.Urls[nextId%int64(len(c.params.Urls))] 163 | } 164 | 165 | func (c *Handler) OnOpen(socket *gws.Conn) { _ = socket.SetNoDelay(false) } 166 | 167 | func (c *Handler) OnClose(socket *gws.Conn, err error) { 168 | if _, ok := err.(*gws.CloseError); !ok { 169 | log.Error().Msg(err.Error()) 170 | } 171 | os.Exit(0) 172 | } 173 | 174 | func (c *Handler) OnPing(socket *gws.Conn, payload []byte) {} 175 | 176 | func (c *Handler) OnPong(socket *gws.Conn, payload []byte) {} 177 | 178 | func (c *Handler) OnMessage(socket *gws.Conn, message *gws.Message) { 179 | defer message.Close() 180 | atomic.AddInt64(&c.num, 1) 181 | } 182 | 183 | func (c *Handler) ShowProgress() { 184 | ticker := time.NewTicker(time.Duration(c.params.Interval) * time.Second) 185 | defer ticker.Stop() 186 | for { 187 | <-ticker.C 188 | requests := atomic.LoadInt64(&c.num) 189 | log.Info().Int64("Requests", requests).Msg("") 190 | } 191 | } 192 | -------------------------------------------------------------------------------- /pkg/echo/echo.go: -------------------------------------------------------------------------------- 1 | package echo 2 | 3 | import ( 4 | "bytes" 5 | "crypto/tls" 6 | "encoding/binary" 7 | "encoding/json" 8 | "fmt" 9 | "github.com/lxzan/concurrency" 10 | "github.com/lxzan/gws" 11 | "github.com/lxzan/wsbench/internal" 12 | "github.com/rs/zerolog/log" 13 | "github.com/urfave/cli/v2" 14 | "os" 15 | "sync" 16 | "sync/atomic" 17 | "time" 18 | ) 19 | 20 | const M = 10000 21 | 22 | type Params struct { 23 | Serial int64 // 序列号 24 | Urls []string // 服务器地址列表 25 | Compress bool // 是否压缩 26 | Latency bool // 是否统计延迟 27 | PayloadSize int // 载荷大小 28 | NumClient int // 客户端数量 29 | NumMessage int64 // 消息数量 30 | Output string // 输出JSON文件目录 31 | Stats [M]uint64 // 统计 32 | Payload []byte // 载荷 33 | Concurrency int // 单连接并发度 34 | } 35 | 36 | func NewCommand() *cli.Command { 37 | return &cli.Command{ 38 | Name: "echo", 39 | Flags: []cli.Flag{ 40 | &cli.StringSliceFlag{ 41 | Name: "u", 42 | Aliases: []string{"urls"}, 43 | Usage: "server address", 44 | }, 45 | &cli.IntFlag{ 46 | Name: "c", 47 | Usage: "number of multiple requests to make at a time", 48 | DefaultText: "100", 49 | Value: 100, 50 | Aliases: []string{"connection"}, 51 | }, 52 | &cli.IntFlag{ 53 | Name: "n", 54 | Usage: "number of requests per connection to perform", 55 | DefaultText: "10000", 56 | Value: 10000, 57 | Aliases: []string{"message_num"}, 58 | }, 59 | &cli.IntFlag{ 60 | Name: "p", 61 | Usage: "payload Size", 62 | DefaultText: "4000", 63 | Value: 4000, 64 | Aliases: []string{"payload_size"}, 65 | }, 66 | &cli.StringFlag{ 67 | Name: "f", 68 | Usage: "load payload content from file", 69 | DefaultText: "", 70 | Aliases: []string{"file"}, 71 | }, 72 | &cli.BoolFlag{ 73 | Name: "compress", 74 | Usage: "whether to turn on compression", 75 | DefaultText: "false", 76 | Value: false, 77 | }, 78 | &cli.BoolFlag{ 79 | Name: "latency", 80 | Usage: "whether to turn on latency", 81 | DefaultText: "false", 82 | Value: false, 83 | }, 84 | &cli.IntFlag{ 85 | Name: "concurrency", 86 | Usage: "single-connection concurrency", 87 | DefaultText: "8", 88 | Value: 8, 89 | }, 90 | &cli.StringFlag{ 91 | Name: "o", 92 | Usage: "output json file path", 93 | DefaultText: "", 94 | Value: "", 95 | Aliases: []string{"output"}, 96 | }, 97 | }, 98 | Action: Run, 99 | } 100 | } 101 | 102 | func Run(ctx *cli.Context) error { 103 | var params = &Params{} 104 | params.Urls = ctx.StringSlice("urls") 105 | params.NumClient = ctx.Int("connection") 106 | params.NumMessage = ctx.Int64("message_num") 107 | params.PayloadSize = ctx.Int("payload_size") 108 | params.Compress = ctx.Bool("compress") 109 | params.Latency = ctx.Bool("latency") 110 | params.Output = ctx.String("output") 111 | params.Concurrency = ctx.Int("concurrency") 112 | params.Payload = internal.AlphabetNumeric.Generate(params.PayloadSize) 113 | 114 | if dir := ctx.String("file"); dir != "" { 115 | b, err := os.ReadFile(dir) 116 | if err != nil { 117 | return err 118 | } 119 | params.Payload = b 120 | params.PayloadSize = len(b) 121 | } 122 | 123 | var handler = &Handler{ 124 | params: params, 125 | pool: &sync.Pool{New: func() any { 126 | return bytes.NewBuffer(make([]byte, 0, params.PayloadSize+8)) 127 | }}, 128 | done: make(chan struct{}), 129 | sessions: &sync.Map{}, 130 | } 131 | 132 | var cc = concurrency.NewWorkerGroup[int]() 133 | for i := 0; i < params.NumClient; i++ { 134 | cc.Push(i) 135 | } 136 | cc.OnMessage = func(args int) error { 137 | socket, _, err := gws.NewClient(handler, &gws.ClientOption{ 138 | ParallelEnabled: true, 139 | ParallelGolimit: params.Concurrency, 140 | ReadBufferSize: 8 * 1024, 141 | PermessageDeflate: gws.PermessageDeflate{ 142 | Enabled: params.Compress, 143 | ServerContextTakeover: true, 144 | ClientContextTakeover: true, 145 | }, 146 | Addr: handler.SelectURL(), 147 | TlsConfig: &tls.Config{InsecureSkipVerify: true}, 148 | }) 149 | if err != nil { 150 | return err 151 | } 152 | handler.sessions.Store(socket, 1) 153 | go socket.ReadLoop() 154 | return nil 155 | } 156 | if err := cc.Start(); err != nil { 157 | return err 158 | } 159 | 160 | var t0 = time.Now() 161 | handler.sessions.Range(func(key, value any) bool { 162 | go func() { 163 | for i := 0; i < params.Concurrency; i++ { 164 | handler.SendMessage(key.(*gws.Conn), params.Payload) 165 | } 166 | }() 167 | return true 168 | }) 169 | 170 | go handler.ShowProgress() 171 | 172 | <-handler.done 173 | log.Info().Str("Percentage", "100.00%").Int("Requests", int(params.NumMessage)).Msg("") 174 | 175 | var iops = int(float64(params.NumMessage) / time.Since(t0).Seconds()) 176 | var logger = log.Info().Int("IOPS", iops).Str("Duration", time.Since(t0).String()) 177 | var output = map[string]any{ 178 | "iops": iops, 179 | "payload": params.PayloadSize, 180 | } 181 | if params.Latency { 182 | var p50 = handler.Report(50) 183 | var p90 = handler.Report(90) 184 | var p95 = handler.Report(95) 185 | var p99 = handler.Report(99) 186 | output["p50"] = p50 187 | output["p90"] = p90 188 | output["p95"] = p95 189 | output["p99"] = p99 190 | logger. 191 | Str("P50", p50). 192 | Str("P90", p90). 193 | Str("P95", p95). 194 | Str("P99", p99) 195 | } 196 | logger.Msg("") 197 | 198 | if params.Output != "" { 199 | file, err := os.OpenFile(params.Output, os.O_APPEND|os.O_CREATE|os.O_RDWR, 0644) 200 | if err != nil { 201 | return err 202 | } 203 | b, _ := json.Marshal(output) 204 | b = append(b, '\n') 205 | _, _ = file.Write(b) 206 | } 207 | return nil 208 | } 209 | 210 | type Handler struct { 211 | params *Params 212 | pool *sync.Pool 213 | numReceived int64 214 | numSend int64 215 | sessions *sync.Map 216 | done chan struct{} 217 | } 218 | 219 | func (c *Handler) SelectURL() string { 220 | nextId := atomic.AddInt64(&c.params.Serial, 1) 221 | return c.params.Urls[nextId%int64(len(c.params.Urls))] 222 | } 223 | 224 | func (c *Handler) OnOpen(socket *gws.Conn) { _ = socket.SetNoDelay(false) } 225 | 226 | func (c *Handler) OnClose(socket *gws.Conn, err error) { 227 | if _, ok := err.(*gws.CloseError); !ok { 228 | log.Error().Msg(err.Error()) 229 | } 230 | os.Exit(0) 231 | } 232 | 233 | func (c *Handler) OnPing(socket *gws.Conn, payload []byte) { _ = socket.WritePong(nil) } 234 | 235 | func (c *Handler) OnPong(socket *gws.Conn, payload []byte) {} 236 | 237 | func (c *Handler) OnMessage(socket *gws.Conn, message *gws.Message) { 238 | defer message.Close() 239 | 240 | n := message.Data.Len() 241 | if c.params.Latency { 242 | cost := uint64(M - 1) 243 | if n >= 8 { 244 | p := message.Bytes()[n-8:] 245 | cost = (uint64(time.Now().UnixNano()) - binary.BigEndian.Uint64(p)) / 1000000 246 | } 247 | if cost >= M { 248 | cost = M - 1 249 | } 250 | atomic.AddUint64(&c.params.Stats[cost], 1) 251 | } 252 | 253 | if x := atomic.AddInt64(&c.numReceived, 1); x == c.params.NumMessage { 254 | c.done <- struct{}{} 255 | return 256 | } 257 | c.SendMessage(socket, c.params.Payload) 258 | } 259 | 260 | func (c *Handler) SendMessage(socket *gws.Conn, payload []byte) { 261 | if atomic.AddInt64(&c.numSend, 1) <= c.params.NumMessage { 262 | buf := c.pool.Get().(*bytes.Buffer) 263 | buf.Reset() 264 | buf.Write(payload) 265 | p := buf.Bytes() 266 | if c.params.Latency { 267 | p = binary.BigEndian.AppendUint64(buf.Bytes(), uint64(time.Now().UnixNano())) 268 | } 269 | _ = socket.WriteMessage(gws.OpcodeBinary, p) 270 | c.pool.Put(buf) 271 | } 272 | } 273 | 274 | func (c *Handler) Report(rate int) string { 275 | sum := uint64(0) 276 | threshold := uint64(int64(rate) * c.params.NumMessage / 100) 277 | for i, v := range c.params.Stats { 278 | if v == 0 { 279 | continue 280 | } 281 | sum += v 282 | if sum >= threshold { 283 | if i == M-1 { 284 | return "∞" 285 | } 286 | return fmt.Sprintf("%dms", i) 287 | } 288 | } 289 | return "" 290 | } 291 | 292 | func (c *Handler) ShowProgress() { 293 | ticker := time.NewTicker(time.Second) 294 | defer ticker.Stop() 295 | for { 296 | <-ticker.C 297 | requests := atomic.LoadInt64(&c.numReceived) 298 | percentage := fmt.Sprintf("%.2f", float64(100*requests)/float64(c.params.NumMessage)) + "%" 299 | log.Info().Str("Percentage", percentage).Int64("Requests", requests).Msg("") 300 | } 301 | } 302 | --------------------------------------------------------------------------------