101 lines
2.7 KiB
Go
101 lines
2.7 KiB
Go
|
package socket
|
||
|
|
|
||
|
|
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 := &SocketConn{
|
||
|
|
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 }
|