diff --git a/conn_test.go b/conn_test.go
index 7514540dfdb8bd66e84ee0eac33edbde68eaf627..68dc837da96639868314eee78dbfb16d242ea886 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -271,12 +271,11 @@ func TestWasm(t *testing.T) {
t.Skip("skipping on CI")
}
- var g websocket.Grace
- defer g.Close()
- s := httptest.NewServer(g.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ // TODO grace
+ s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
c, err := websocket.Accept(w, r, &websocket.AcceptOptions{
- Subprotocols: []string{"echo"},
- InsecureSkipVerify: true,
+ Subprotocols: []string{"echo"},
+ OriginPatterns: []string{"*"},
})
if err != nil {
t.Errorf("echo server failed: %v", err)
@@ -291,7 +290,7 @@ func TestWasm(t *testing.T) {
t.Errorf("echo server failed: %v", err)
return
}
- })))
+ }))
defer s.Close()
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
diff --git a/examples/chat/chat_test.go b/examples/chat/chat_test.go
index 2cbc995eb54abe08ed815ae4c3f03f071eb4a2e7..79523d2a81dabc52782e642d8d67fc3ab1547b0a 100644
--- a/examples/chat/chat_test.go
+++ b/examples/chat/chat_test.go
@@ -130,11 +130,10 @@ func setupTest(t *testing.T) (url string, closeFn func()) {
cs.subscriberMessageBuffer = 4096
cs.publishLimiter.SetLimit(rate.Inf)
- var g websocket.Grace
- s := httptest.NewServer(g.Handler(cs))
+ // TODO grace
+ s := httptest.NewServer(cs)
return s.URL, func() {
s.Close()
- g.Close()
}
}
diff --git a/examples/chat/main.go b/examples/chat/main.go
index 1b6f3266cab9d37d320f74a62572830c29bbd47b..cc2d01e8d37a7e51ea22b45d18e06069957ac06c 100644
--- a/examples/chat/main.go
+++ b/examples/chat/main.go
@@ -9,8 +9,6 @@ import (
"os"
"os/signal"
"time"
-
- "nhooyr.io/websocket"
)
func main() {
@@ -36,9 +34,9 @@ func run() error {
log.Printf("listening on http://%v", l.Addr())
cs := newChatServer()
- var g websocket.Grace
+ // TODO grace
s := http.Server{
- Handler: g.Handler(cs),
+ Handler: cs,
ReadTimeout: time.Second * 10,
WriteTimeout: time.Second * 10,
}
@@ -59,8 +57,5 @@ func run() error {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
defer cancel()
- s.Shutdown(ctx)
- g.Shutdown(ctx)
-
- return nil
+ return s.Shutdown(ctx)
}
diff --git a/examples/echo/echo.go b/examples/echo/main.go
similarity index 96%
rename from examples/echo/echo.go
rename to examples/echo/main.go
index 0f31235df93576b0568b46fea90462d3f9ae8fc0..db2d06c9adfd0741b9e5d378f15a3cf27d13bb6a 100644
--- a/examples/echo/echo.go
+++ b/examples/echo/main.go
@@ -24,7 +24,7 @@ import (
// This example starts a WebSocket echo server,
// dials the server and then sends 5 different messages
// and prints out the server's responses.
-func Example_echo() {
+func main() {
// First we listen on port 0 which means the OS will
// assign us a random free port. This is the listener
// the server will serve on and the client will connect to.
@@ -34,15 +34,14 @@ func Example_echo() {
}
defer l.Close()
- var g websocket.Grace
- defer g.Close()
+ // TODO grace
s := &http.Server{
- Handler: g.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+ Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
err := echoServer(w, r)
if err != nil {
log.Printf("echo server: %v", err)
}
- })),
+ }),
ReadTimeout: time.Second * 15,
WriteTimeout: time.Second * 15,
}
@@ -61,6 +60,7 @@ func Example_echo() {
if err != nil {
log.Fatalf("client failed: %v", err)
}
+
// Output:
// received: map[i:0]
// received: map[i:1]
diff --git a/grace.go b/grace.go
deleted file mode 100644
index a0ec8969bbbc415dec8ae1537644951c2fa3c844..0000000000000000000000000000000000000000
--- a/grace.go
+++ /dev/null
@@ -1,120 +0,0 @@
-package websocket
-
-import (
- "context"
- "fmt"
- "net/http"
- "sync"
- "time"
-)
-
-// Grace enables graceful shutdown of accepted WebSocket connections.
-//
-// Use Handler to wrap WebSocket handlers to record accepted connections
-// 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.
-//
-// Make sure to Close or Shutdown the *http.Server first as you don't want to accept
-// any new connections while the existing websockets are being shut down.
-type Grace struct {
- handlersMu sync.Mutex
- closing bool
- handlers map[context.Context]context.CancelFunc
-}
-
-// Handler returns a handler that wraps around h to record
-// all WebSocket connections accepted.
-//
-// Use Close or Shutdown to gracefully close recorded connections.
-// Make sure to Close or Shutdown the *http.Server first.
-func (g *Grace) Handler(h http.Handler) http.Handler {
- return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
- ctx, cancel := context.WithCancel(r.Context())
- defer cancel()
-
- r = r.WithContext(ctx)
-
- ok := g.add(w, ctx, cancel)
- if !ok {
- return
- }
- defer g.del(ctx)
-
- h.ServeHTTP(w, r)
- })
-}
-
-func (g *Grace) add(w http.ResponseWriter, ctx context.Context, cancel context.CancelFunc) bool {
- g.handlersMu.Lock()
- defer g.handlersMu.Unlock()
-
- if g.closing {
- http.Error(w, "shutting down", http.StatusServiceUnavailable)
- return false
- }
-
- if g.handlers == nil {
- g.handlers = make(map[context.Context]context.CancelFunc)
- }
- g.handlers[ctx] = cancel
-
- return true
-}
-
-func (g *Grace) del(ctx context.Context) {
- g.handlersMu.Lock()
- defer g.handlersMu.Unlock()
-
- delete(g.handlers, ctx)
-}
-
-// Close prevents the acceptance of new connections with
-// http.StatusServiceUnavailable and closes all accepted
-// connections with StatusGoingAway.
-//
-// Make sure to Close or Shutdown the *http.Server first.
-func (g *Grace) Close() error {
- g.handlersMu.Lock()
- for _, cancel := range g.handlers {
- cancel()
- }
- g.handlersMu.Unlock()
-
- // Wait for all goroutines to exit.
- g.Shutdown(context.Background())
-
- return nil
-}
-
-// Shutdown prevents the acceptance of new connections and waits until
-// all connections close. If the context is cancelled before that, it
-// calls Close to close all connections immediately.
-//
-// Make sure to Close or Shutdown the *http.Server first.
-func (g *Grace) Shutdown(ctx context.Context) error {
- defer g.Close()
-
- // Same poll period used by net/http.
- t := time.NewTicker(500 * time.Millisecond)
- defer t.Stop()
- for {
- if g.zeroHandlers() {
- return nil
- }
-
- select {
- case <-t.C:
- case <-ctx.Done():
- return fmt.Errorf("failed to shutdown WebSockets: %w", ctx.Err())
- }
- }
-}
-
-func (g *Grace) zeroHandlers() bool {
- g.handlersMu.Lock()
- defer g.handlersMu.Unlock()
- return len(g.handlers) == 0
-}
diff --git a/ws_js.go b/ws_js.go
index 69019e61600635fd829e52a0193df2540d2b5a95..b87e32cdafb2eafa4bccf381a491283e0f96b70e 100644
--- a/ws_js.go
+++ b/ws_js.go
@@ -39,8 +39,6 @@ type Conn struct {
readSignal chan struct{}
readBufMu sync.Mutex
readBuf []wsjs.MessageEvent
-
- g *Grace
}
func (c *Conn) close(err error, wasClean bool) {