global server state, individual client state, websocket, go

main.go

package main

import (
	"encoding/json"
	"fmt"
	"github.com/gorilla/mux"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
	"time"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

const (
	writeWait      = 10 * time.Second
	pongWait       = 60 * time.Second
	pingPeriod     = (pongWait * 9) / 10
	maxMessageSize = 512
)

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

func (c *connection) write(mt int, payload []byte) error {
	c.ws.SetWriteDeadline(time.Now().Add(writeWait))
	return c.ws.WriteMessage(mt, payload)
}

type Server struct {
	addConn   chan *connection
	delConn   chan *connection
	broadcast chan []byte

	manager *State
}

type State struct {
	sync.RWMutex
	blocks map[string]bool
}

func (s *State) AddKey(key string) {
	s.Lock()
	s.blocks[key] = true
	s.Unlock()
}

func (s *State) DelKey(key string) {
	s.Lock()
	delete(s.blocks, key)
	s.Unlock()
}

func (s *State) JSON() []byte {
	s.RLock()
	j, _ := json.Marshal(s.blocks)
	s.RUnlock()
	return j
}

func (s *Server) websocketRouter() {
	hub := make(map[*connection]bool)
	for {
		select {
		case c := <-s.addConn:
			hub[c] = true
		case c := <-s.delConn:
			log.Println("goodbye")
			delete(hub, c)
		case m := <-s.broadcast:
			for c, _ := range hub {
				c.send <- m
			}
		}
	}
}

func (s *Server) websocketReadPump(c *connection) {
	defer func() {
		s.delConn <- c
		c.ws.Close()
	}()
	c.ws.SetReadLimit(maxMessageSize)
	c.ws.SetReadDeadline(time.Now().Add(pongWait))
	c.ws.SetPongHandler(func(string) error { c.ws.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := c.ws.ReadMessage()
		if err != nil {
			break
		}

		// example state check & response
		if string(message) == "list" {
			c.send <- s.manager.JSON()
		}
	}
}

func (s *Server) websocketWritePump(c *connection) {
	ticker := time.NewTicker(pingPeriod)
	defer func() {
		ticker.Stop()
		c.ws.Close()
	}()
	for {
		select {
		case message, ok := <-c.send:
			if !ok {
				c.write(websocket.CloseMessage, []byte{})
				return
			}
			if err := c.write(websocket.TextMessage, message); err != nil {
				return
			}
		case <-ticker.C:
			if err := c.write(websocket.PingMessage, []byte{}); err != nil {
				return
			}
		}
	}
}

func (s *Server) websocketHandler(w http.ResponseWriter, r *http.Request) {
	ws, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	c := &connection{send: make(chan []byte, 256), ws: ws}

	s.addConn <- c
	go s.websocketWritePump(c)
	go s.websocketReadPump(c)
}

func (s *Server) addHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	key := vars["id"]

	s.manager.AddKey(key)
	s.broadcast <- []byte(fmt.Sprintf(`{"server":"added %s"}`, key))

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(`{"server":"ok"}`))
}

func (s *Server) delHandler(w http.ResponseWriter, r *http.Request) {
	vars := mux.Vars(r)
	key := vars["id"]

	s.manager.DelKey(key)
	s.broadcast <- []byte(fmt.Sprintf(`{"server":"deleted %s"}`, key))

	w.Header().Set("Content-Type", "application/json")
	w.Header().Set("Access-Control-Allow-Origin", "*")
	w.WriteHeader(http.StatusOK)
	w.Write([]byte(`{"server":"ok"}`))
}

func (s *Server) rootHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Write([]byte(`
    <style> p { margin: 0px; font-family: courier; font-size: 14px } </style>
    <div id="log"></div>
    <script>
    var ws = new WebSocket('ws://localhost:8080/ws');
    var log = document.getElementById('log');
    ws.onmessage = function(m){
        var child = document.createElement('p')
        child.innerHTML = m.data;
        log.insertBefore(child, log.firstChild);
    }
    var id = Math.floor(Math.random() * 100);
    ws.onopen = function(){
        ws.send('i am client ' + id);
        ws.send('list');
    }
    </script>
    `))
}

func main() {
	s := &Server{
		addConn:   make(chan *connection),
		delConn:   make(chan *connection),
		broadcast: make(chan []byte),
		manager: &State{
			blocks: make(map[string]bool),
		},
	}

	go s.websocketRouter()

	r := mux.NewRouter()
	r.HandleFunc("/ws", s.websocketHandler)
	r.HandleFunc("/add/{id}", s.addHandler)
	r.HandleFunc("/del/{id}", s.delHandler)
	r.HandleFunc("/", s.rootHandler)
	http.Handle("/", r)

	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		log.Fatal("ListenAndServe: ", err)
	}
}