Comment chat example

Author: nhooyr

conn_test.go

@@ -329,7 +329,7 @@ func newConnTest(t testing.TB, dialOpts *websocket.DialOptions, acceptOpts *webs
 	}
 	t.Helper()
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second * 30)
 	tt = &connTest{t: t, ctx: ctx}
 	tt.appendDone(cancel)
 

example/chat.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"context"
+	"errors"
 	"io"
 	"io/ioutil"
 	"log"
@@ -12,21 +13,36 @@ import (
 	"nhooyr.io/websocket"
 )
 
+// chatServer enables broadcasting to a set of subscribers.
 type chatServer struct {
 	subscribersMu sync.RWMutex
-	subscribers   map[chan []byte]struct{}
+	subscribers   map[chan<- []byte]struct{}
 }
 
+// subscribeHandler accepts the WebSocket connection and then subscribes
+// it to all future messages.
 func (cs *chatServer) subscribeHandler(w http.ResponseWriter, r *http.Request) {
 	c, err := websocket.Accept(w, r, nil)
 	if err != nil {
 		log.Print(err)
 		return
 	}
 
-	cs.subscribe(r.Context(), c)
+	err = cs.subscribe(r.Context(), c)
+	if errors.Is(err, context.Canceled) {
+		return
+	}
+	if websocket.CloseStatus(err) == websocket.StatusNormalClosure ||
+			websocket.CloseStatus(err) == websocket.StatusGoingAway {
+		return
+	}
+	if err != nil {
+		log.Print(err)
+	}
 }
 
+// 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)
@@ -35,12 +51,44 @@ func (cs *chatServer) publishHandler(w http.ResponseWriter, r *http.Request) {
 	body := io.LimitReader(r.Body, 8192)
 	msg, err := ioutil.ReadAll(body)
 	if err != nil {
+		http.Error(w, http.StatusText(http.StatusRequestEntityTooLarge), http.StatusRequestEntityTooLarge)
 		return
 	}
 
 	cs.publish(msg)
 }
 
+// subscribe subscribes the given WebSocket to all broadcast messages.
+// It creates a msgs chan with a buffer of 16 to give some room to slower
+// connections and then registers it. It then listens for all messages
+// and writes them to the WebSocket. If the context is cancelled or
+// an error occurs, it returns and deletes the subscription.
+//
+// It uses CloseRead to keep reading from the connection to process control
+// messages and cancel the context if the connection drops.
+func (cs *chatServer) subscribe(ctx context.Context, c *websocket.Conn) error {
+	ctx = c.CloseRead(ctx)
+
+	msgs := make(chan []byte, 16)
+	cs.addSubscriber(msgs)
+	defer cs.deleteSubscriber(msgs)
+
+	for {
+		select {
+		case msg := <-msgs:
+			err := writeTimeout(ctx, time.Second*5, c, msg)
+			if err != nil {
+				return err
+			}
+		case <-ctx.Done():
+			return ctx.Err()
+		}
+	}
+}
+
+// publish publishes the msg to all subscribers.
+// It never blocks and so messages to slow subscribers
+// are dropped.
 func (cs *chatServer) publish(msg []byte) {
 	cs.subscribersMu.RLock()
 	defer cs.subscribersMu.RUnlock()
@@ -53,41 +101,24 @@ func (cs *chatServer) publish(msg []byte) {
 	}
 }
 
-func (cs *chatServer) addSubscriber(msgs chan []byte) {
+// addSubscriber registers a subscriber with a channel
+// on which to send messages.
+func (cs *chatServer) addSubscriber(msgs chan<- []byte) {
 	cs.subscribersMu.Lock()
 	if cs.subscribers == nil {
-		cs.subscribers = make(map[chan []byte]struct{})
+		cs.subscribers = make(map[chan<- []byte]struct{})
 	}
 	cs.subscribers[msgs] = struct{}{}
 	cs.subscribersMu.Unlock()
 }
 
+// deleteSubscriber deletes the subscriber with the given msgs channel.
 func (cs *chatServer) deleteSubscriber(msgs chan []byte) {
 	cs.subscribersMu.Lock()
 	delete(cs.subscribers, msgs)
 	cs.subscribersMu.Unlock()
 }
 
-func (cs *chatServer) subscribe(ctx context.Context, c *websocket.Conn) error {
-	ctx = c.CloseRead(ctx)
-
-	msgs := make(chan []byte, 16)
-	cs.addSubscriber(msgs)
-	defer cs.deleteSubscriber(msgs)
-
-	for {
-		select {
-		case msg := <-msgs:
-			err := writeTimeout(ctx, time.Second*5, c, msg)
-			if err != nil {
-				return err
-			}
-		case <-ctx.Done():
-			return ctx.Err()
-		}
-	}
-}
-
 func writeTimeout(ctx context.Context, timeout time.Duration, c *websocket.Conn, msg []byte) error {
 	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()

example/index.js

@@ -1,8 +1,10 @@
 ;(() => {
-  let conn
+  // expectingMessage is set to true
+  // if the user has just submitted a message
+  // and so we should scroll the next message into view when received.
   let expectingMessage = false
   function dial() {
-    conn = new WebSocket(`ws://${location.host}/subscribe`)
+    const conn = new WebSocket(`ws://${location.host}/subscribe`)
 
     conn.addEventListener("close", ev => {
       console.info("websocket disconnected, reconnecting in 1000ms", ev)
@@ -11,6 +13,8 @@
     conn.addEventListener("open", ev => {
       console.info("websocket connected")
     })
+
+    // This is where we handle messages received.
     conn.addEventListener("message", ev => {
       if (typeof ev.data !== "string") {
         console.error("unexpected message type", typeof ev.data)
@@ -29,15 +33,18 @@
   const publishForm = document.getElementById("publish-form")
   const messageInput = document.getElementById("message-input")
 
+  // appendLog appends the passed text to messageLog.
   function appendLog(text) {
     const p = document.createElement("p")
+    // Adding a timestamp to each message makes the log easier to read.
     p.innerText = `${new Date().toLocaleTimeString()}: ${text}`
     messageLog.append(p)
     return p
   }
   appendLog("Submit a message to get started!")
 
-  publishForm.onsubmit = async ev => {
+  // onsubmit publishes the message from the user when the form is submitted.
+  publishForm.onsubmit = ev => {
     ev.preventDefault()
 
     const msg = messageInput.value

example/main.go

@@ -18,6 +18,8 @@ func main() {
 	}
 }
 
+// run initializes the chatServer and routes and then
+// starts a http.Server for the passed in address.
 func run() error {
 	if len(os.Args) < 2 {
 		return errors.New("please provide an address to listen on as the first argument")