From 206d98e84be6cc309f3fbe9eb34844f0c7883a28 Mon Sep 17 00:00:00 2001 From: Brad Fitzpatrick Date: Thu, 2 Oct 2025 18:29:54 -0700 Subject: [PATCH] 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 --- control/controlclient/direct.go | 4 ++-- control/ts2021/client.go | 32 +++++++++++++++++++++++++++----- control/ts2021/conn.go | 14 ++++++++++++-- util/set/handle.go | 16 ++++++++++++---- 4 files changed, 53 insertions(+), 13 deletions(-) diff --git a/control/controlclient/direct.go b/control/controlclient/direct.go index de577bea4..482affe33 100644 --- a/control/controlclient/direct.go +++ b/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 diff --git a/control/ts2021/client.go b/control/ts2021/client.go index 9a9a3ded8..e0b82b89c 100644 --- a/control/ts2021/client.go +++ b/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) { diff --git a/control/ts2021/conn.go b/control/ts2021/conn.go index ecf184d3c..52d663272 100644 --- a/control/ts2021/conn.go +++ b/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. // diff --git a/util/set/handle.go b/util/set/handle.go index 471ceeba2..9c6b6dab0 100644 --- a/util/set/handle.go +++ b/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) }