writer.go

package gws

import (
	"bytes"
	"errors"
	"github.com/lxzan/gws/internal"
	"math"
	"sync"
	"sync/atomic"
)

// WriteClose
// code: https://developer.mozilla.org/zh-CN/docs/Web/API/CloseEvent#status_codes
// &#36890;&#36807;emitError&#21457;&#36865;&#20851;&#38381;&#24103;, &#23558;&#36830;&#25509;&#29366;&#24577;&#32622;&#20026;&#20851;&#38381;, &#29992;&#20110;&#26381;&#21153;&#31471;&#20027;&#21160;&#26029;&#24320;&#36830;&#25509;
// &#27809;&#26377;&#29305;&#27530;&#21407;&#22240;&#30340;&#35805;, &#24314;&#35758;code=0, reason=nil
// Send a close frame via emitError to set the connection state to closed, for server-initiated disconnection
// If there is no special reason, we suggest code=0, reason=nil
func (c *Conn) WriteClose(code uint16, reason []byte) {
	var err = internal.NewError(internal.StatusCode(code), internal.GwsError(""))
	if len(reason) > 0 {
		err.Err = errors.New(string(reason))
	}
	c.emitError(err)
}

// WritePing write ping frame
func (c *Conn) WritePing(payload []byte) error {
	return c.WriteMessage(OpcodePing, payload)
}

// WritePong write pong frame
func (c *Conn) WritePong(payload []byte) error {
	return c.WriteMessage(OpcodePong, payload)
}

// WriteString write text frame
func (c *Conn) WriteString(s string) error {
	return c.WriteMessage(OpcodeText, internal.StringToBytes(s))
}

// WriteAsync &#24322;&#27493;&#38750;&#38459;&#22622;&#22320;&#20889;&#20837;&#28040;&#24687;
// Write messages asynchronously and non-blockingly
func (c *Conn) WriteAsync(opcode Opcode, payload []byte) error {
	frame, index, err := c.genFrame(opcode, payload)
	if err != nil {
		c.emitError(err)
		return err
	}

	c.writeQueue.Push(func() {
		if c.isClosed() {
			return
		}
		err = internal.WriteN(c.conn, frame.Bytes(), frame.Len())
		myBufferPool.Put(frame, index)
		c.emitError(err)
	})
	return nil
}

// WriteMessage &#21457;&#36865;&#28040;&#24687;
func (c *Conn) WriteMessage(opcode Opcode, payload []byte) error {
	if c.isClosed() {
		return internal.ErrConnClosed
	}
	err := c.doWrite(opcode, payload)
	c.emitError(err)
	return err
}

// &#25191;&#34892;&#20889;&#20837;&#36923;&#36753;, &#20851;&#38381;&#29366;&#24577;&#32622;&#20026;1&#21518;&#36824;&#33021;&#20889;, &#20197;&#20415;&#21457;&#36865;&#20851;&#38381;&#24103;
// Execute the write logic, and write after the close state is set to 1, so that the close frame can be sent
func (c *Conn) doWrite(opcode Opcode, payload []byte) error {
	frame, index, err := c.genFrame(opcode, payload)
	if err != nil {
		return err
	}

	err = internal.WriteN(c.conn, frame.Bytes(), frame.Len())
	myBufferPool.Put(frame, index)
	return err
}

// &#24103;&#29983;&#25104;
func (c *Conn) genFrame(opcode Opcode, payload []byte) (*bytes.Buffer, int, error) {
	// &#19981;&#35201;&#21024;&#38500; opcode == OpcodeText
	if opcode == OpcodeText && !c.isTextValid(opcode, payload) {
		return nil, 0, internal.NewError(internal.CloseUnsupportedData, internal.ErrTextEncoding)
	}

	if c.compressEnabled && opcode.isDataFrame() && len(payload) >= c.config.CompressThreshold {
		return c.compressData(opcode, payload)
	}

	var n = len(payload)
	if n > c.config.WriteMaxPayloadSize {
		return nil, 0, internal.CloseMessageTooLarge
	}

	var header = frameHeader{}
	headerLength, maskBytes := header.GenerateHeader(c.isServer, true, false, opcode, n)
	var totalSize = n + headerLength
	var buf, index = myBufferPool.Get(totalSize)
	buf.Write(header[:headerLength])
	buf.Write(payload)
	var contents = buf.Bytes()
	if !c.isServer {
		internal.MaskXOR(contents[headerLength:], maskBytes)
	}
	return buf, index, nil
}

