Chat

chat.go

package main

import (
	"flag"
	"html/template"
	"log"
	"net/http"
	"path/filepath"
	"sync"

	"github.com/gorilla/websocket"
)

const (
	connectionBufferSize = 1024
	messageBufferSize    = 256
)

// Turns on debugging for the application; especially useful during
// development:
//
// 1. Templates are compiled per request instead of per process instance
//    allowing reloading the templates as you make changes.
var debug = flag.Bool("debug", false, "turns on debugging")

// The host and port to which the HTTP server will bind.
var listenAddr = flag.String("addr", ":8080", "listen on <host:port>")

// compileTemplate compiles a template given its filename.
func compileHTMLTemplate(filename string) *template.Template {
	return template.Must(template.ParseFiles(filepath.Join("templates", filename)))
}


// client represents a single chatting connection.
type client struct {
	conn *websocket.Conn
	send chan []byte
	room *room
}

// Blocking read from the conn.
func (c *client) read() {
	for {
		if _, msg, err := c.conn.ReadMessage(); err != nil {
			break
		} else {
			c.room.forward <- msg
		}
	}
	c.conn.Close()
}

// Blocking write to the conn.
func (c *client) write() {
	for msg := range c.send {
		if err := c.conn.WriteMessage(websocket.TextMessage, msg); err != nil {
			break
		}
	}
	c.conn.Close()
}


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

type room struct {
	// forward is a channel that holds incoming messages that should be
	// forwarded to other clients.
	forward chan []byte
	join    chan *client
	leave   chan *client
	clients map[*client]bool
}

// Creates an new instance of a room.
func newRoom() *room {
	return &room{
		forward: make(chan []byte),
		join:    make(chan *client),
		leave:   make(chan *client),
		clients: make(map[*client]bool),
	}
}

// Kick starts a room and blocks for a conversation.
func (r *room) run() {
	for {
		select {
		case client := <-r.join:
			r.clients[client] = true
		case client := <-r.leave:
			delete(r.clients, client)
			close(client.send)
		case msg := <-r.forward:
			for client := range r.clients {
				select {
				case client.send <- msg:
					// send the message
				default:
					// failed to send
					delete(r.clients, client)
					close(client.send)
				}
			}
		}
	}
}

// The room is an HTTP handler.
func (r *room) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	conn, err := upgrader.Upgrade(w, req, nil)
	if err != nil {
		log.Fatal("ServeHTTP:", err)
		return
	}
	client := &client{
		conn: conn,
		send: make(chan []byte, messageBufferSize),
		room: r,
	}
	r.join <- client
	defer func() { r.leave <- client }() // If the connection closes, clean up.
	go client.write()                    // Keep writing to the client in a goroutine.
	client.read()                        // Block reading.
}


// A templateHandler executes a template to generate an HTTP response.
type templateHandler struct {
	// Enabling debug mode causes the template to be recompiled
	// per request (useful during development).
	debug bool
	once  sync.Once

	filename string
	templ    *template.Template
}

// ServeHTTP implements the http.Handler interface.
func (t *templateHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if t.debug {
		t.templ = compileHTMLTemplate(t.filename)
	} else {
		t.once.Do(func() {
			t.templ = compileHTMLTemplate(t.filename)
		})
	}
	t.templ.Execute(w, nil)
}

// Entry-point function.
func main() {
	flag.Parse()

	r := newRoom()

	http.Handle("/", &templateHandler{filename: "chat.html", debug: *debug})
	http.Handle("/room", r)

	go r.run()

	log.Println("Started HTTP server on: ", *listenAddr)
	if err := http.ListenAndServe(*listenAddr, nil); err != nil {
		log.Fatal("ListenAndServe:", err)
	}
}