123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102 |
- // 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
- 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 struct{}, 1)
- mu <- struct{}{}
- 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 }
|