package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"net/http"
	"os"
	"time"

	"golang.org/x/net/http2"
)

func assert(err error) {
	if err != nil {
		panic(err)
	}
}

type handler struct{}

func (h *handler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	res.WriteHeader(200)
	res.(http.Flusher).Flush()
	<-req.Context().Done()
}

func makeServer() *http.Server {
	httpServer := &http.Server{
		Addr: ":" + os.Args[1],
	}

	httpServer.ConnState = func(_ net.Conn, state http.ConnState) {
		if state == http.StateNew {
			fmt.Println("connections++")
		} else if state == http.StateClosed {
			fmt.Println("connections--")
		}
	}

	err := http2.ConfigureServer(httpServer, &http2.Server{})
	assert(err)

	httpServer.Handler = &handler{}

	go func() {
		err = httpServer.ListenAndServeTLS(os.Args[2], os.Args[3])
		assert(err)
	}()

	return httpServer
}

func makeRequest(httpClient *http.Client) context.CancelFunc {
	ctx, cancel := context.WithCancel(context.Background())
	request, err := http.NewRequest("GET", "https://localhost:"+os.Args[1], nil)
	assert(err)

	request = request.WithContext(ctx)
	response, err := httpClient.Do(request)
	assert(err)

	err = response.Body.Close()
	assert(err)

	return cancel
}

func makeClient() *http.Client {
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}

	err := http2.ConfigureTransport(transport)
	assert(err)

	return &http.Client{Transport: transport}
}

func main() {
	if len(os.Args) != 4 {
		fmt.Println("Usage: ", os.Args[0], " <port> <cert> <key>")
		os.Exit(1)
	}

	httpServer := makeServer()
	time.Sleep(time.Second)
	httpClient := makeClient()

	for i := 0; i < 1001; i++ {
		cancel := makeRequest(httpClient)
		cancel()
		if i%50 == 0 {
			fmt.Println("Requests: ", i)
		}
	}

	err := httpServer.Close()
	assert(err)
}

diff --git a/http2/transport.go b/http2/transport.go
index 850d7ae..fc97566 100644
--- a/http2/transport.go
+++ b/http2/transport.go
@@ -221,7 +221,7 @@ type clientStream struct {
 // request to be done (any way it might be removed from the cc.streams
 // map: peer reset, successful completion, TCP connection breakage,
 // etc)
-func (cs *clientStream) awaitRequestCancel(req *http.Request) {
+func (cs *clientStream) awaitRequestCancel(req *http.Request, rl *clientConnReadLoop) {
        ctx := reqContext(req)
        if req.Cancel == nil && ctx.Done() == nil {
                return
@@ -229,10 +229,10 @@ func (cs *clientStream) awaitRequestCancel(req *http.Request) {
        select {
        case <-req.Cancel:
                cs.cancelStream()
-               cs.bufPipe.CloseWithError(errRequestCanceled)
+               rl.endStreamError(cs, errRequestCanceled)
        case <-ctx.Done():
                cs.cancelStream()
-               cs.bufPipe.CloseWithError(ctx.Err())
+               rl.endStreamError(cs, ctx.Err())
        case <-cs.done:
        }
 }
@@ -243,6 +243,8 @@ func (cs *clientStream) cancelStream() {
        cs.didReset = true
        cs.cc.mu.Unlock()
 
+       cs.cc.forgetStreamID(cs.ID)
+
        if !didReset {
                cs.cc.writeStreamReset(cs.ID, ErrCodeCancel, nil)
        }
@@ -1531,7 +1533,7 @@ func (rl *clientConnReadLoop) handleResponse(cs *clientStream, f *MetaHeadersFra
        cs.bufPipe = pipe{b: &dataBuffer{expected: res.ContentLength}}
        cs.bytesRemain = res.ContentLength
        res.Body = transportResponseBody{cs}
-       go cs.awaitRequestCancel(cs.req)
+       go cs.awaitRequestCancel(cs.req, rl)
 
        if cs.requestedGzip && res.Header.Get("Content-Encoding") == "gzip" {
                res.Header.Del("Content-Encoding")

go version go1.8.3 linux/amd64
GOARCH="amd64"
GOBIN=""
GOEXE=""
GOHOSTARCH="amd64"
GOHOSTOS="linux"
GOOS="linux"
GOPATH="/home/ben/development"
GORACE=""
GOROOT="/usr/local/go"
GOTOOLDIR="/usr/local/go/pkg/tool/linux_amd64"
GCCGO="gccgo"
CC="gcc"
GOGCCFLAGS="-fPIC -m64 -pthread -fmessage-length=0 -fdebug-prefix-map=/tmp/go-build098570469=/tmp/go-build -gno-record-gcc-switches"
CXX="g++"
CGO_ENABLED="1"
PKG_CONFIG="pkg-config"
CGO_CFLAGS="-g -O2"
CGO_CPPFLAGS=""
CGO_CXXFLAGS="-g -O2"
CGO_FFLAGS="-g -O2"
CGO_LDFLAGS="-g -O2"
GOROOT/bin/go version: go version go1.8.3 linux/amd64
GOROOT/bin/go tool compile -V: compile version go1.8.3 X:framepointer
uname -sr: Linux 4.12.4-041204-generic
LSB Version:	core-9.20160110ubuntu5-amd64:core-9.20160110ubuntu5-noarch:printing-9.20160110ubuntu5-amd64:printing-9.20160110ubuntu5-noarch:security-9.20160110ubuntu5-amd64:security-9.20160110ubuntu5-noarch
Distributor ID:	Ubuntu
Description:	Ubuntu 17.04
Release:	17.04
Codename:	zesty
/lib/x86_64-linux-gnu/libc.so.6: GNU C Library (Ubuntu GLIBC 2.24-9ubuntu2.2) stable release version 2.24, by Roland McGrath et al.
gdb --version: GNU gdb (Ubuntu 7.12.50.20170314-0ubuntu1.1) 7.12.50.20170314-git