control/controlclient: restore aggressive Direct.Close teardown

In the earlier http2 package migration (1d93bdce20ddd2, #17394) I had
removed Direct.Close's tracking of the connPool, thinking it wasn't
necessary.

Some tests (in another repo) are strict and like it to tear down the
world and wait, to check for leaked goroutines. And they caught this
letting some goroutines idle past Close, even if they'd eventually
close down on their own.

This restores the connPool accounting and the aggressife close.

Updates #17305
Updates #17394

Change-Id: I5fed283a179ff7c3e2be104836bbe58b05130cc7
Signed-off-by: Brad Fitzpatrick <[email protected]>

control/controlclient/direct.go

@@ -64,7 +64,7 @@ import (
 
 // Direct is the client that connects to a tailcontrol server for a node.
 type Direct struct {
-	httpc                 *http.Client // HTTP client used to talk to tailcontrol
+	httpc                 *http.Client // HTTP client used to do TLS requests to control (just https://controlplane.tailscale.com/key?v=123)
 	interceptedDial       *atomic.Bool // if non-nil, pointer to bool whether ScreenTime intercepted our dial
 	dialer                *tsdial.Dialer
 	dnsCache              *dnscache.Resolver
@@ -97,7 +97,7 @@ type Direct struct {
 	serverNoiseKey  key.MachinePublic
 
 	sfGroup     singleflight.Group[struct{}, *ts2021.Client] // protects noiseClient creation.
-	noiseClient *ts2021.Client
+	noiseClient *ts2021.Client                               // also protected by mu
 
 	persist                 persist.PersistView
 	authKey                 string

control/ts2021/client.go

@@ -28,6 +28,8 @@ import (
 	"tailscale.com/tstime"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
+	"tailscale.com/util/mak"
+	"tailscale.com/util/set"
 )
 
 // Client provides a http.Client to connect to tailcontrol over
@@ -44,8 +46,9 @@ type Client struct {
 	httpsPort string // the fallback Noise-over-https port or empty if none
 
 	// mu protects the following
-	mu     sync.Mutex
-	closed bool
+	mu       sync.Mutex
+	closed   bool
+	connPool set.HandleSet[*Conn] // all live connections
 }
 
 // ClientOpts contains options for the [NewClient] function. All fields are
@@ -175,9 +178,15 @@ func NewClient(opts ClientOpts) (*Client, error) {
 // It is a no-op and returns nil if the connection is already closed.
 func (nc *Client) Close() error {
 	nc.mu.Lock()
-	defer nc.mu.Unlock()
+	live := nc.connPool
 	nc.closed = true
+	nc.mu.Unlock()
+
+	for _, c := range live {
+		c.Close()
+	}
 	nc.Client.CloseIdleConnections()
+
 	return nil
 }
 
@@ -249,18 +258,31 @@ func (nc *Client) dial(ctx context.Context) (*Conn, error) {
 		return nil, err
 	}
 
-	ncc := NewConn(clientConn.Conn)
-
 	nc.mu.Lock()
+
+	handle := set.NewHandle()
+	ncc := NewConn(clientConn.Conn, func() { nc.noteConnClosed(handle) })
+	mak.Set(&nc.connPool, handle, ncc)
+
 	if nc.closed {
 		nc.mu.Unlock()
 		ncc.Close() // Needs to be called without holding the lock.
 		return nil, errors.New("noise client closed")
 	}
+
 	defer nc.mu.Unlock()
 	return ncc, nil
 }
 
+// noteConnClosed notes that the *Conn with the given handle has closed and
+// should be removed from the live connPool (which is usually of size 0 or 1,
+// except perhaps briefly 2 during a network failure and reconnect).
+func (nc *Client) noteConnClosed(handle set.Handle) {
+	nc.mu.Lock()
+	defer nc.mu.Unlock()
+	nc.connPool.Delete(handle)
+}
+
 // post does a POST to the control server at the given path, JSON-encoding body.
 // The provided nodeKey is an optional load balancing hint.
 func (nc *Client) Post(ctx context.Context, path string, nodeKey key.NodePublic, body any) (*http.Response, error) {

control/ts2021/conn.go

@@ -31,6 +31,7 @@ import (
 type Conn struct {
 	*controlbase.Conn
 
+	onClose           func()        // or nil
 	readHeaderOnce    sync.Once     // guards init of reader field
 	reader            io.Reader     // (effectively Conn.Reader after header)
 	earlyPayloadReady chan struct{} // closed after earlyPayload is set (including set to nil)
@@ -44,11 +45,12 @@ type Conn struct {
 // http2.ClientConn will be created that reads from the returned Conn.
 //
 // connID should be a unique ID for this connection. When the Conn is closed,
-// the onClose function will be called with the connID if it is non-nil.
-func NewConn(conn *controlbase.Conn) *Conn {
+// the onClose function will be called if it is non-nil.
+func NewConn(conn *controlbase.Conn, onClose func()) *Conn {
 	return &Conn{
 		Conn:              conn,
 		earlyPayloadReady: make(chan struct{}),
+		onClose:           sync.OnceFunc(onClose),
 	}
 }
 
@@ -103,6 +105,14 @@ func (c *Conn) Read(p []byte) (n int, err error) {
 	return c.reader.Read(p)
 }
 
+// Close closes the connection.
+func (c *Conn) Close() error {
+	if c.onClose != nil {
+		defer c.onClose()
+	}
+	return c.Conn.Close()
+}
+
 // readHeader reads the optional "early payload" from the server that arrives
 // after the Noise handshake but before the HTTP/2 session begins.
 //

util/set/handle.go

@@ -9,20 +9,28 @@ package set
 type HandleSet[T any] map[Handle]T
 
 // Handle is an opaque comparable value that's used as the map key in a
-// HandleSet. The only way to get one is to call HandleSet.Add.
+// HandleSet.
 type Handle struct {
 	v *byte
 }
 
+// NewHandle returns a new handle value.
+func NewHandle() Handle {
+	return Handle{new(byte)}
+}
+
 // Add adds the element (map value) e to the set.
 //
-// It returns the handle (map key) with which e can be removed, using a map
-// delete.
+// It returns a new handle (map key) with which e can be removed, using a map
+// delete or the [HandleSet.Delete] method.
 func (s *HandleSet[T]) Add(e T) Handle {
-	h := Handle{new(byte)}
+	h := NewHandle()
 	if *s == nil {
 		*s = make(HandleSet[T])
 	}
 	(*s)[h] = e
 	return h
 }
+
+// Delete removes the element with handle h from the set.
+func (s HandleSet[T]) Delete(h Handle) { delete(s, h) }