Cancelable web socket in Go
This is my learning attempt to create a websocket client/server model with cancel support. I have used the following resources to learn and create this project:
The code is hosted at:
Concept
A client can use a connection to connect to a hub. The hub takes the incoming request and upgrades it to a connection and registers the new connection.
Every time the client sends a message the hub reads from the connection and streams a request. A request allows a service to decode the message as json and to reply to the sender or to broadcast to all connected clients.
- Connection - Shared connection code between client and server hub
- Hub - Manages all incoming connections and streams requests
- Request - A server side socket message with reply/broadcast features
Connection
The most important pieces are the connection which is used on the client as also on the hub. The connection blocks when reading a message in a for loop with a select. The blocking happens in the default case and the select case ends when the context is done.
// Read blocks and reads a message from the socket.
func (c *Connection) Read() ([]byte, error) {
for {
select {
case <-c.ctx.Done():
return nil, c.ctx.Err()
default:
_, message, err := c.socket.ReadMessage()
if err != nil {
log.Debugf("conn: read error %s\n", err)
c.closeSocket(err)
return nil, err
}
log.Debugf("conn: read %s\n", message)
return message, nil
}
}
}
Hub
The hub creates a new connection for each incoming client. It sends a request for each incoming socket message as request stream.
Most of this happens in the ServeHTTP function of the hub.
func (h *Hub) ServeHTTP(w http.ResponseWriter, r *http.Request) {
socket, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Infoln("upgrade:", err)
return
}
log.Infof("new connection: %s\n", socket.RemoteAddr())
conn := NewConnection(h.ctx, socket)
go func() {
// unregister connection when it is done
<-conn.Done()
log.Infof("unregister conn: %s\n", conn.Id())
h.unregister <- conn
}()
go func() {
// read requests from connection and stream them to the hub
defer func() {
h.unregister <- conn
}()
for {
select {
case <-h.ctx.Done():
return
default:
data, err := conn.Read()
log.Debugf("read: %s\n", data)
if err != nil {
log.Debugf("read error: %s\n", err)
return
}
req := &Request{
data: data,
connection: conn,
err: err,
}
h.requests <- req
}
}
}()
h.register <- conn
}
My first attempt was to place the unregister functionality inside the connection s also the request generation. But using closures on the hub makes the code much easier. As the closures encode behaviour and not provide an API surface.
The request allows the user to retrieve the raw data or decode it as json. Also the request allows the user to reply raw data and json or to broadcast to all connected client.
package ws
import "encoding/json"
// Request represents a request from a client.
type Request struct {
data []byte
connection *Connection
hub *Hub
err error
}
// Error returns the error that occurred while processing the request.
func (r Request) Error() error {
return r.err
}
// Reply writes a response to the client.
func (r Request) Reply(data []byte) {
r.connection.Write(data)
}
// ReplyJSON writes a JSON response to the client.
func (r Request) ReplyJSON(data interface{}) error {
return r.connection.WriteJSON(data)
}
// AsData returns the request data as a string.
func (r Request) AsData() []byte {
return r.data
}
// AsJSON decodes the request data as JSON.
func (r Request) AsJSON(v interface{}) error {
return json.Unmarshal(r.data, v)
}
// Broadcast sends data to all connections
func (r Request) Broadcast(data []byte) {
r.hub.Broadcast(data)
}
// BroadcastJSON encodes v as JSON and broadcasts it to all connections.
func (r Request) BroadcastJSON(v interface{}) error {
return r.hub.BroadcastJSON(v)
}
Server and Request Stream
Requests are only available from the hub as RequestStream. So the server can range over the stream and handle each message.
package main
import (
"context"
"flag"
"fmt"
"net/http"
"wswc/log"
"wswc/model"
"wswc/ws"
)
var addr = flag.String("addr", ":8080", "http service address")
func main() {
flag.Parse()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
h := ws.NewHub(ctx)
go func() {
for r := range h.Requests() {
var msg model.Message
err := r.AsJSON(&msg)
if err != nil {
log.Warnf("unmarshal message failed: %v", err)
continue
}
log.Debugf("server received: %d", msg.Count)
msg.Count++
log.Debugf("server send: %d", msg.Count)
r.ReplyJSON(&msg)
}
}()
go h.Run()
http.HandleFunc("/", h.ServeHTTP)
fmt.Printf("listen at %s\n", *addr)
err := http.ListenAndServe(*addr, nil)
if err != nil {
panic(err)
}
}
Client and cancel on exit condition
The client exits when the context is canceled. The canceling can happen either on an interrupt (e.g. Ctrl+C), a timeout, a severe error or when the exit condition, in this case max count has been reached.
There is no wait group in the client, the whole program just waits for <-ctx.Done(). I find this really elegant.
package main
import (
"context"
"flag"
"os"
"os/signal"
"syscall"
"time"
"wswc/log"
"wswc/model"
"wswc/ws"
)
const (
maxCount = 100000
timeout = 10 * time.Second
)
var addr = flag.String("addr", "ws://127.0.0.1:8080", "ws service address")
func main() {
flag.Parse()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
client, err := ws.Dial(ctx, *addr)
if err != nil {
cancel()
panic(err)
}
// handle incoming messages
go func() {
defer cancel()
for {
select {
case <-ctx.Done():
return
default:
log.Debugln("waiting for message")
var msg model.Message
err := client.ReadJSON(&msg)
if err != nil {
log.Warnf("read message failed: %v", err)
return
}
log.Debugf("count: %d\n", msg.Count)
if msg.Count >= maxCount {
log.Infof("client received %d messages, exit\n", maxCount)
return
}
msg.Count++
client.WriteJSON(&msg)
if err != nil {
log.Errorf("write message failed: %v", err)
return
}
}
}
}()
// send initial message
msg := model.Message{Count: 0}
client.WriteJSON(msg)
go func() {
sigs := make(chan os.Signal, 1)
signal.Notify(sigs, os.Interrupt, syscall.SIGTERM)
<-sigs
cancel()
}()
time.AfterFunc(timeout, func() {
log.Infof("timeout, ...")
cancel()
})
<-ctx.Done()
log.Infoln("client exit")
}