commit fb91a400328be9f48299451bdf35a5451bebb72c Author: Bartal Laearsson Date: Fri Jul 24 21:38:32 2026 +0100 initial commit diff --git a/README.md b/README.md new file mode 100644 index 0000000..4d35458 --- /dev/null +++ b/README.md @@ -0,0 +1,35 @@ +# Chat Example + +This directory contains a full stack example of a simple chat webapp using github.com/coder/websocket. + +```bash +$ cd examples/chat +$ go run . localhost:0 +listening on ws://127.0.0.1:51055 +``` + +Visit the printed URL to submit and view broadcasted messages in a browser. + +![Image of Example](https://i.imgur.com/VwJl9Bh.png) + +## Structure + +The frontend is contained in `index.html`, `index.js` and `index.css`. It sets up the +DOM with a scrollable div at the top that is populated with new messages as they are broadcast. +At the bottom it adds a form to submit messages. + +The messages are received via the WebSocket `/subscribe` endpoint and published via +the HTTP POST `/publish` endpoint. The reason for not publishing messages over the WebSocket +is so that you can easily publish a message with curl. + +The server portion is `main.go` and `chat.go` and implements serving the static frontend +assets, the `/subscribe` WebSocket endpoint and the HTTP POST `/publish` endpoint. + +The code is well commented. I would recommend starting in `main.go` and then `chat.go` followed by +`index.html` and then `index.js`. + +There are two automated tests for the server included in `chat_test.go`. The first is a simple one +client echo test. It publishes a single message and ensures it's received. + +The second is a complex concurrency test where 10 clients send 128 unique messages +of max 128 bytes concurrently. The test ensures all messages are seen by every client. diff --git a/Taskfile.yml b/Taskfile.yml new file mode 100644 index 0000000..fbba28f --- /dev/null +++ b/Taskfile.yml @@ -0,0 +1,30 @@ +version: 3 +tasks: + + db-up: + desc: start ws-chat postgres and dbgate + cmds: + - echo "Postgres - localhost 54545 || DBGate localhost 7777" + - podman compose up -d + + db-down: + desc: stop ws-chat postgres and dbgate + cmds: + - podman compose down + + db-logs: + desc: view ws-chat postgres and dbgate logs + cmds: + - podman compose logs -f + + db-status: + desc: show ws-chat postgres and dbgate status + cmds: + - podman compose ps + + db-reset: + desc: delete the database data + cmds: + - podman compose down + - podman volume rm postgres-data-wschat + - podman volume rm dbgate-data-wschat diff --git a/cmd/server/chat.go b/cmd/server/chat.go new file mode 100644 index 0000000..11d305c --- /dev/null +++ b/cmd/server/chat.go @@ -0,0 +1,227 @@ +package main + +import ( + "context" + "errors" + "io" + "log" + "net" + "net/http" + "sync" + "time" + + "golang.org/x/time/rate" + + "github.com/coder/websocket" +) + +// chatServer enables broadcasting to a set of subscribers. +type chatServer struct { + // subscriberMessageBuffer controls the max number + // of messages that can be queued for a subscriber + // before it is kicked. + // + // Defaults to 16. + subscriberMessageBuffer int + + // publishLimiter controls the rate limit applied to the publish endpoint. + // + // Defaults to one publish every 100ms with a burst of 8. + publishLimiter *rate.Limiter + + // logf controls where logs are sent. + // Defaults to log.Printf. + logf func(f string, v ...any) + + // serveMux routes the various endpoints to the appropriate handler. + serveMux http.ServeMux + + subscribersMu sync.Mutex + subscribers map[*subscriber]struct{} +} + +// newChatServer constructs a chatServer with the defaults. +func newChatServer() *chatServer { + cs := &chatServer{ + subscriberMessageBuffer: 16, + logf: log.Printf, + subscribers: make(map[*subscriber]struct{}), + publishLimiter: rate.NewLimiter(rate.Every(time.Millisecond*100), 8), + } + // Serve static files (index.html, index.js, index.css) from cwd + cs.serveMux.Handle("/", http.FileServer(http.Dir("."))) + // WebSocket endpoint: clients connect here to receive messages + cs.serveMux.HandleFunc("/subscribe", cs.subscribeHandler) + // HTTP POST endpoint: clients POST here to send a message + cs.serveMux.HandleFunc("/publish", cs.publishHandler) + + return cs +} + +// subscriber represents a single connected client. +// Messages are sent on the msgs channel. If the client +// cannot keep up (buffer full), closeSlow is called +// to disconnect them. +type subscriber struct { + msgs chan []byte + closeSlow func() +} + +// ServeHTTP makes chatServer implement http.Handler, +// delegating to the internal serveMux. +func (cs *chatServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + cs.serveMux.ServeHTTP(w, r) +} + +// subscribeHandler accepts the WebSocket connection and then subscribes +// it to all future messages. +func (cs *chatServer) subscribeHandler(w http.ResponseWriter, r *http.Request) { + err := cs.subscribe(w, r) + if errors.Is(err, context.Canceled) { + return + } + if websocket.CloseStatus(err) == websocket.StatusNormalClosure || + websocket.CloseStatus(err) == websocket.StatusGoingAway { + return + } + if err != nil { + cs.logf("%v", err) + return + } +} + +// publishHandler reads the request body with a limit of 8192 bytes +// and then publishes the received message. +func (cs *chatServer) publishHandler(w http.ResponseWriter, r *http.Request) { + if r.Method != "POST" { + http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed) + return + } + body := http.MaxBytesReader(w, r.Body, 8192) + msg, err := io.ReadAll(body) + if err != nil { + http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge) + return + } + + cs.publish(msg) + + w.WriteHeader(http.StatusAccepted) +} + +// subscribe upgrades the HTTP connection to a WebSocket, +// registers the subscriber, and then loops forever writing +// messages from the subscriber's channel to the WebSocket. +// +// It uses CloseRead to keep reading control frames (close, ping, pong) +// and cancel the context if the connection drops. This means we +// never actually read message data from the WebSocket — +// messages come in via /publish instead. +func (cs *chatServer) subscribe(w http.ResponseWriter, r *http.Request) error { + // We need a mutex here because the WebSocket connection might + // be closed by closeSlow (from another goroutine) before we've + // assigned it to c. The mutex ensures we don't race. + var mu sync.Mutex + var c *websocket.Conn + var closed bool + + s := &subscriber{ + msgs: make(chan []byte, cs.subscriberMessageBuffer), + closeSlow: func() { + mu.Lock() + defer mu.Unlock() + closed = true + if c != nil { + c.Close(websocket.StatusPolicyViolation, "connection too slow to keep up with messages") + } + }, + } + + // Register the subscriber BEFORE accepting the WebSocket. + // This means we won't miss any messages published between + // the accept and the registration. + cs.addSubscriber(s) + defer cs.deleteSubscriber(s) + + // Upgrade HTTP to WebSocket + c2, err := websocket.Accept(w, r, nil) + if err != nil { + return err + } + + // Check if closeSlow was already called before we got the connection + mu.Lock() + if closed { + mu.Unlock() + return net.ErrClosed + } + c = c2 + mu.Unlock() + + defer c.CloseNow() + + // CloseRead returns a context that is cancelled when: + // - The client sends a close frame + // - The connection drops + // This frees us from manually handling ping/pong and deadlines. + ctx := c.CloseRead(context.Background()) + + // Main loop: write messages to the WebSocket as they arrive + for { + select { + case msg := <-s.msgs: + err := writeTimeout(ctx, time.Second*5, c, msg) + if err != nil { + return err + } + case <-ctx.Done(): + // Context cancelled = client disconnected + return ctx.Err() + } + } +} + +// publish sends a message to ALL subscribers. +// It never blocks — if a subscriber's buffer is full, +// they get kicked via closeSlow. +func (cs *chatServer) publish(msg []byte) { + cs.subscribersMu.Lock() + defer cs.subscribersMu.Unlock() + + // Rate limit: wait for a token before broadcasting + cs.publishLimiter.Wait(context.Background()) + + for s := range cs.subscribers { + select { + case s.msgs <- msg: + // Message delivered to subscriber's buffer + default: + // Buffer is full — kick the slow subscriber + go s.closeSlow() + } + } +} + +// addSubscriber registers a subscriber. +func (cs *chatServer) addSubscriber(s *subscriber) { + cs.subscribersMu.Lock() + cs.subscribers[s] = struct{}{} + cs.subscribersMu.Unlock() +} + +// deleteSubscriber removes a subscriber. +func (cs *chatServer) deleteSubscriber(s *subscriber) { + cs.subscribersMu.Lock() + delete(cs.subscribers, s) + cs.subscribersMu.Unlock() +} + +// writeTimeout writes a message to the WebSocket with a deadline. +// If the write doesn't complete within the timeout, the context +// is cancelled and the operation aborts. +func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, msg []byte) error { + ctx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + return c.Write(ctx, websocket.MessageText, msg) +} diff --git a/cmd/server/main.go b/cmd/server/main.go new file mode 100644 index 0000000..2220b5c --- /dev/null +++ b/cmd/server/main.go @@ -0,0 +1,28 @@ +package main + +import ( + "log/slog" + "net/http" + "os" +) + +const webPort = ":7654" + +func main() { + logger := slog.New(slog.NewJSONHandler(os.Stdout, nil)) + slog.SetDefault(logger) + + mux := http.NewServeMux() + + mux.Handle("/", http.FileServer(http.Dir("./static"))) + + cs := newChatServer() + mux.Handle("/ws/", http.StripPrefix("/ws", cs)) + + slog.Info("starting wschat web server", "port", webPort) + + if err := http.ListenAndServe(webPort, mux); err != nil { + slog.Error("unable to listen and serve", "port", webPort, "error", err) + os.Exit(1) + } +} diff --git a/compose.yml b/compose.yml new file mode 100644 index 0000000..1b19a37 --- /dev/null +++ b/compose.yml @@ -0,0 +1,57 @@ +services: + postgres: + container_name: postgres_wschat + image: docker.io/library/postgres:18 + environment: + POSTGRES_USER: wschat + POSTGRES_PASSWORD: example + POSTGRES_DB: wschat_db + volumes: + - postgres-data-wschat:/var/lib/postgresql + ports: + - "127.0.0.1:54325:5432" + shm_size: 128m + command: ["postgres", "-c", "listen_addresses=*"] + networks: + - postgres-net + restart: always + healthcheck: + test: ["CMD-SHELL", "pg_isready -U wschat -d wschat_db"] + interval: 10s + timeout: 5s + retries: 5 + start_period: 900s + + dbgate: + container_name: dbgate_wschat + image: docker.io/dbgate/dbgate:latest + environment: + CONNECTIONS: pg + LABEL_pg: Postgres + SERVER_pg: postgres + PORT_pg: 5432 + USER_pg: wschat + PASSWORD_pg: example + DATABASE_pg: wschat_db + ENGINE_pg: postgres@dbgate-plugin-postgres + volumes: + - dbgate-data-wschat:/root/.dbgate + ports: + - "127.0.0.1:7777:3000" + depends_on: + postgres: + condition: service_healthy + networks: + - postgres-net + restart: always + +networks: + postgres-net: + name: postgres-net-wschat + driver: bridge + +volumes: + postgres-data-wschat: + name: postgres-data-wschat + dbgate-data-wschat: + name: dbgate-data-wschat diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..047e1fb --- /dev/null +++ b/go.mod @@ -0,0 +1,8 @@ +module ws-chat + +go 1.26.4 + +require ( + github.com/coder/websocket v1.8.15 + golang.org/x/time v0.15.0 +) diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..6dccedf --- /dev/null +++ b/go.sum @@ -0,0 +1,4 @@ +github.com/coder/websocket v1.8.15 h1:6B2JPeOGlpff2Uz6vOEH1Vzpi0iUz20A+lPVhPHtNUA= +github.com/coder/websocket v1.8.15/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= diff --git a/static/index.css b/static/index.css new file mode 100644 index 0000000..734cd1e --- /dev/null +++ b/static/index.css @@ -0,0 +1,59 @@ +* { + margin: 0; + padding: 0; + box-sizing: border-box; +} + +body { + font-family: monospace; + display: flex; + flex-direction: column; + height: 100vh; + background: #1a1a2e; + color: #e0e0e0; +} + +#messages { + flex: 1; + overflow-y: auto; + padding: 1rem; +} + +.message { + padding: 0.4rem 0.6rem; + margin-bottom: 0.3rem; + background: #16213e; + border-radius: 4px; + word-break: break-word; +} + +#publish-form { + display: flex; + padding: 1rem; + background: #0f0f23; + border-top: 1px solid #333; +} + +#message-input { + flex: 1; + padding: 0.5rem; + background: #16213e; + border: 1px solid #333; + border-radius: 4px; + color: #e0e0e0; + font-family: monospace; +} + +#publish-form button { + margin-left: 0.5rem; + padding: 0.5rem 1rem; + background: #6d4aff; + color: white; + border: none; + border-radius: 4px; + cursor: pointer; +} + +#publish-form button:hover { + background: #5a3acc; +} diff --git a/static/index.html b/static/index.html new file mode 100644 index 0000000..76229e6 --- /dev/null +++ b/static/index.html @@ -0,0 +1,19 @@ + + + + + + WS-Chat + + + +
+ +
+ + +
+ + + + diff --git a/static/index.js b/static/index.js new file mode 100644 index 0000000..de33fa0 --- /dev/null +++ b/static/index.js @@ -0,0 +1,73 @@ +(() => { + let expectingMessage = false + + function dial() { + const conn = new WebSocket(`ws://${location.host}/ws/subscribe`) + + conn.addEventListener('close', ev => { + appendLog(`WebSocket Disconnected code: ${ev.code}, reason: ${ev.reason}`, true) + if (ev.code !== 1001) { + appendLog('Reconnecting in 1s', true) + setTimeout(dial, 1000) + } + }) + + conn.addEventListener('open', ev => { + console.info('WebSocket connected') + }) + + conn.addEventListener('message', ev => { + if (typeof ev.data !== 'string') { + console.error('unexpected message type', typeof ev.data) + return + } + const p = appendLog(ev.data) + if (expectingMessage) { + p.scrollIntoView() + expectingMessage = false + } + }) + } + + dial() + + const messageLog = document.getElementById('message-log') + const publishForm = document.getElementById('publish-form') + const messageInput = document.getElementById('message-input') + + function appendLog(text, error) { + const p = document.createElement('p') + p.innerText = `${new Date().toLocaleTimeString()}: ${text}` + if (error) { + p.style.color = 'red' + p.style.fontStyle = 'bold' + } + messageLog.append(p) + return p + } + + appendLog('Submit a message to get started!') + + publishForm.onsubmit = async ev => { + ev.preventDefault() + + const msg = messageInput.value + if (msg === '') { + return + } + messageInput.value = '' + + expectingMessage = true + try { + const resp = await fetch('/ws/publish', { + method: 'POST', + body: msg, + }) + if (resp.status !== 202) { + throw new Error(`Unexpected HTTP Status ${resp.status} ${resp.statusText}`) + } + } catch (err) { + appendLog(`Publish failed: ${err.message}`, true) + } + } +})()