Blog
WebSocket Server in Go

WebSocket Server in Go

October 5, 2023

We will create a simple WebSocket server with a tiny WebSocket library for Go gobwas/ws. Learn how to handle concurrent client connections.

As a base project, we will use the example app from First App with Go.

Create component for WebSocket handler

Make directory for new component called ws. This component will implement an HTTP handler, active connections storage, and queue to send messages.

mkdir -p internal/ws

WebSocket handler

Install WebSocket library to handle HTTP requests:

go get github.com/gobwas/ws

Create new file internal/ws/handler.go with HTTP handler:

package ws

import (
	"io"
	"log"
	"net"
	"net/http"

	"github.com/gobwas/ws"
	"github.com/gobwas/ws/wsutil"
)

func wsLoop(conn net.Conn) {
	client := newClient(conn)
	defer client.close()

	r := wsutil.NewReader(conn, ws.StateServerSide)

	for {
		// Wait for next frame from client
		hdr, err := r.NextFrame()
		if err != nil {
			if err != io.EOF {
				log.Panicln("failed to read frame", err)
			}
			return
		}

		if hdr.OpCode == ws.OpClose {
			return
		}

		// Drop any data
		_, _ = io.Copy(io.Discard, r)
	}
}

func Handler(w http.ResponseWriter, r *http.Request) {
	conn, _, _, err := ws.UpgradeHTTP(r, w)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	go wsLoop(conn)
}

The ws.UpgradeHTTP() function detaches the connection from the HTTP handler and starts the wsLoop() function in the new goroutine. The wsLoop() function appends new client to the storage and waits for packets from the client. In this example, we discarded all incoming packets.

Client storage

Create new file internal/ws/client.go with client definition:

package ws

import (
	"log"
	"net"
	"sync"

	"github.com/gobwas/ws"
	"github.com/gobwas/ws/wsutil"
)

type wsClient struct {
	conn net.Conn
}

var (
	lock    sync.RWMutex
	clients = make(map[*wsClient]struct{})
)

func newClient(conn net.Conn) *wsClient {
	c := &wsClient{
		conn: conn,
	}

	lock.Lock()
	clients[c] = struct{}{}
	lock.Unlock()

	return c
}

func (c *wsClient) close() {
	c.conn.Close()

	lock.Lock()
	delete(clients, c)
	lock.Unlock()
}

func (c *wsClient) send(data []byte) error {
	w := wsutil.NewWriter(c.conn, ws.StateServerSide, ws.OpText)

	if _, err := w.Write(data); err != nil {
		return err
	}

	if err := w.Flush(); err != nil {
		return err
	}

	return nil
}

func broadcast(data []byte) {
	lock.RLock()
	defer lock.RUnlock()

	for client := range clients {
		if err := client.send(data); err != nil {
			log.Println("send to client failed", err)
			continue
		}

	}
}

For client storage, we use map. It’s a nice trick. Appending and removing new objects with map is very simple and effective.

Message Sending

Now we create new public function to send messages to all clients. Create new file internal/ws/send.go:

package ws

import (
	"context"
	"log"
)

var queue = make(chan []byte, 1000)

// Send pushes data to the queue
func Send(data []byte) {
	select {
	case queue <- data:
	default:
		log.Println("send failed: queue is full")
	}
}

// SendLoop is an infinity loop, send messages from queue to WebSocket clients
func SendLoop(ctx context.Context) {
	for {
		select {
		case <-ctx.Done():
			return
		case msg := <-queue:
			broadcast(msg)
		}
	}
}

HTTP Server

You may append WebSocket handler to your HTTP router. For this example, we will create a basic HTTP server. Open file internal/app/app.go in any text editor:

package app

import (
	"context"
	_ "embed"
	"net/http"
	"time"

	"example/internal/ws"
)

//go:embed index.html
var indexHTML []byte

// webSPA serves index.html
func webSPA(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	_, _ = w.Write(indexHTML)
}

// exampleBroadcast sends new messages to WebSocket clients each second
func exampleBroadcast(ctx context.Context) {
	tick := time.NewTicker(time.Second)
	defer tick.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case v := <-tick.C:
			currentTime := v.Format("2006-01-02 15:04:05")
			ws.Send([]byte(currentTime))
		}
	}
}

func Start() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go ws.SendLoop(ctx)
	go exampleBroadcast(ctx)

	mux := http.NewServeMux()
	mux.HandleFunc("/ws/", ws.Handler)
	mux.HandleFunc("/", webSPA)

	return http.ListenAndServe(":8080", mux)
}

And finally, create simple HTML file internal/app/index.html to start WebSocket client:

<!doctype html>
<html>
<head>
	<meta charset="UTF-8">
	<title>WebSocket Client</title>
</head>
<body>
	<h1>WebSocket Client</h1>
	<pre id="output"></pre>
	<script>
		const url = new URL(window.location)
		url.pathname = '/ws/'
		url.protocol = url.protocol.replace('http', 'ws')
		const socket = new WebSocket(url)

		socket.addEventListener('message', function (event) {
			const codeElement = document.createElement('code')
			codeElement.textContent = event.data + '\n'
			document.getElementById('output').appendChild(codeElement)
		})
	</script>
</body>
</html>

Done. WebSocket server ready to launch

Launch

Start your application:

go run ./cmd/example

And open in your browser http://localhost:8080

Check out this project in our repository.

Effective Code Style in Golocalhost vs 127.0.0.1