Implementing Chat with WebSockets and Go

This example application shows how to use WebSockets, the Go programming language and jQuery to create a simple web chat application.

The source for this example is available as a gist.

Running the example

The example requires a working Go development environment. The Getting Started page describes how to install the development environment.

Once you have Go up and running, you can download, build and run the example using the commands:

go get gary.burd.info/go-websocket-chat
go-websocket-chat

Open http://127.0.0.1:8080/ in a websocket capable browser to try the application.

Server

The server application is implemented using the http package included with the Go distribution and the Gorilla Project's websocket package.

The application defines two types, connection and hub. The server creates an instance of the connection type for each webscocket connection. Connections act as an intermediary between the websocket and a single instance of the hub type. The hub maintains a set of registered connections and broadcasts messages to the connections.

The application runs one goroutine for the hub and two goroutines for each connection. The goroutines communicate with each other using channels. The hub has channels for registering connections, unregistering connections and broadcasting messages. A connection has a buffered channel of outbound messages. One of the connection's goroutines reads messages from this channel and writes the messages to the webscoket. The other connection goroutine reads messages from the websocket and sends them to the hub.

Here's the code for the hub type. A description of the code follows.

package main

type hub struct {
    // Registered connections.
    connections map[*connection]bool

    // Inbound messages from the connections.
    broadcast chan []byte

    // Register requests from the connections.
    register chan *connection

    // Unregister requests from connections.
    unregister chan *connection
}

var h = hub{
    broadcast:   make(chan []byte),
    register:    make(chan *connection),
    unregister:  make(chan *connection),
    connections: make(map[*connection]bool),
}

func (h *hub) run() {
    for {
        select {
        case c := <-h.register:
            h.connections[c] = true
        case c := <-h.unregister:
            delete(h.connections, c)
            close(c.send)
        case m := <-h.broadcast:
            for c := range h.connections {
                select {
                case c.send <- m:
                default:
                    delete(h.connections, c)
                    close(c.send)
                    go c.ws.Close()
                }
            }
        }
    }
}

The application's main function starts the hub run method as a goroutine. Connections send requests to the hub using the register, unregister and broadcast channels.

The hub registers connections by adding the connection pointer as a key in the connections map. The map value is always true.

The unregister code is a little more complicated. In addition to deleting the connection pointer from the connections map, the hub closes the connection's send channel to signal the connection that no more messages will be sent to the connection.

The hub handles messages by looping over the registered connections and sending the message to the connection's send channel. If the connection's send buffer is full, then the hub assumes that the client is dead or stuck. In this case, the hub unregisters the connection and closes the websocket.

Here's the code related to the connection type.

package main

import (
    "github.com/gorilla/websocket"
    "net/http"
)

type connection struct {
    // The websocket connection.
    ws *websocket.Conn

    // Buffered channel of outbound messages.
    send chan []byte
}

func (c *connection) reader() {
    for {
        _, message, err := c.ws.ReadMessage()
        if err != nil {
            break
        }
        h.broadcast <- message
    }
    c.ws.Close()
}

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

func wsHandler(w http.ResponseWriter, r *http.Request) {
    ws, err := websocket.Upgrade(w, r, nil, 1024, 1024)
    if _, ok := err.(websocket.HandshakeError); ok {
        http.Error(w, "Not a websocket handshake", 400)
        return
    } else if err != nil {
        return
    }
    c := &connection{send: make(chan []byte, 256), ws: ws}
    h.register <- c
    defer func() { h.unregister <- c }()
    go c.writer()
    c.reader()
}

The wsHandler function is registered by the application's main function as a http handler. The upgrades the HTTP connection to the WebSocket protocol, creates a connection object, registers the connection with the hub and schedules the connection to be unregistered using a defer statement.

Next, the wsHandler function starts the connection's writer method as a goroutine. The writer method transfers messages from the connection's send channel to the websocket. The writer method exits when the channel is closed by the hub or there's an error writing to the websocket.

