cloud-game/ws.go

package main

import (
	"encoding/json"
	"log"
	"sync"
	"time"

	"github.com/gorilla/websocket"
	uuid "github.com/satori/go.uuid"
)

type Client struct {
	conn *websocket.Conn

	sendLock sync.Mutex
	// sendCallback is callback based on packetID
	sendCallback     map[string]func(req WSPacket)
	sendCallbackLock sync.Mutex
	// recvCallback is callback when receive based on ID of the packet
	recvCallback map[string]func(req WSPacket)
}

type WSPacket struct {
	ID   string `json:"id"`
	Data string `json:"data"`

	RoomID      string `json:"room_id"`
	PlayerIndex int    `json:"player_index"`

	TargetHostID string `json:"target_id"`
	PacketID     string `json:"packet_id"`
	// Globally ID of a session
	SessionID string `json:"session_id"`
}

var EmptyPacket = WSPacket{}

func NewClient(conn *websocket.Conn) *Client {
	sendCallback := map[string]func(WSPacket){}
	recvCallback := map[string]func(WSPacket){}
	return &Client{
		conn: conn,

		sendCallback: sendCallback,
		recvCallback: recvCallback,
	}
}

// send sends a packet and trigger callback when the packet comes back
func (c *Client) send(request WSPacket, callback func(response WSPacket)) {
	request.PacketID = uuid.Must(uuid.NewV4()).String()
	data, err := json.Marshal(request)
	if err != nil {
		return
	}

	// TODO: Consider using lock free
	c.sendLock.Lock()
	c.conn.WriteMessage(websocket.TextMessage, data)
	c.sendLock.Unlock()
	wrapperCallback := func(resp WSPacket) {
		resp.PacketID = request.PacketID
		resp.SessionID = request.SessionID
		callback(resp)
	}
	if callback == nil {
		return
	}
	c.sendCallbackLock.Lock()
	c.sendCallback[request.PacketID] = wrapperCallback
	c.sendCallbackLock.Unlock()
}

// receive receive and response back
func (c *Client) receive(id string, f func(response WSPacket) (request WSPacket)) {
	c.recvCallback[id] = func(response WSPacket) {
		req := f(response)
		// Add Meta data
		req.PacketID = response.PacketID
		req.SessionID = response.SessionID

		// Skip rqeuest if it is EmptyPacket
		if req == EmptyPacket {
			return
		}
		resp, err := json.Marshal(req)
		if err != nil {
			log.Println("[!] json marshal error:", err)
		}
		//c.conn.SetWriteDeadline(time.Now().Add(writeWait))
		c.sendLock.Lock()
		c.conn.WriteMessage(websocket.TextMessage, resp)
		c.sendLock.Unlock()
	}
}

// syncSend sends a packet and wait for callback till the packet comes back
func (c *Client) syncSend(request WSPacket) (response WSPacket) {
	res := make(chan WSPacket)
	f := func(resp WSPacket) {
		res <- resp
	}
	c.send(request, f)
	return <-res
}

// heartbeat maintains connection to server
func (c *Client) heartbeat() {
	// send heartbeat every 1s
	timer := time.Tick(time.Second)

	for range timer {
		c.send(WSPacket{ID: "heartbeat"}, nil)
	}
}

func (c *Client) listen() {
	for {
		// log.Println("Waiting for message ...")
		_, rawMsg, err := c.conn.ReadMessage()
		if err != nil {
			log.Println("[!] read:", err)
			break
		}
		wspacket := WSPacket{}
		err = json.Unmarshal(rawMsg, &wspacket)
		if err != nil {
			continue
		}

		// Check if some async send is waiting for the response based on packetID
		// TODO: Change to read lock
		c.sendCallbackLock.Lock()
		callback, ok := c.sendCallback[wspacket.PacketID]
		c.sendCallbackLock.Unlock()
		if ok {
			go callback(wspacket)
			c.sendCallbackLock.Lock()
			delete(c.sendCallback, wspacket.PacketID)
			c.sendCallbackLock.Unlock()
			// Skip receiveCallback to avoid duplication
			continue
		}
		// Check if some receiver with the ID is registered
		if callback, ok := c.recvCallback[wspacket.ID]; ok {
			go callback(wspacket)
		}
	}
}