Use grace in chat example

Author: nhooyr

accept.go

@@ -76,8 +76,8 @@ func accept(w http.ResponseWriter, r *http.Request, opts *AcceptOptions) (_ *Con
 	defer errd.Wrap(&err, "failed to accept WebSocket connection")
 
 	g := graceFromRequest(r)
-	if g != nil && g.isClosing() {
-		err := errors.New("server closing")
+	if g != nil && g.isShuttingdown() {
+		err := errors.New("server shutting down")
 		http.Error(w, err.Error(), http.StatusServiceUnavailable)
 		return nil, err
 	}

chat-example/go.mod

@@ -2,4 +2,6 @@ module nhooyr.io/websocket/example-chat
 
 go 1.13
 
-require nhooyr.io/websocket v1.8.2
+require nhooyr.io/websocket v0.0.0
+
+replace nhooyr.io/websocket => ../

chat-example/go.sum

@@ -1,12 +1,18 @@
+github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee h1:s+21KNqlpePfkah2I+gwHF8xmJWRjooY+5248k6m4A0=
 github.com/gobwas/httphead v0.0.0-20180130184737-2c6c146eadee/go.mod h1:L0fX3K22YWvt/FAX9NnzrNzcI4wNYi9Yku4O0LKYflo=
+github.com/gobwas/pool v0.2.0 h1:QEmUOlnSjWtnpRGHF3SauEiOsy82Cup83Vf2LcMlnc8=
 github.com/gobwas/pool v0.2.0/go.mod h1:q8bcK0KcYlCgd9e7WYLm9LpyS+YeLd8JVDW6WezmKEw=
+github.com/gobwas/ws v1.0.2 h1:CoAavW/wd/kulfZmSIBt6p24n4j7tHgNVCjsfHVNUbo=
 github.com/gobwas/ws v1.0.2/go.mod h1:szmBTxLgaFppYjEmNtny/v3w89xOydFnnZMcgRRu/EM=
+github.com/golang/protobuf v1.3.3 h1:gyjaxf+svBWX08ZjK86iN9geUJF0H6gp2IRKX6Nf6/I=
 github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw=
+github.com/google/go-cmp v0.4.0 h1:xsAVV57WRhGj6kEIi8ReJzQlHHqcBYCElAvkovg3B/4=
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/gorilla/websocket v1.4.1 h1:q7AeDBpnBk8AogcD4DSag/Ukw/KV+YhzLj2bP5HvKCM=
 github.com/gorilla/websocket v1.4.1/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/klauspost/compress v1.10.0 h1:92XGj1AcYzA6UrVdd4qIIBrT8OroryvRvdmg/IfmC7Y=
 github.com/klauspost/compress v1.10.0/go.mod h1:aoV0uJVorq1K+umq18yTdKaF57EivdYsUV+/s2qKfXs=
+golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
 golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543 h1:E7g+9GITq07hpfrRu66IVDexMakfv52eLZ2CXBWiKr4=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-nhooyr.io/websocket v1.8.2 h1:LwdzfyyOZKtVFoXay6A39Acu03KmidSZ3YUUvPa13PA=
-nhooyr.io/websocket v1.8.2/go.mod h1:LiqdCg1Cu7TPWxEvPjPa0TGYxCsy4pHNTN9gGluwBpQ=

chat-example/index.css

@@ -5,7 +5,7 @@ body {
 
 #root {
   padding: 40px 20px;
-  max-width: 480px;
+  max-width: 600px;
   margin: auto;
   height: 100vh;
 

chat-example/index.js

@@ -7,8 +7,11 @@
     const conn = new WebSocket(`ws://${location.host}/subscribe`)
 
     conn.addEventListener("close", ev => {
-      console.info("websocket disconnected, reconnecting in 1000ms", ev)
-      setTimeout(dial, 1000)
+      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")
@@ -34,10 +37,14 @@
   const messageInput = document.getElementById("message-input")
 
   // appendLog appends the passed text to messageLog.
-  function appendLog(text) {
+  function appendLog(text, error) {
     const p = document.createElement("p")
     // Adding a timestamp to each message makes the log easier to read.
     p.innerText = `${new Date().toLocaleTimeString()}: ${text}`
+    if (error) {
+      p.style.color = "red"
+      p.style.fontStyle = "bold"
+    }
     messageLog.append(p)
     return p
   }

chat-example/main.go

@@ -1,12 +1,16 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"log"
 	"net"
 	"net/http"
 	"os"
+	"os/signal"
 	"time"
+
+	"nhooyr.io/websocket"
 )
 
 func main() {
@@ -38,10 +42,31 @@ func run() error {
 	m.HandleFunc("/subscribe", ws.subscribeHandler)
 	m.HandleFunc("/publish", ws.publishHandler)
 
+	var g websocket.Grace
 	s := http.Server{
-		Handler:      m,
+		Handler:      g.Handler(m),
 		ReadTimeout:  time.Second * 10,
 		WriteTimeout: time.Second * 10,
 	}
-	return s.Serve(l)
+	errc := make(chan error, 1)
+	go func() {
+		errc <- s.Serve(l)
+	}()
+
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, os.Interrupt)
+	select {
+	case err := <-errc:
+		log.Printf("failed to serve: %v", err)
+	case sig := <-sigs:
+		log.Printf("terminating: %v", sig)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+	defer cancel()
+
+	s.Shutdown(ctx)
+	g.Shutdown(ctx)
+
+	return nil
 }

example_test.go

@@ -167,12 +167,20 @@ func ExampleGrace() {
 		ReadTimeout:  time.Second * 15,
 		WriteTimeout: time.Second * 15,
 	}
-	go s.ListenAndServe()
+
+	errc := make(chan error, 1)
+	go func() {
+		errc <- s.ListenAndServe()
+	}()
 
 	sigs := make(chan os.Signal, 1)
 	signal.Notify(sigs, os.Interrupt)
-	sig := <-sigs
-	log.Printf("recieved %v, shutting down", sig)
+	select {
+	case err := <-errc:
+		log.Printf("failed to listen and serve: %v", err)
+	case sig := <-sigs:
+		log.Printf("terminating: %v", sig)
+	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()

grace.go

@@ -15,10 +15,13 @@ import (
 // and then use Close or Shutdown to gracefully close these connections.
 //
 // Grace is intended to be used in harmony with net/http.Server's Shutdown and Close methods.
+// It's required as net/http's Shutdown and Close methods do not keep track of WebSocket
+// connections.
 type Grace struct {
-	mu      sync.Mutex
-	closing bool
-	conns   map[*Conn]struct{}
+	mu           sync.Mutex
+	closed       bool
+	shuttingDown bool
+	conns        map[*Conn]struct{}
 }
 
 // Handler returns a handler that wraps around h to record
@@ -33,10 +36,10 @@ func (g *Grace) Handler(h http.Handler) http.Handler {
 	})
 }
 
-func (g *Grace) isClosing() bool {
+func (g *Grace) isShuttingdown() bool {
 	g.mu.Lock()
 	defer g.mu.Unlock()
-	return g.closing
+	return g.shuttingDown
 }
 
 func graceFromRequest(r *http.Request) *Grace {
@@ -47,7 +50,7 @@ func graceFromRequest(r *http.Request) *Grace {
 func (g *Grace) addConn(c *Conn) error {
 	g.mu.Lock()
 	defer g.mu.Unlock()
-	if g.closing {
+	if g.closed {
 		c.Close(StatusGoingAway, "server shutting down")
 		return errors.New("server shutting down")
 	}
@@ -72,7 +75,8 @@ type gracefulContextKey struct{}
 // connections with StatusGoingAway.
 func (g *Grace) Close() error {
 	g.mu.Lock()
-	g.closing = true
+	g.shuttingDown = true
+	g.closed = true
 	var wg sync.WaitGroup
 	for c := range g.conns {
 		wg.Add(1)
@@ -97,7 +101,7 @@ func (g *Grace) Shutdown(ctx context.Context) error {
 	defer g.Close()
 
 	g.mu.Lock()
-	g.closing = true
+	g.shuttingDown = true
 	g.mu.Unlock()
 
 	// Same poll period used by net/http.