gosocketio/vendor/github.com/googollee/go-socket.io/redis_broadcast.go

package socketio

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"sync"

	"github.com/gomodule/redigo/redis"
)

// redisBroadcast gives Join, Leave & BroadcastTO server API support to socket.io along with room management
// map of rooms where each room contains a map of connection id to connections in that room
type redisBroadcast struct {
	pub *redis.PubSubConn
	sub *redis.PubSubConn

	nsp        string
	uid        string
	key        string
	reqChannel string
	resChannel string

	requests map[string]interface{}

	rooms map[string]map[string]Conn

	lock sync.RWMutex
}

// request types
const (
	roomLenReqType   = "0"
	clearRoomReqType = "1"
	allRoomReqType   = "2"
)

// request structs
type roomLenRequest struct {
	RequestType string
	RequestID   string
	Room        string
	numSub      int        `json:"-"`
	msgCount    int        `json:"-"`
	connections int        `json:"-"`
	mutex       sync.Mutex `json:"-"`
	done        chan bool  `json:"-"`
}

type clearRoomRequest struct {
	RequestType string
	RequestID   string
	Room        string
	UUID        string
}

type allRoomRequest struct {
	RequestType string
	RequestID   string
	rooms       map[string]bool `json:"-"`
	numSub      int             `json:"-"`
	msgCount    int             `json:"-"`
	mutex       sync.Mutex      `json:"-"`
	done        chan bool       `json:"-"`
}

// response struct
type roomLenResponse struct {
	RequestType string
	RequestID   string
	Connections int
}

type allRoomResponse struct {
	RequestType string
	RequestID   string
	Rooms       []string
}

func newRedisBroadcast(nsp string, opts *RedisAdapterOptions) (*redisBroadcast, error) {
	addr := opts.getAddr()
	var redisOpts []redis.DialOption
	if len(opts.Password) > 0 {
		redisOpts = append(redisOpts, redis.DialPassword(opts.Password))
	}
	if opts.DB > 0 {
		redisOpts = append(redisOpts, redis.DialDatabase(opts.DB))
	}

	pub, err := redis.Dial(opts.Network, addr, redisOpts...)
	if err != nil {
		return nil, err
	}

	sub, err := redis.Dial(opts.Network, addr, redisOpts...)
	if err != nil {
		return nil, err
	}

	subConn := &redis.PubSubConn{Conn: sub}
	pubConn := &redis.PubSubConn{Conn: pub}

	if err = subConn.PSubscribe(fmt.Sprintf("%s#%s#*", opts.Prefix, nsp)); err != nil {
		return nil, err
	}

	uid := newV4UUID()
	rbc := &redisBroadcast{
		rooms:      make(map[string]map[string]Conn),
		requests:   make(map[string]interface{}),
		sub:        subConn,
		pub:        pubConn,
		key:        fmt.Sprintf("%s#%s#%s", opts.Prefix, nsp, uid),
		reqChannel: fmt.Sprintf("%s-request#%s", opts.Prefix, nsp),
		resChannel: fmt.Sprintf("%s-response#%s", opts.Prefix, nsp),
		nsp:        nsp,
		uid:        uid,
	}

	if err = subConn.Subscribe(rbc.reqChannel, rbc.resChannel); err != nil {
		return nil, err
	}

	go rbc.dispatch()

	return rbc, nil
}

// AllRooms gives list of all rooms available for redisBroadcast.
func (bc *redisBroadcast) AllRooms() []string {
	req := allRoomRequest{
		RequestType: allRoomReqType,
		RequestID:   newV4UUID(),
	}
	reqJSON, _ := json.Marshal(&req)

	req.rooms = make(map[string]bool)
	numSub, _ := bc.getNumSub(bc.reqChannel)
	req.numSub = numSub
	req.done = make(chan bool, 1)

	bc.requests[req.RequestID] = &req
	_, err := bc.pub.Conn.Do("PUBLISH", bc.reqChannel, reqJSON)
	if err != nil {
		return []string{} // if error occurred,return empty
	}

	<-req.done

	rooms := make([]string, 0, len(req.rooms))
	for room := range req.rooms {
		rooms = append(rooms, room)
	}

	delete(bc.requests, req.RequestID)
	return rooms
}

// Join joins the given connection to the redisBroadcast room.
func (bc *redisBroadcast) Join(room string, connection Conn) {
	bc.lock.Lock()
	defer bc.lock.Unlock()

	if _, ok := bc.rooms[room]; !ok {
		bc.rooms[room] = make(map[string]Conn)
	}

	bc.rooms[room][connection.ID()] = connection
}

