mtsatellite/cmd/mtwebmapper/forwardupdates.go

// Copyright 2014 by Sascha L. Teichmann
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package main

import (
	"bytes"
	"encoding/json"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

type websocketForwarder struct {
	upgrader    *websocket.Upgrader
	register    chan *connection
	unregister  chan *connection
	broadcast   chan msg
	connections map[*connection]bool
}

type connection struct {
	ws   *websocket.Conn
	send chan []byte
}

type msg struct {
	tiles map[xz]bool
	pls   []*player
}

func newWebsocketForwarder() *websocketForwarder {
	upgrader := &websocket.Upgrader{ReadBufferSize: 512, WriteBufferSize: 2048}
	return &websocketForwarder{
		upgrader:    upgrader,
		register:    make(chan *connection),
		unregister:  make(chan *connection),
		broadcast:   make(chan msg),
		connections: make(map[*connection]bool)}
}

func (wsf *websocketForwarder) run() {
	for {
		select {
		case c := <-wsf.register:
			wsf.connections[c] = true
		case c := <-wsf.unregister:
			if _, ok := wsf.connections[c]; ok {
				delete(wsf.connections, c)
				close(c.send)
			}
		case message := <-wsf.broadcast:
			if len(wsf.connections) == 0 {
				continue
			}
			encMsg := map[string]interface{}{}

			if tiles := message.tiles; tiles != nil {
				xzs := make([]xz, 0, len(tiles))
				for xz := range tiles {
					xzs = append(xzs, xz)
				}
				encMsg["tiles"] = xzs
			}

			if message.pls != nil {
				encMsg["players"] = message.pls
			}

			var buf bytes.Buffer
			encoder := json.NewEncoder(&buf)
			if err := encoder.Encode(encMsg); err != nil {
				log.Printf("encoding changes failed: %s\n", err)
				continue
			}
			m := buf.Bytes()
			for c := range wsf.connections {
				select {
				case c.send <- m:
				default:
					delete(wsf.connections, c)
					close(c.send)
				}
			}
		}
	}
}

func (wsf *websocketForwarder) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
	ws, err := wsf.upgrader.Upgrade(rw, r, nil)
	if err != nil {
		log.Printf("Cannot upgrade to websocket: %s\n", err)
		return
	}
	c := &connection{ws: ws, send: make(chan []byte, 8)}
	wsf.register <- c
	defer func() { wsf.unregister <- c }()
	go c.writer()
	c.reader()
}

func (wsf *websocketForwarder) BaseTilesUpdated(changes map[xz]bool) {
	wsf.broadcast <- msg{tiles: changes}
}

func (wsf *websocketForwarder) BroadcastPlayers(pls []*player) {
	wsf.broadcast <- msg{pls: pls}
}

func (c *connection) writer() {
	defer c.ws.Close()
	for msg := range c.send {
		if c.ws.WriteMessage(websocket.TextMessage, msg) != nil {
			break
		}
	}
}

func (c *connection) reader() {
	defer c.ws.Close()
	for {
		// Just read the message and ignore it.
		if _, _, err := c.ws.NextReader(); err != nil {
			break
		}
	}
}