func (c *Conn) compressData(opcode Opcode, payload []byte) (*bytes.Buffer, int, error) {
	var buf, index = myBufferPool.Get(len(payload) / compressionRate)
	buf.Write(myPadding[0:])
	err := c.config.compressors.Select().Compress(payload, buf)
	if err != nil {
		return nil, 0, err
	}
	var contents = buf.Bytes()
	var payloadSize = buf.Len() - frameHeaderSize
	if payloadSize > c.config.WriteMaxPayloadSize {
		return nil, 0, internal.CloseMessageTooLarge
	}
	var header = frameHeader{}
	headerLength, maskBytes := header.GenerateHeader(c.isServer, true, true, opcode, payloadSize)
	if !c.isServer {
		internal.MaskXOR(contents[frameHeaderSize:], maskBytes)
	}
	copy(contents[frameHeaderSize-headerLength:], header[:headerLength])
	buf.Next(frameHeaderSize - headerLength)
	return buf, index, nil
}

type (
	Broadcaster struct {
		opcode  Opcode
		payload []byte
		msgs    [2]*broadcastMessageWrapper
		state   int64
	}

	broadcastMessageWrapper struct {
		once  sync.Once
		err   error
		index int
		frame *bytes.Buffer
	}
)

// NewBroadcaster
// &#30456;&#27604;WriteAsync, Broadcaster&#21482;&#20250;&#21387;&#32553;&#19968;&#27425;&#28040;&#24687;, &#21487;&#20197;&#33410;&#30465;&#22823;&#37327;CPU&#24320;&#38144;.
// Compared to WriteAsync, Broadcaster compresses the message only once, saving a lot of CPU overhead.
func NewBroadcaster(opcode Opcode, payload []byte) *Broadcaster {
	c := &Broadcaster{
		opcode:  opcode,
		payload: payload,
		msgs:    [2]*broadcastMessageWrapper{},
		state:   int64(math.MaxInt32),
	}
	return c
}

// Broadcast &#24191;&#25773;
// &#21521;&#21333;&#20010;&#23458;&#25143;&#31471;&#21457;&#36865;&#24191;&#25773;&#28040;&#24687;. &#27880;&#24847;: &#19981;&#35201;&#24182;&#34892;&#35843;&#29992;Broadcast&#26041;&#27861;
// Send a broadcast message to a single client. Note: Do not call the Broadcast method in parallel.
func (c *Broadcaster) Broadcast(socket *Conn) error {
	var idx = internal.SelectValue(socket.compressEnabled, 1, 0)
	var msg = c.msgs[idx]
	if msg == nil {
		c.msgs[idx] = &broadcastMessageWrapper{}
		msg = c.msgs[idx]
		msg.frame, msg.index, msg.err = socket.genFrame(c.opcode, c.payload)
	}
	if msg.err != nil {
		return msg.err
	}

	atomic.AddInt64(&c.state, 1)
	socket.writeQueue.Push(func() {
		if !socket.isClosed() {
			socket.emitError(internal.WriteN(socket.conn, msg.frame.Bytes(), msg.frame.Len()))
		}
		if atomic.AddInt64(&c.state, -1) == 0 {
			c.doClose()
		}
	})
	return nil
}

func (c *Broadcaster) doClose() {
	for _, item := range c.msgs {
		if item != nil {
			myBufferPool.Put(item.frame, item.index)
		}
	}
}

// Release
// &#22312;&#23436;&#25104;&#25152;&#26377;Broadcast&#20043;&#21518;&#35843;&#29992;Release&#26041;&#27861;&#37322;&#25918;&#36164;&#28304;.
// Call the Release method after all the Broadcasts have been completed to release the resources.
func (c *Broadcaster) Release() {
	if atomic.AddInt64(&c.state, -1*math.MaxInt32) == 0 {
		c.doClose()
	}
}