Finally, the wsHandler function calls the connection's reader method. The reader method transfers inbound messages from the websocket to the hub.

Here's the remainder of the server code.

package main

import (
    "flag"
    "go/build"
    "log"
    "net/http"
    "path/filepath"
    "text/template"
)

var (
    addr      = flag.String("addr", ":8080", "http service address")
    assets    = flag.String("assets", defaultAssetPath(), "path to assets")
    homeTempl *template.Template
)

func defaultAssetPath() string {
    p, err := build.Default.Import("gary.burd.info/go-websocket-chat", "", build.FindOnly)
    if err != nil {
        return "."
    }
    return p.Dir
}

func homeHandler(c http.ResponseWriter, req *http.Request) {
    homeTempl.Execute(c, req.Host)
}

func main() {
    flag.Parse()
    homeTempl = template.Must(template.ParseFiles(filepath.Join(*assets, "home.html")))
    go h.run()
    http.HandleFunc("/", homeHandler)
    http.HandleFunc("/ws", wsHandler)
    if err := http.ListenAndServe(*addr, nil); err != nil {
        log.Fatal("ListenAndServe:", err)
    }
}

The application's main function starts the hub goroutine. Next, the main function registers handlers for the home page and websocket connections. Finally, the main function starts the HTTP server.

Client

The client is implemented in a single HTML file.

<html>
<head>
<title>Chat Example</title>
<script type="text/javascript" src="http://ajax.googleapis.com/ajax/libs/jquery/1.4.2/jquery.min.js"></script>
<script type="text/javascript">
    $(function() {

    var conn;
    var msg = $("#msg");
    var log = $("#log");

    function appendLog(msg) {
        var d = log[0]
        var doScroll = d.scrollTop == d.scrollHeight - d.clientHeight;
        msg.appendTo(log)
        if (doScroll) {
            d.scrollTop = d.scrollHeight - d.clientHeight;
        }
    }

    $("#form").submit(function() {
        if (!conn) {
            return false;
        }
        if (!msg.val()) {
            return false;
        }
        conn.send(msg.val());
        msg.val("");
        return false
    });

    if (window["WebSocket"]) {
        conn = new WebSocket("ws://{{$}}/ws");
        conn.onclose = function(evt) {
            appendLog($("<div><b>Connection closed.</b></div>"))
        }
        conn.onmessage = function(evt) {
            appendLog($("<div/>").text(evt.data))
        }
    } else {
        appendLog($("<div><b>Your browser does not support WebSockets.</b></div>"))
    }
    });
</script>
<style type="text/css">
html {
    overflow: hidden;
}

body {
    overflow: hidden;
    padding: 0;
    margin: 0;
    width: 100%;
    height: 100%;
    background: gray;
}

#log {
    background: white;
    margin: 0;
    padding: 0.5em 0.5em 0.5em 0.5em;
    position: absolute;
    top: 0.5em;
    left: 0.5em;
    right: 0.5em;
    bottom: 3em;
    overflow: auto;
}

#form {
    padding: 0 0.5em 0 0.5em;
    margin: 0;
    position: absolute;
    bottom: 1em;
    left: 0px;
    width: 100%;
    overflow: hidden;
}

</style>
</head>
<body>
<div id="log"></div>
<form id="form">
    <input type="submit" value="Send" />
    <input type="text" id="msg" size="64"/>
</form>
</body>
</html>

The client uses jQuery to manipulate objects in the browser.

On document load, the script checks for websocket functionality in the browser. If websocket functionality is available, then the script opens a connection to the server and registers a callback to handle messages from the server. The callback appends the message to the chat log using the appendLog function.

To allow the user to manually scroll through the chat log without interruption from new messages, the appendLog function checks the scroll position before adding new content. If the chat log is scrolled to the bottom, then the function scrolls new content into view after adding the content. Otherwise, the scroll position is not changed.

The form handler writes the user input to the websocket and clears the input field.

Feedback

Send a message to gary@beagledreams.com if you have questions or comments about this example. You should follow me on Twitter.