Howdy folks. I was refactoring some of my Go code from another project this past weekend and I realized this would be a great opportunity for a post about working with WebSockets in Go, so here we are.
Background
⌗
First things first, the WebSocket implementation in the Go standard library is generally regarded as too low-level/error prone, and almost everyone uses a third-party library. One very widely used library is the
Gorilla WebSockets library
. The Gorilla library is especially great because it contains a few examples that cover a lot of concepts that are essential for beginners.
However, I found that even this was a little too verbose. I wanted to pull in some WebSocket code from one project into another and it was going to be an absolute headache for this one project, let alone for more projects I have planned. So I ended up refactoring it and wrote my own
github.com/brojonat/websocket
package that wraps the Gorilla WebSocket connection and exposes some convenient interfaces (and has tests so I don’t have to worry about some subtle race condition). Let’s get into it.
Interfaces
⌗
The main thing to internalize about working with WebSockets in Go is that each client connection should get at least two goroutines: one that continuously processes messages coming from the client (i.e., a “read pump”), and one that continuously processes message going out to the client (i.e., a “write pump”). There’s also some lower level protocol stuff you need to handle (i.e., handling ping/pong messages, closing the connection, etc.). We can encapsulate this in a simple
Client
interface:
// Client is an interface for reading from and writing to a websocket
// connection. It is designed to be used as a middleman between a service and a
// client websocket connection.
type Client interface {
io.Writer
io.Closer
// WritePump is responsible for writing message to the client (including
// ping message)
WritePump(func(Client), time.Duration)
// ReadPump is responsible for setting up the client connection, reading
// connections from the client, and passing them to the handlers
ReadPump(func(Client), ...func([]byte))
You’ll notice that this interface embeds the
io.Writer
and
io.Closer
interfaces. These both do the expected thing, which is nice; we can use an implementation of this interface like we would any
Writer
. What about all these extra parameters in
WritePump
and
ReadPump
though? Let’s walk through those:
In
WritePump
:
func(Client)
: The deregistration function. This is
defer
-ed. We’ll see an example below, but callers can implement their own callback to invoke after the connection is closed.
time.Duration
: The interval at which ping messages are sent.
In
ReadPump
:
func(Client)
: The deregistration function (see above).
...func([]byte)
: Zero or more message handlers that will be invoked with any messages received from the client. This is what gives this package flexibility: callers can implement whatever functionality they want here.
The package provides an example/convenience implementation available via
NewClient
that uses a channel under the hood, but consumers of this package are free to implement their own
Client
in whatever way they want (dang, interfaces really
are
useful).
Cool, so we have this
Client
interface, and we’re going to have one instance for every client connection. But in just about every scenario, we’re also going to want an abstraction for managing
all
of the connections. I smell another interface: the
Manager
. This is also a rather simple interface:
// Manager maintains the set of active WebSocket clients.
type Manager interface {
// Clients returns a slice of Clients
Clients() []Client
// RegisterClient adds the client to the Managers internal store
RegisterClient(Client)
// UnregisterClient removes the client from the Managers internal store
UnregisterClient(Client)
// Run runs in its own goroutine. It continuously processes client
// (un)registration. When a client registers, it is added to the manager's
// internal store. When a client unregisters, it is removed from the
// manager's internal store.
Run(context.Context)
As you may have guessed from the comments, this is simply a store of connections. The package provides an example implementation available via
NewManager
, but consumers are free to implement their own
Manager
however they want. Again, that’s sweet.
Putting it all together
⌗
We’re still missing the main entry point handling HTTP requests and upgrading it into a WebSocket connection. Well, the package provides a handler to do exactly that! Check it out:
// ServeWS upgrades HTTP connections to WebSocket, creates the pump, calls the
// registration callback, and starts goroutines that handle reading (writing)
// from (to) the client.
func ServeWS(
// upgrader upgrades the connection
upgrader websocket.Upgrader,
// connSetup is called on the upgraded WebSocket connection to configure
// the connection
connSetup func(*websocket.Conn),
// factory is a function that takes a connection and returns a Client
factory func(*websocket.Conn) Client,
// register is a function to call once the Client is created (e.g.,
// store it in a some collection on the service for later reference)
register func(Client),
// unregister is a function to call after the WebSocket connection is closed
// (e.g., remove it from the collection on the service)
unregister func(Client),
// ping is the interval at which ping messages are aren't sent
ping time.Duration,
// msgHandlers are callbacks that handle messages received from the client
msgHandlers []func([]byte),
) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
// if Upgrade fails it closes the connection, so just return
return
// set connection limits and ping/pong settings
connSetup(conn)
// create the interface
client := factory(conn)
// call the register callback
register(client)
// run writePump in a separate goroutine; all writes will happen here,
// ensuring only one write on the connection at a time
go client.WritePump(unregister, ping)
// run readPump in a separate goroutine; all reads will happen here,
// ensuring only one reader on the connection at a time
go client.ReadPump(unregister, msgHandlers...)
Ok, so this
ServeWS
returns an
http.HandlerFunc
that we can use like any other HTTP handler and it will do all the necessary stuff to turn the client’s vanilla HTTP connection into a real time bi-directional byte stream…sick. Most importantly, all the dependencies are made explicit in this
ServeWS
signature (even some that probably don’t need to be…who wants to futz with ping intervals anyway?), so anyone can use this package for whatever their use case entails.
Ok, I want to go use this thing so I’m gonna wrap it up here. If you want more details, check out the package itself
here
. In case the package isn’t public yet or you’re just too tired to click a link, I’ve added the package code (and test package code) below:
package websocket
import (
"context"
"net/http"
"sync"
"time"
"github.com/gorilla/websocket"
// DefaultSetupClient is an example implementation of a function that sets up a
// websocket connection.
func DefaultSetupConn(c *websocket.Conn) {
pw := 60 * time.Second
c.SetReadLimit(512)
c.SetReadDeadline(time.Now().Add(pw))
c.SetPongHandler(func(string) error {
c.SetReadDeadline(time.Now().Add(pw))
return nil
// Client is an interface for reading from and writing to a websocket
// connection. It is designed to be used as a middleman between a service and a
// client websocket connection.
type Client interface {
io.Writer
io.Closer
// WritePump is responsible for writing message to the client (including
// ping message)
WritePump(func(Client), time.Duration)
// ReadPump is responsible for setting up the client connection, reading
// connections from the client, and passing them to the handlers
ReadPump(func(Client), ...func([]byte))
// ServeWS upgrades HTTP connections to WebSocket, creates the pump, calls the
// registration callback, and starts goroutines that handle reading (writing)
// from (to) the client.
func ServeWS(
// upgrader upgrades the connection
upgrader websocket.Upgrader,
// connSetup is called on the upgraded WebSocket connection to configure
// the connection
connSetup func(*websocket.Conn),
// factory is a function that takes a connection and returns a Client
factory func(*websocket.Conn) Client,
// register is a function to call once the Client is created (e.g.,
// store it in a some collection on the service for later reference)
register func(Client),
// unregister is a function to call after the WebSocket connection is closed
// (e.g., remove it from the collection on the service)
unregister func(Client),
// ping is the interval at which ping messages are aren't sent
ping time.Duration,
// msgHandlers are callbacks that handle messages received from the client
msgHandlers []func([]byte),
) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
// if Upgrade fails it closes the connection, so just return
return
// set connection limits and ping/pong settings
connSetup(conn)
// create the interface
client := factory(conn)
// call the register callback
register(client)
// run writePump in a separate goroutine; all writes will happen here,
// ensuring only one write on the connection at a time
go client.WritePump(unregister, ping)
// run readPump in a separate goroutine; all reads will happen here,
// ensuring only one reader on the connection at a time
go client.ReadPump(unregister, msgHandlers...)
// This is a convenient example implementation of the Client interface.
type client struct {
// the underlying websocket connection
conn *websocket.Conn
// Egress is a buffered channel for writes to the client; callers can push
// messages into this channel and they'll be written to the client
egress chan []byte
// NewClient is a convenience function for returning a new Client
func NewClient(c *websocket.Conn) Client {
return &client{
conn: c,
egress: make(chan []byte, 32),
// Write implements the Writer interface
func (c *client) Write(p []byte) (int, error) {
c.egress <- p
return len(p), nil
// Close implements the Closer interface. Note the behavior of calling Close()
// multiple times is undefined, so we're just going to swallow all errors for
// now
func (c *client) Close() error {
c.conn.WriteControl(websocket.CloseMessage, []byte{}, time.Time{})
c.conn.Close()
return nil
// WritePump pumps messages from the egress channel (typically originating from
// the service) to the underlying websocket connection.
// A goroutine running WritePump is started for each connection. The application
// ensures that there is at most one writer to a connection by executing all
// writes from this goroutine.
func (c *client) WritePump(unregisterFunc func(Client), ping time.Duration) {
// Create a ticker that triggers a ping at given interval
pingTicker := time.NewTicker(ping)
defer func() {
pingTicker.Stop()
unregisterFunc(c)
for {
select {
case msgBytes, ok := <-c.egress:
// ok will be false in case the egress channel is closed
if !ok {
// manager has closed this connection channel, so send a close
// message and return which will initiate the connection
// shutdown process in the manager
c.conn.WriteMessage(websocket.CloseMessage, nil)
return
// write a message to the connection
if err := c.conn.WriteMessage(websocket.TextMessage, msgBytes); err != nil {
// just return to (closes the connection) indiscriminantly in
// the case of errors
return
case <-pingTicker.C:
if err := c.conn.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
// just return to (closes the connection) indiscriminantly in
// the case of errors
return
// ReadPump pumps messages from the websocket connection to the service.
// The application runs ReadPump in a per-connection goroutine. The application
// ensures that there is at most one reader on a connection by executing all
// reads from this goroutine.
func (c *client) ReadPump(
unregisterFunc func(Client),
handlers ...func([]byte),
// unregister and close before exit
defer func() {
unregisterFunc(c)
// read forever
for {
_, payload, err := c.conn.ReadMessage()
// handle close (we could check the error here but it doesn't really
// matter, something is botched, so just close the connection)
if err != nil {
break
// do things with the message
for _, h := range handlers {
h(payload)
// Manager maintains the set of active WebSocket clients.
type Manager interface {
// Clients returns a slice of Clients
Clients() []Client
// RegisterClient adds the client to the Managers internal store
RegisterClient(Client)
// UnregisterClient removes the client from the Managers internal store
UnregisterClient(Client)
// Run runs in its own goroutine. It continuously processes client
// (un)registration. When a client registers, it is added to the manager's
// internal store. When a client unregisters, it is removed from the
// manager's internal store.
Run(context.Context)
type manager struct {
lock sync.RWMutex
clients map[Client]struct{}
register chan regreq
unregister chan regreq
// Helper struct for signaling registration/unregistration. The Run() goroutine
// can signal the operation is done by sending on the done chan.
type regreq struct {
wp Client
done chan struct{}
func NewManager() Manager {
return &manager{
lock: sync.RWMutex{},
clients: make(map[Client]struct{}),
register: make(chan regreq),
unregister: make(chan regreq),
func (m *manager) Clients() []Client {
res := []Client{}
m.lock.RLock()
defer m.lock.RUnlock()
for c, _ := range m.clients {
res = append(res, c)
return res
func (m *manager) RegisterClient(wp Client) {
done := make(chan struct{})
rr := regreq{
wp: wp,
done: done,
m.register <- rr
<-done
func (m *manager) UnregisterClient(wp Client) {
done := make(chan struct{})
rr := regreq{
wp: wp,
done: done,
m.unregister <- rr
<-done
func (m *manager) Run(ctx context.Context) {
// helper fn for cleaning up client
cleanupClient := func(c Client) {
// delete from map
delete(m.clients, c)
// close connections
c.Close()
// run forever registering, unregistering, and listening for cleanup
for {
select {
// register new client
case rr := <-m.register:
m.lock.Lock()
m.clients[rr.wp] = struct{}{}
m.lock.Unlock()
rr.done <- struct{}{}
// cleanup single client
case rr := <-m.unregister:
m.lock.Lock()
if _, ok := m.clients[rr.wp]; ok {
cleanupClient(rr.wp)
m.lock.Unlock()
rr.done <- struct{}{}
// handle service shutdown
case <-ctx.Done():
m.lock.Lock()
for client := range m.clients {
cleanupClient(client)
m.lock.Unlock()
// Broadcaster is an example implementation of Manager that has a
// Broadcast method that writes the supplied message to all clients.
type Broadcaster struct {
*manager
func NewBroadcaster() Manager {
m := manager{
lock: sync.RWMutex{},
clients: make(map[Client]struct{}),
register: make(chan regreq),
unregister: make(chan regreq),
return &Broadcaster{
manager: &m,
func (bb *Broadcaster) Broadcast(b []byte) {
bb.lock.RLock()
defer bb.lock.RUnlock()
for w := range bb.clients {
w.Write(b)
Tests, as promised:
package websocket_test
import (
"context"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/gorilla/websocket"
"github.com/matryer/is"
ws "github.com/brojonat/websocket"
func TestWSHandler(t *testing.T) {
is := is.New(t)
ctx := context.Background()
var p ws.Client
reg := make(chan ws.Client)
dereg := make(chan ws.Client)
manager := ws.NewManager()
upgrader := websocket.Upgrader{ReadBufferSize: 1024, WriteBufferSize: 1024}
upgrader.CheckOrigin = func(r *http.Request) bool { return true }
testBytes := []byte("testing")
// run the manager
go manager.Run(ctx)
// setup the handler
h := ws.ServeWS(
upgrader,
ws.DefaultSetupConn,
ws.NewClient,
func(_p ws.Client) {
p = _p
manager.RegisterClient(p)
reg <- p
func(_p ws.Client) {
manager.UnregisterClient(_p)
dereg <- _p
50*time.Second,
[]func([]byte){func(b []byte) { p.Write(b) }},
// create test server
s := httptest.NewServer(h)
defer s.Close()
// connect to the server
client, _, err := websocket.DefaultDialer.Dial(
"ws"+strings.TrimPrefix(s.URL, "http"), nil)
is.NoErr(err)
defer client.Close()
// the manager should have one registered client
<-reg
is.Equal(len(manager.Clients()), 1)
// write a message to the server; this will be echoed back
err = client.WriteMessage(websocket.TextMessage, testBytes)
is.NoErr(err)
// server should have echoed the message back
_, msg, err := client.ReadMessage()
is.NoErr(err)
is.Equal(msg, testBytes)
// close the connection, this should trigger the server/handler to
// cleanup and unregister the client connection
client.WriteControl(websocket.CloseMessage, []byte{}, time.Now())
client.Close()
// block until the unregistration loop has finished, then assert
// that the deregistration worked as expected
_p := <-dereg
is.Equal(len(manager.Clients()), 0)
is.Equal(_p, p)
