123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242 |
- // Copyright (c) 2012-2014 Jeremy Latt
- // Copyright (c) 2016-2017 Daniel Oaks <daniel@danieloaks.net>
- // released under the MIT license
-
- package irc
-
- import (
- "errors"
- "io"
- "sync"
- )
-
- var (
- errSendQExceeded = errors.New("SendQ exceeded")
-
- sendQExceededMessage = []byte("\r\nERROR :SendQ Exceeded\r\n")
- )
-
- // Socket represents an IRC socket.
- type Socket struct {
- sync.Mutex
-
- conn IRCConn
-
- maxSendQBytes int
-
- // this is a trylock enforcing that only one goroutine can write to `conn` at a time
- writeLock sync.Mutex
-
- buffers [][]byte
- totalLength int
- closed bool
- sendQExceeded bool
- finalData []byte // what to send when we die
- finalized bool
- }
-
- // NewSocket returns a new Socket.
- func NewSocket(conn IRCConn, maxSendQBytes int) *Socket {
- result := Socket{
- conn: conn,
- maxSendQBytes: maxSendQBytes,
- }
- return &result
- }
-
- // Close stops a Socket from being able to send/receive any more data.
- func (socket *Socket) Close() {
- socket.Lock()
- socket.closed = true
- socket.Unlock()
-
- socket.wakeWriter()
- }
-
- // Read returns a single IRC line from a Socket.
- func (socket *Socket) Read() (string, error) {
- // immediately fail if Close() has been called, even if there's
- // still data in a bufio.Reader or websocket buffer:
- if socket.IsClosed() {
- return "", io.EOF
- }
-
- lineBytes, err := socket.conn.ReadLine()
- line := string(lineBytes)
-
- if err == io.EOF {
- socket.Close()
- }
-
- return line, err
- }
-
- // Write sends the given string out of Socket. Requirements:
- // 1. MUST NOT block for macroscopic amounts of time
- // 2. MUST NOT reorder messages
- // 3. MUST provide mutual exclusion for socket.conn.Write
- // 4. SHOULD NOT tie up additional goroutines, beyond the one blocked on socket.conn.Write
- func (socket *Socket) Write(data []byte) (err error) {
- if len(data) == 0 {
- return
- }
-
- socket.Lock()
- if socket.closed {
- err = io.EOF
- } else {
- prospectiveLen := socket.totalLength + len(data)
- if prospectiveLen > socket.maxSendQBytes {
- socket.sendQExceeded = true
- socket.closed = true
- err = errSendQExceeded
- } else {
- socket.buffers = append(socket.buffers, data)
- socket.totalLength = prospectiveLen
- }
- }
- socket.Unlock()
-
- socket.wakeWriter()
- return
- }
-
- // BlockingWrite sends the given string out of Socket. Requirements:
- // 1. MUST block until the message is sent
- // 2. MUST bypass sendq (calls to BlockingWrite cannot, on their own, cause a sendq overflow)
- // 3. MUST provide mutual exclusion for socket.conn.Write
- // 4. MUST respect the same ordering guarantees as Write (i.e., if a call to Write that sends
- // message m1 happens-before a call to BlockingWrite that sends message m2,
- // m1 must be sent on the wire before m2
- //
- // Callers MUST be writing to the client's socket from the client's own goroutine;
- // other callers must use the nonblocking Write call instead. Otherwise, a client
- // with a slow/unreliable connection risks stalling the progress of the system as a whole.
- func (socket *Socket) BlockingWrite(data []byte) (err error) {
- if len(data) == 0 {
- return
- }
-
- // after releasing the semaphore, we must check for fresh data, same as `send`
- defer func() {
- if socket.readyToWrite() {
- socket.wakeWriter()
- }
- }()
-
- // blocking acquire of the trylock
- socket.writeLock.Lock()
- defer socket.writeLock.Unlock()
-
- // first, flush any buffered data, to preserve the ordering guarantees
- closed := socket.performWrite()
- if closed {
- return io.EOF
- }
-
- err = socket.conn.WriteLine(data)
- if err != nil {
- socket.finalize()
- }
- return
- }
-
- // wakeWriter starts the goroutine that actually performs the write, without blocking
- func (socket *Socket) wakeWriter() {
- if socket.writeLock.TryLock() {
- // acquired the trylock; send() will release it
- go socket.send()
- }
- // else: do nothing, the holder will check for more data after releasing it
- }
-
- // SetFinalData sets the final data to send when the SocketWriter closes.
- func (socket *Socket) SetFinalData(data []byte) {
- socket.Lock()
- defer socket.Unlock()
- socket.finalData = data
- }
-
- // IsClosed returns whether the socket is closed.
- func (socket *Socket) IsClosed() bool {
- socket.Lock()
- defer socket.Unlock()
- return socket.closed
- }
-
- // is there data to write?
- func (socket *Socket) readyToWrite() bool {
- socket.Lock()
- defer socket.Unlock()
- // on the first time observing socket.closed, we still have to write socket.finalData
- return !socket.finalized && (socket.totalLength > 0 || socket.closed)
- }
-
- // send actually writes messages to socket.Conn; it may block
- func (socket *Socket) send() {
- for {
- // we are holding the trylock: actually do the write
- socket.performWrite()
- // surrender the trylock, avoiding a race where a write comes in after we've
- // checked readyToWrite() and it returned false, but while we still hold the trylock:
- socket.writeLock.Unlock()
- // check if more data came in while we held the trylock:
- if !socket.readyToWrite() {
- return
- }
- if !socket.writeLock.TryLock() {
- // failed to acquire; exit and wait for the holder to observe readyToWrite()
- // after releasing it
- return
- }
- // got the lock again, loop back around and write
- }
- }
-
- // write the contents of the buffer, then see if we need to close
- // returns whether we closed
- func (socket *Socket) performWrite() (closed bool) {
- // retrieve the buffered data, clear the buffer
- socket.Lock()
- buffers := socket.buffers
- socket.buffers = nil
- socket.totalLength = 0
- closed = socket.closed
- socket.Unlock()
-
- var err error
- if 0 < len(buffers) {
- err = socket.conn.WriteLines(buffers)
- }
-
- closed = closed || err != nil
- if closed {
- socket.finalize()
- }
- return
- }
-
- // mark closed and send final data. you must be holding the semaphore to call this:
- func (socket *Socket) finalize() {
- // mark the socket closed (if someone hasn't already), then write error lines
- socket.Lock()
- socket.closed = true
- finalized := socket.finalized
- socket.finalized = true
- finalData := socket.finalData
- if socket.sendQExceeded {
- finalData = sendQExceededMessage
- }
- socket.Unlock()
-
- if finalized {
- return
- }
-
- if len(finalData) != 0 {
- socket.conn.WriteLine(finalData)
- }
-
- // close the connection
- socket.conn.Close()
- }
|