roundtripper: give back the stream an unused round tripper holds

e2e_test.go

@@ -1093,6 +1093,64 @@ func TestNewRoundTripperFailsWhenDestinationUnreachable(t *testing.T) {
 	}
 }
 
+// An unused round tripper holds a stream until it's closed. A race that discards its losers
+// must get those streams back, or the session runs out of them for good.
+func TestClosedRoundTrippersGiveBackTheirStreams(t *testing.T) {
+	// Only Close may give the streams back here, not their expiry.
+	defer func(ttl time.Duration) { preDialedTTL = ttl }(preDialedTTL)
+	preDialedTTL = time.Hour
+	srv, pub := startServer(t)
+	hs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	t.Cleanup(hs.Close)
+	addr := strings.TrimPrefix(hs.URL, "http://")
+	c, _ := NewClient(Config{Zone: testZone, ServerKey: pub, Resolvers: []string{srv}, MaxStreams: 2})
+	defer c.Close()
+	for range 2 {
+		rt, err := c.NewRoundTripper(context.Background(), addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+		rt.(io.Closer).Close()
+	}
+	if _, err := connectWithin(c, addr, 10*time.Second); err != nil {
+		t.Fatalf("connecting after closing every unused round tripper: %v", err)
+	}
+}
+
+func TestUnusedRoundTripperGivesBackItsStream(t *testing.T) {
+	defer func(ttl time.Duration) { preDialedTTL = ttl }(preDialedTTL)
+	preDialedTTL = 200 * time.Millisecond
+	srv, pub := startServer(t)
+	hs := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	t.Cleanup(hs.Close)
+	addr := strings.TrimPrefix(hs.URL, "http://")
+	c, _ := NewClient(Config{Zone: testZone, ServerKey: pub, Resolvers: []string{srv}, MaxStreams: 2})
+	defer c.Close()
+	for range 2 {
+		if _, err := c.NewRoundTripper(context.Background(), addr); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if _, err := connectWithin(c, addr, 10*time.Second); err != nil {
+		t.Fatalf("connecting after the unused round trippers expired: %v", err)
+	}
+}
+
+// connectWithin retries NewRoundTripper until it succeeds: a closed stream gives back its slot
+// only once the server acknowledges its FIN.
+func connectWithin(c *Client, addr string, d time.Duration) (http.RoundTripper, error) {
+	deadline := time.Now().Add(d)
+	for {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		rt, err := c.NewRoundTripper(ctx, addr)
+		cancel()
+		if err == nil || time.Now().After(deadline) {
+			return rt, err
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
 // errInjector forwards to server but first answers every query with a forged error.
 func errInjector(t *testing.T, server string, rcode byte, tc bool) string {
 	front, _ := net.ListenPacket("udp", "127.0.0.1:0")

roundtripper.go

@@ -14,6 +14,10 @@ import (
 // phase and falls back to another transport only when connecting fails. TLS runs end to end, so
 // the server only relays ciphertext. With Close, MaxLength and RequestTimeout, this satisfies
 // kindling's DNS-tunnel transport (dnstt.DNSTT).
+//
+// The confirmed connection holds one of the session's streams until a request takes it. Close
+// the returned transport (it implements io.Closer) when it won't be used, as a race does when
+// another transport wins; otherwise the connection is closed after preDialedTTL.
 func (c *Client) NewRoundTripper(ctx context.Context, addr string) (http.RoundTripper, error) {
 	var pre *preDialed
 	if addr != "" {
@@ -26,29 +30,54 @@ func (c *Client) NewRoundTripper(ctx context.Context, addr string) (http.RoundTr
 			return nil, err
 		}
 		pre = &preDialed{conn: conn, addr: addr}
+		pre.mu.Lock()
+		pre.timer = time.AfterFunc(preDialedTTL, pre.close)
+		pre.mu.Unlock()
 	} else if _, err := c.session(ctx); err != nil {
 		return nil, err
 	}
-	return &http.Transport{
-		DialContext: func(ctx context.Context, network, a string) (net.Conn, error) {
-			if conn := pre.take(a); conn != nil {
-				return conn, nil
-			}
-			return c.DialContext(ctx, network, a)
+	return &roundTripper{
+		Transport: &http.Transport{
+			DialContext: func(ctx context.Context, network, a string) (net.Conn, error) {
+				if conn := pre.take(a); conn != nil {
+					return conn, nil
+				}
+				return c.DialContext(ctx, network, a)
+			},
+			MaxIdleConns:        4,
+			MaxIdleConnsPerHost: 2,
+			// Idle keep-alives hold a stream open, and an open stream keeps a poll outstanding.
+			IdleConnTimeout:     5 * time.Second,
+			TLSHandshakeTimeout: 30 * time.Second,
 		},
-		MaxIdleConns:        4,
-		MaxIdleConnsPerHost: 2,
-		// Idle keep-alives hold a stream open, and an open stream keeps a poll outstanding.
-		IdleConnTimeout:     5 * time.Second,
-		TLSHandshakeTimeout: 30 * time.Second,
+		pre: pre,
 	}, nil
 }
 
+// preDialedTTL bounds how long a confirmed connection waits for its first request. A session
+// has only MaxStreams streams, so one never taken must not hold its stream for good.
+var preDialedTTL = 10 * time.Second
+
+// roundTripper is the transport NewRoundTripper returns.
+type roundTripper struct {
+	*http.Transport
+	pre *preDialed
+}
+
+// Close releases the confirmed connection if no request took it, and closes idle connections.
+// Requests already in flight are unaffected.
+func (rt *roundTripper) Close() error {
+	rt.pre.close()
+	rt.CloseIdleConnections()
+	return nil
+}
+
 // preDialed hands a confirmed connection to the first dial for its address.
 type preDialed struct {
-	mu   sync.Mutex
-	conn net.Conn
-	addr string
+	mu    sync.Mutex
+	conn  net.Conn
+	addr  string
+	timer *time.Timer
 }
 
 func (p *preDialed) take(addr string) net.Conn {
@@ -62,9 +91,27 @@ func (p *preDialed) take(addr string) net.Conn {
 	}
 	conn := p.conn
 	p.conn = nil
+	p.timer.Stop()
 	return conn
 }
 
+// close closes the connection unless a dial already took it.
+func (p *preDialed) close() {
+	if p == nil {
+		return
+	}
+	p.mu.Lock()
+	conn := p.conn
+	p.conn = nil
+	if p.timer != nil {
+		p.timer.Stop()
+	}
+	p.mu.Unlock()
+	if conn != nil {
+		conn.Close()
+	}
+}
+
 // MaxLength caps request bodies: uploads ride ~50-byte queries, so big ones belong elsewhere.
 func (c *Client) MaxLength() int { return 64 << 10 }