melody/session.go

package melody

import (
	"github.com/gorilla/websocket"
	"time"
)

type Session struct {
	Conn   *websocket.Conn
	output chan *envelope
	config *Config
}

func newSession(config *Config, conn *websocket.Conn) *Session {
	return &Session{
		Conn:   conn,
		output: make(chan *envelope, config.MessageBufferSize),
		config: config,
	}
}

func (s *Session) writeMessage(message *envelope) {
	s.output <- message
}

func (s *Session) writeRaw(message *envelope) error {
	s.Conn.SetWriteDeadline(time.Now().Add(s.config.WriteWait))
	return s.Conn.WriteMessage(message.t, message.msg)
}

func (s *Session) close() {
	s.writeRaw(&envelope{t: websocket.CloseMessage, msg: []byte{}})
}

func (s *Session) ping() {
	s.writeMessage(&envelope{t: websocket.PingMessage, msg: []byte{}})
}

func (s *Session) writePump(errorHandler handleErrorFunc) {
	defer s.Conn.Close()

	ticker := time.NewTicker(s.config.PingPeriod)
	defer ticker.Stop()

	for {
		select {
		case msg, ok := <-s.output:
			if !ok {
				s.close()
				return
			}
			if err := s.writeRaw(msg); err != nil {
				go errorHandler(s, err)
				return
			}
		case <-ticker.C:
			s.ping()
		}
	}
}

func (s *Session) readPump(messageHandler handleMessageFunc, errorHandler handleErrorFunc) {
	defer s.Conn.Close()

	s.Conn.SetReadLimit(s.config.MaxMessageSize)
	s.Conn.SetReadDeadline(time.Now().Add(s.config.PongWait))

	s.Conn.SetPongHandler(func(string) error {
		s.Conn.SetReadDeadline(time.Now().Add(s.config.PongWait))
		return nil
	})

	for {
		_, message, err := s.Conn.ReadMessage()

		if err != nil {
			go errorHandler(s, err)
			break
		}

		go messageHandler(s, message)
	}
}

func (s *Session) Write(msg []byte) {
	s.writeMessage(&envelope{t: websocket.TextMessage, msg: msg})
}

func (s *Session) Close() {
	s.writeMessage(&envelope{t: websocket.CloseMessage, msg: []byte{}})
}