Go gorilla/websocket program demonstrating server doesn't cancel context on connection close

websocket-problem.go

package main

import (
        "context"
        "log"
        "net/http"
        "sync"
        "time"

        "github.com/gorilla/websocket"
)

var (
        upgrader = websocket.Upgrader{
                ReadBufferSize:  1024,
                WriteBufferSize: 1024,
                CheckOrigin: func(r *http.Request) bool {
                        return true
                },
        }
)

func handler(w http.ResponseWriter, r *http.Request) {
        tick := time.NewTicker(10 * time.Second)
        defer tick.Stop()
        select {
        case <-r.Context().Done():
                log.Print("server> context finished, yay!")
        case <-tick.C:
                log.Print("server> request timeout, boo! <--------------")
        }
}

func websocketHandler(w http.ResponseWriter, r *http.Request) {
        conn, err := upgrader.Upgrade(w, r, nil)
        if err != nil {
                log.Fatalf("wsserver> websocket upgrader: %v", err)
        }
        defer conn.Close()

        conn.SetCloseHandler(func(code int, text string) error {
                log.Printf("wsserver> close handler called")
                return nil
        })

        tick := time.NewTicker(10 * time.Second)
        defer tick.Stop()
        select {
        case <-r.Context().Done():
                log.Print("wsserver> context finished, yay!")
        case <-tick.C:
                log.Print("wsserver> request timeout, boo! <-----------")
        }
}

func websocketClient() {
        log.Print("wsclient> starting")
        ctx, cancel := context.WithTimeout(context.Background(), time.Second)
        defer cancel()

        conn, _, err := websocket.DefaultDialer.DialContext(ctx, "ws://127.0.0.1:4711/ws", nil)
        if err != nil {
                log.Fatalf("wsclient> failed to connect: %v", err)
        }
        if false {
                conn.ReadMessage()
        }
        select {
        case <-ctx.Done():
        }
        log.Printf("wsclient> client context has been cancelled (server should know)")
        conn.Close()
        log.Printf("wsclient> client has now closed connection (server *definitely* should know)")
}

func client() {
        log.Print("client> starting")
        ctx, cancel := context.WithTimeout(context.Background(), time.Second)
        defer cancel()
        req, err := http.NewRequestWithContext(ctx, "GET", "http://localhost:4711/req", nil)
        if err != nil {
                log.Fatalf("client> Failed to create request: %v", err)
        }
        _, err = http.DefaultClient.Do(req)
        log.Printf("client> request finished with: %v", err)
}

func main() {
        http.HandleFunc("/req", handler)
        http.HandleFunc("/ws", websocketHandler)
        var wg sync.WaitGroup
        wg.Add(1)
        go http.ListenAndServe(":4711", nil)

        time.Sleep(time.Second) // Give time for server to start.

        log.Printf("---------First a normal request-------")
        client()
        log.Printf("---------Now a websocket request-------")
        websocketClient()

        wg.Wait() // block forever
}