make done channel private and create public Shutdown method

Author: ttjiaa

websocketproxy.go

@@ -2,6 +2,7 @@
 package websocketproxy
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"log"
@@ -47,9 +48,9 @@ type (
 		// If nil, DefaultDialer is used.
 		Dialer *websocket.Dialer
 
-		// Done specifies a channel for which all proxied websocket connections
+		// done specifies a channel for which all proxied websocket connections
 		// can be closed on demand by closing the channel.
-		Done chan struct{}
+		done chan struct{}
 	}
 
 	websocketMsg struct {
@@ -214,7 +215,7 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 					errc <- err
 					break
 				}
-			case <-w.Done:
+			case <-w.done:
 				m := websocket.FormatCloseMessage(websocket.CloseGoingAway, "websocketproxy: closing connection")
 				dst.WriteMessage(websocket.CloseMessage, m)
 				break
@@ -234,10 +235,23 @@ func (w *WebsocketProxy) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 		if e, ok := err.(*websocket.CloseError); !ok || e.Code == websocket.CloseAbnormalClosure {
 			log.Printf("websocketproxy: Error when copying from client to backend: %v", err)
 		}
-	case <-w.Done:
+	case <-w.done:
 	}
 }
 
+// Shutdown gracefully closes proxied websocket connections by closing the
+// Done channel they are subscribed to.
+func (w *WebsocketProxy) Shutdown(ctx context.Context) error {
+	// TODO: support using context for control and return error when applicable
+	// Currently implemented such that the method signature matches http.Server.Shutdown()
+
+	if w.done == nil {
+		w.done = make(chan struct{})
+	}
+	close(w.done)
+	return nil
+}
+
 func copyHeader(dst, src http.Header) {
 	for k, vv := range src {
 		for _, v := range vv {

websocketproxy_test.go

@@ -1,6 +1,7 @@
 package websocketproxy
 
 import (
+	"context"
 	"log"
 	"net/http"
 	"net/url"
@@ -30,7 +31,6 @@ func TestProxy(t *testing.T) {
 	u, _ := url.Parse(backendURL)
 	proxy := NewProxy(u)
 	proxy.Upgrader = upgrader
-	proxy.Done = make(chan struct{})
 
 	mux := http.NewServeMux()
 	mux.Handle("/proxy", proxy)
@@ -123,5 +123,5 @@ func TestProxy(t *testing.T) {
 		t.Errorf("expecting: %s, got: %s", msg, string(p))
 	}
 
-	close(proxy.Done)
+	proxy.Shutdown(context.Background())
 }