// Leave leaves the given connection from given room (if exist)
func (bc *redisBroadcast) Leave(room string, connection Conn) {
	bc.lock.Lock()
	defer bc.lock.Unlock()

	if connections, ok := bc.rooms[room]; ok {
		delete(connections, connection.ID())

		if len(connections) == 0 {
			delete(bc.rooms, room)
		}
	}
}

// LeaveAll leaves the given connection from all rooms.
func (bc *redisBroadcast) LeaveAll(connection Conn) {
	bc.lock.Lock()
	defer bc.lock.Unlock()

	for room, connections := range bc.rooms {
		delete(connections, connection.ID())

		if len(connections) == 0 {
			delete(bc.rooms, room)
		}
	}
}

// Clear clears the room.
func (bc *redisBroadcast) Clear(room string) {
	bc.lock.Lock()
	defer bc.lock.Unlock()

	delete(bc.rooms, room)
	go bc.publishClear(room)
}

// Send sends given event & args to all the connections in the specified room.
func (bc *redisBroadcast) Send(room, event string, args ...interface{}) {
	bc.lock.RLock()
	defer bc.lock.RUnlock()

	connections, ok := bc.rooms[room]
	if ok {
		for _, connection := range connections {
			connection.Emit(event, args...)
		}
	}

	bc.publishMessage(room, event, args...)
}

// SendAll sends given event & args to all the connections to all the rooms.
func (bc *redisBroadcast) SendAll(event string, args ...interface{}) {
	bc.lock.RLock()
	defer bc.lock.RUnlock()

	for _, connections := range bc.rooms {
		for _, connection := range connections {
			connection.Emit(event, args...)
		}
	}
	bc.publishMessage("", event, args...)
}

// ForEach sends data returned by DataFunc, if room does not exits sends nothing.
func (bc *redisBroadcast) ForEach(room string, f EachFunc) {
	bc.lock.RLock()
	defer bc.lock.RUnlock()

	occupants, ok := bc.rooms[room]
	if !ok {
		return
	}

	for _, connection := range occupants {
		f(connection)
	}
}

// Len gives number of connections in the room.
func (bc *redisBroadcast) Len(room string) int {
	req := roomLenRequest{
		RequestType: roomLenReqType,
		RequestID:   newV4UUID(),
		Room:        room,
	}

	reqJSON, err := json.Marshal(&req)
	if err != nil {
		return -1
	}

	numSub, err := bc.getNumSub(bc.reqChannel)
	if err != nil {
		return -1
	}

	req.numSub = numSub

	req.done = make(chan bool, 1)

	bc.requests[req.RequestID] = &req
	_, err = bc.pub.Conn.Do("PUBLISH", bc.reqChannel, reqJSON)
	if err != nil {
		return -1
	}

	<-req.done

	delete(bc.requests, req.RequestID)
	return req.connections
}

// Rooms gives the list of all the rooms available for redisBroadcast in case of
// no connection is given, in case of a connection is given, it gives
// list of all the rooms the connection is joined to.
func (bc *redisBroadcast) Rooms(connection Conn) []string {
	bc.lock.RLock()
	defer bc.lock.RUnlock()

	if connection == nil {
		return bc.AllRooms()
	}

	return bc.getRoomsByConn(connection)
}

func (bc *redisBroadcast) onMessage(channel string, msg []byte) error {
	channelParts := strings.Split(channel, "#")
	nsp := channelParts[len(channelParts)-2]
	if bc.nsp != nsp {
		return nil
	}

	uid := channelParts[len(channelParts)-1]
	if bc.uid == uid {
		return nil
	}

	var bcMessage map[string][]interface{}
	err := json.Unmarshal(msg, &bcMessage)
	if err != nil {
		return errors.New("invalid broadcast message")
	}

	args := bcMessage["args"]
	opts := bcMessage["opts"]

	room, ok := opts[0].(string)
	if !ok {
		return errors.New("invalid room")
	}

	event, ok := opts[1].(string)
	if !ok {
		return errors.New("invalid event")
	}

	if room != "" {
		bc.send(room, event, args...)
	} else {
		bc.sendAll(event, args...)
	}

	return nil
}

// Get the number of subscribers of a channel.
func (bc *redisBroadcast) getNumSub(channel string) (int, error) {
	rs, err := bc.pub.Conn.Do("PUBSUB", "NUMSUB", channel)
	if err != nil {
		return 0, err
	}

	numSub64, ok := rs.([]interface{})[1].(int64)
	if !ok {
		return 0, errors.New("redis reply cast to int error")
	}
	return int(numSub64), nil
}

