forked from gorilla/websocket
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
6 changed files
with
357 additions
and
7 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,134 @@ | ||
| // Copyright 2017 The Gorilla WebSocket Authors. All rights reserved. | ||
| // Use of this source code is governed by a BSD-style | ||
| // license that can be found in the LICENSE file. | ||
|
|
||
| // +build go1.7 | ||
|
|
||
| package websocket | ||
|
|
||
| import ( | ||
| "io" | ||
| "io/ioutil" | ||
| "sync/atomic" | ||
| "testing" | ||
| ) | ||
|
|
||
| // broadcastBench allows to run broadcast benchmarks. | ||
| // In every broadcast benchmark we create many connections, then send the same | ||
| // message into every connection and wait for all writes complete. This emulates | ||
| // an application where many connections listen to the same data - i.e. PUB/SUB | ||
| // scenarios with many subscribers in one channel. | ||
| type broadcastBench struct { | ||
| w io.Writer | ||
| message *broadcastMessage | ||
| closeCh chan struct{} | ||
| doneCh chan struct{} | ||
| count int32 | ||
| conns []*broadcastConn | ||
| compression bool | ||
| usePrepared bool | ||
| } | ||
|
|
||
| type broadcastMessage struct { | ||
| payload []byte | ||
| prepared *PreparedMessage | ||
| } | ||
|
|
||
| type broadcastConn struct { | ||
| conn *Conn | ||
| msgCh chan *broadcastMessage | ||
| } | ||
|
|
||
| func newBroadcastConn(c *Conn) *broadcastConn { | ||
| return &broadcastConn{ | ||
| conn: c, | ||
| msgCh: make(chan *broadcastMessage, 1), | ||
| } | ||
| } | ||
|
|
||
| func newBroadcastBench(usePrepared, compression bool) *broadcastBench { | ||
| bench := &broadcastBench{ | ||
| w: ioutil.Discard, | ||
| doneCh: make(chan struct{}), | ||
| closeCh: make(chan struct{}), | ||
| usePrepared: usePrepared, | ||
| compression: compression, | ||
| } | ||
| msg := &broadcastMessage{ | ||
| payload: textMessages(1)[0], | ||
| } | ||
| if usePrepared { | ||
| pm, _ := NewPreparedMessage(TextMessage, msg.payload) | ||
| msg.prepared = pm | ||
| } | ||
| bench.message = msg | ||
| bench.makeConns(10000) | ||
| return bench | ||
| } | ||
|
|
||
| func (b *broadcastBench) makeConns(numConns int) { | ||
| conns := make([]*broadcastConn, numConns) | ||
|
|
||
| for i := 0; i < numConns; i++ { | ||
| c := newConn(fakeNetConn{Reader: nil, Writer: b.w}, true, 1024, 1024) | ||
| if b.compression { | ||
| c.enableWriteCompression = true | ||
| c.newCompressionWriter = compressNoContextTakeover | ||
| } | ||
| conns[i] = newBroadcastConn(c) | ||
| go func(c *broadcastConn) { | ||
| for { | ||
| select { | ||
| case msg := <-c.msgCh: | ||
| if b.usePrepared { | ||
| c.conn.WritePreparedMessage(msg.prepared) | ||
| } else { | ||
| c.conn.WriteMessage(TextMessage, msg.payload) | ||
| } | ||
| val := atomic.AddInt32(&b.count, 1) | ||
| if val%int32(numConns) == 0 { | ||
| b.doneCh <- struct{}{} | ||
| } | ||
| case <-b.closeCh: | ||
| return | ||
| } | ||
| } | ||
| }(conns[i]) | ||
| } | ||
| b.conns = conns | ||
| } | ||
|
|
||
| func (b *broadcastBench) close() { | ||
| close(b.closeCh) | ||
| } | ||
|
|
||
| func (b *broadcastBench) runOnce() { | ||
| for _, c := range b.conns { | ||
| c.msgCh <- b.message | ||
| } | ||
| <-b.doneCh | ||
| } | ||
|
|
||
| func BenchmarkBroadcast(b *testing.B) { | ||
| benchmarks := []struct { | ||
| name string | ||
| usePrepared bool | ||
| compression bool | ||
| }{ | ||
| {"NoCompression", false, false}, | ||
| {"WithCompression", false, true}, | ||
| {"NoCompressionPrepared", true, false}, | ||
| {"WithCompressionPrepared", true, true}, | ||
| } | ||
| for _, bm := range benchmarks { | ||
| b.Run(bm.name, func(b *testing.B) { | ||
| bench := newBroadcastBench(bm.usePrepared, bm.compression) | ||
| defer bench.close() | ||
| b.ResetTimer() | ||
| for i := 0; i < b.N; i++ { | ||
| bench.runOnce() | ||
| } | ||
| b.ReportAllocs() | ||
| }) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| // Copyright 2017 The Gorilla WebSocket Authors. All rights reserved. | ||
| // Use of this source code is governed by a BSD-style | ||
| // license that can be found in the LICENSE file. | ||
|
|
||
| package websocket | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "net" | ||
| "sync" | ||
| "time" | ||
| ) | ||
|
|
||
| // PreparedMessage caches on the wire representations of a message payload. | ||
| // Use PreparedMessage to efficiently send a message payload to multiple | ||
| // connections. PreparedMessage is especially useful when compression is used | ||
| // because the CPU and memory expensive compression operation can be executed | ||
| // once for a given set of compression options. | ||
| type PreparedMessage struct { | ||
| messageType int | ||
| data []byte | ||
| err error | ||
| mu sync.Mutex | ||
| frames map[prepareKey]*preparedFrame | ||
| } | ||
|
|
||
| // prepareKey defines a unique set of options to cache prepared frames in PreparedMessage. | ||
| type prepareKey struct { | ||
| isServer bool | ||
| compress bool | ||
| compressionLevel int | ||
| } | ||
|
|
||
| // preparedFrame contains data in wire representation. | ||
| type preparedFrame struct { | ||
| once sync.Once | ||
| data []byte | ||
| } | ||
|
|
||
| // NewPreparedMessage returns an initialized PreparedMessage. You can then send | ||
| // it to connection using WritePreparedMessage method. Valid wire | ||
| // representation will be calculated lazily only once for a set of current | ||
| // connection options. | ||
| func NewPreparedMessage(messageType int, data []byte) (*PreparedMessage, error) { | ||
| pm := &PreparedMessage{ | ||
| messageType: messageType, | ||
| frames: make(map[prepareKey]*preparedFrame), | ||
| data: data, | ||
| } | ||
|
|
||
| // Prepare a plain server frame. | ||
| _, frameData, err := pm.frame(prepareKey{isServer: true, compress: false}) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| // To protect against caller modifying the data argument, remember the data | ||
| // copied to the plain server frame. | ||
| pm.data = frameData[len(frameData)-len(data):] | ||
| return pm, nil | ||
| } | ||
|
|
||
| func (pm *PreparedMessage) frame(key prepareKey) (int, []byte, error) { | ||
| pm.mu.Lock() | ||
| frame, ok := pm.frames[key] | ||
| if !ok { | ||
| frame = &preparedFrame{} | ||
| pm.frames[key] = frame | ||
| } | ||
| pm.mu.Unlock() | ||
|
|
||
| var err error | ||
| frame.once.Do(func() { | ||
| // Prepare a frame using a 'fake' connection. | ||
| // TODO: Refactor code in conn.go to allow more direct construction of | ||
| // the frame. | ||
| mu := make(chan bool, 1) | ||
| mu <- true | ||
| var nc prepareConn | ||
| c := &Conn{ | ||
| conn: &nc, | ||
| mu: mu, | ||
| isServer: key.isServer, | ||
| compressionLevel: key.compressionLevel, | ||
| enableWriteCompression: true, | ||
| writeBuf: make([]byte, defaultWriteBufferSize+maxFrameHeaderSize), | ||
| } | ||
| if key.compress { | ||
| c.newCompressionWriter = compressNoContextTakeover | ||
| } | ||
| err = c.WriteMessage(pm.messageType, pm.data) | ||
| frame.data = nc.buf.Bytes() | ||
| }) | ||
| return pm.messageType, frame.data, err | ||
| } | ||
|
|
||
| type prepareConn struct { | ||
| buf bytes.Buffer | ||
| net.Conn | ||
| } | ||
|
|
||
| func (pc *prepareConn) Write(p []byte) (int, error) { return pc.buf.Write(p) } | ||
| func (pc *prepareConn) SetWriteDeadline(t time.Time) error { return nil } |
Oops, something went wrong.