// Handle request from redis channel.
func (bc *redisBroadcast) onRequest(msg []byte) {
	var req map[string]string

	if err := json.Unmarshal(msg, &req); err != nil {
		return
	}

	var res interface{}
	switch req["RequestType"] {
	case roomLenReqType:
		res = roomLenResponse{
			RequestType: req["RequestType"],
			RequestID:   req["RequestID"],
			Connections: len(bc.rooms[req["Room"]]),
		}
		bc.publish(bc.resChannel, &res)

	case allRoomReqType:
		res := allRoomResponse{
			RequestType: req["RequestType"],
			RequestID:   req["RequestID"],
			Rooms:       bc.allRooms(),
		}
		bc.publish(bc.resChannel, &res)

	case clearRoomReqType:
		if bc.uid == req["UUID"] {
			return
		}
		bc.clear(req["Room"])

	default:
	}
}

func (bc *redisBroadcast) publish(channel string, msg interface{}) {
	resJSON, err := json.Marshal(msg)
	if err != nil {
		return
	}

	_, err = bc.pub.Conn.Do("PUBLISH", channel, resJSON)
	if err != nil {
		return
	}
}

// Handle response from redis channel.
func (bc *redisBroadcast) onResponse(msg []byte) {
	var res map[string]interface{}

	err := json.Unmarshal(msg, &res)
	if err != nil {
		return
	}

	req, ok := bc.requests[res["RequestID"].(string)]
	if !ok {
		return
	}

	switch res["RequestType"] {
	case roomLenReqType:
		roomLenReq := req.(*roomLenRequest)

		roomLenReq.mutex.Lock()
		roomLenReq.msgCount++
		roomLenReq.connections += int(res["Connections"].(float64))
		roomLenReq.mutex.Unlock()

		if roomLenReq.numSub == roomLenReq.msgCount {
			roomLenReq.done <- true
		}

	case allRoomReqType:
		allRoomReq := req.(*allRoomRequest)
		rooms, ok := res["Rooms"].([]interface{})
		if !ok {
			allRoomReq.done <- true
			return
		}

		allRoomReq.mutex.Lock()
		allRoomReq.msgCount++
		for _, room := range rooms {
			allRoomReq.rooms[room.(string)] = true
		}
		allRoomReq.mutex.Unlock()

		if allRoomReq.numSub == allRoomReq.msgCount {
			allRoomReq.done <- true
		}

	default:
	}
}

func (bc *redisBroadcast) publishClear(room string) {
	req := clearRoomRequest{
		RequestType: clearRoomReqType,
		RequestID:   newV4UUID(),
		Room:        room,
		UUID:        bc.uid,
	}

	bc.publish(bc.reqChannel, &req)
}

func (bc *redisBroadcast) clear(room string) {
	bc.lock.Lock()
	defer bc.lock.Unlock()

	delete(bc.rooms, room)
}

func (bc *redisBroadcast) send(room string, event string, args ...interface{}) {
	bc.lock.RLock()
	defer bc.lock.RUnlock()

	connections, ok := bc.rooms[room]
	if !ok {
		return
	}

	for _, connection := range connections {
		connection.Emit(event, args...)
	}
}

func (bc *redisBroadcast) publishMessage(room string, event string, args ...interface{}) {
	opts := make([]interface{}, 2)
	opts[0] = room
	opts[1] = event

	bcMessage := map[string][]interface{}{
		"opts": opts,
		"args": args,
	}
	bcMessageJSON, err := json.Marshal(bcMessage)
	if err != nil {
		return
	}

	_, err = bc.pub.Conn.Do("PUBLISH", bc.key, bcMessageJSON)
	if err != nil {
		return
	}
}

func (bc *redisBroadcast) sendAll(event string, args ...interface{}) {
	bc.lock.RLock()
	defer bc.lock.RUnlock()

	for _, connections := range bc.rooms {
		for _, connection := range connections {
			connection.Emit(event, args...)
		}
	}
}

func (bc *redisBroadcast) allRooms() []string {
	bc.lock.RLock()
	defer bc.lock.RUnlock()

	rooms := make([]string, 0, len(bc.rooms))
	for room := range bc.rooms {
		rooms = append(rooms, room)
	}

	return rooms
}

func (bc *redisBroadcast) getRoomsByConn(connection Conn) []string {
	var rooms []string

	for room, connections := range bc.rooms {
		if _, ok := connections[connection.ID()]; ok {
			rooms = append(rooms, room)
		}
	}

	return rooms
}

func (bc *redisBroadcast) dispatch() {
	for {
		switch m := bc.sub.Receive().(type) {
		case redis.Message:
			if m.Channel == bc.reqChannel {
				bc.onRequest(m.Data)
				break
			} else if m.Channel == bc.resChannel {
				bc.onResponse(m.Data)
				break
			}

			err := bc.onMessage(m.Channel, m.Data)
			if err != nil {
				return
			}

		case redis.Subscription:
			if m.Count == 0 {
				return
			}

		case error:
			return
		}
	}
}