[client] Hand off dialed connections to the sweeper atomically

WrapDialContext and WrapConn registered the dial and the connection
independently, so a sweep landing between the dial finishing and WrapConn
cancelled only the dial registration: the connection dialed on the old
network entered the fresh registry and survived the network change.

Replace the pair with a Dial handle. Sweep marks pending dials under the
sweeper mutex, and WrapConn decides under the same mutex: a swept dial's
connection is closed and ErrSwept returned, so the caller redials on the
new network; otherwise the connection transfers to the registry with no
window in between.

client/netsweep/netsweep_test.go

@@ -12,8 +12,8 @@ import (
 func TestSweepClosesRegisteredConns(t *testing.T) {
 	sweeper := New()
 
-	c1 := sweeper.WrapConn(connPair(t))
-	c2 := sweeper.WrapConn(connPair(t))
+	c1 := wrap(t, sweeper, connPair(t))
+	c2 := wrap(t, sweeper, connPair(t))
 
 	assert.Equal(t, 2, sweeper.Sweep(), "both live connections should be closed")
 
@@ -30,7 +30,7 @@ func TestSweepClosesRegisteredConns(t *testing.T) {
 func TestCloseDeregisters(t *testing.T) {
 	sweeper := New()
 
-	conn := sweeper.WrapConn(connPair(t))
+	conn := wrap(t, sweeper, connPair(t))
 	require.NoError(t, conn.Close())
 
 	assert.Equal(t, 0, sweeper.Sweep(), "closed connection must leave the registry")
@@ -39,7 +39,7 @@ func TestCloseDeregisters(t *testing.T) {
 func TestCloseIsIdempotent(t *testing.T) {
 	sweeper := New()
 
-	conn := sweeper.WrapConn(connPair(t))
+	conn := wrap(t, sweeper, connPair(t))
 	require.NoError(t, conn.Close())
 	assert.Error(t, conn.Close(), "double close surfaces the underlying error but must not panic")
 }
@@ -47,50 +47,86 @@ func TestCloseIsIdempotent(t *testing.T) {
 func TestSweepOnlyAffectsOlderConns(t *testing.T) {
 	sweeper := New()
 
-	_ = sweeper.WrapConn(connPair(t))
+	_ = wrap(t, sweeper, connPair(t))
 	assert.Equal(t, 1, sweeper.Sweep())
 
 	// A connection dialed after the sweep must survive until the next one.
-	_ = sweeper.WrapConn(connPair(t))
+	_ = wrap(t, sweeper, connPair(t))
 	assert.Equal(t, 1, sweeper.Sweep(), "post-sweep connection belongs to the next sweep")
 }
 
 func TestSweepAbortsInFlightDials(t *testing.T) {
 	sweeper := New()
 
-	dialCtx, release := sweeper.WrapDialContext(context.Background())
-	defer release()
+	dial := sweeper.StartDial(context.Background())
+	defer dial.Release()
 
 	sweeper.Sweep()
 
-	assert.ErrorIs(t, dialCtx.Err(), context.Canceled, "sweep must cancel the in-flight dial context")
+	assert.ErrorIs(t, dial.Ctx().Err(), context.Canceled, "sweep must cancel the in-flight dial context")
 }
 
 func TestReleasedDialIsNotAborted(t *testing.T) {
 	sweeper := New()
 
 	// Simulate a dial that finished before the sweep.
-	_, release := sweeper.WrapDialContext(context.Background())
-	release()
+	released := sweeper.StartDial(context.Background())
+	released.Release()
 
 	// A dial still in flight during the sweep.
-	pendingCtx, pendingRelease := sweeper.WrapDialContext(context.Background())
-	defer pendingRelease()
+	pending := sweeper.StartDial(context.Background())
+	defer pending.Release()
 
 	sweeper.Sweep()
-	assert.ErrorIs(t, pendingCtx.Err(), context.Canceled, "pending dial must be aborted")
+	assert.ErrorIs(t, pending.Ctx().Err(), context.Canceled, "pending dial must be aborted")
+}
+
+func TestSweepBetweenDialAndHandoffClosesConn(t *testing.T) {
+	sweeper := New()
+
+	dial := sweeper.StartDial(context.Background())
+	defer dial.Release()
+
+	// The dial succeeds on the old network, then the sweep lands before the
+	// connection is handed over.
+	conn := connPair(t)
+	assert.Equal(t, 0, sweeper.Sweep(), "the connection is not registered yet")
+
+	wrapped, err := dial.WrapConn(conn)
+	require.ErrorIs(t, err, ErrSwept)
+	require.Nil(t, wrapped)
+
+	buf := make([]byte, 1)
+	_, err = conn.Read(buf)
+	assert.Error(t, err, "the old-network connection must be closed, not leaked")
+
+	assert.Equal(t, 0, sweeper.Sweep(), "nothing may leak into the next sweep")
 }
 
 func TestNilSweeperIsNoop(t *testing.T) {
 	var sweeper *Sweeper
 
 	conn := connPair(t)
-	assert.Equal(t, conn, sweeper.WrapConn(conn), "nil sweeper must return the conn unchanged")
-	assert.Equal(t, 0, sweeper.Sweep(), "nil sweeper closes nothing")
+	dial := sweeper.StartDial(context.Background())
+	defer dial.Release()
 
-	ctx, release := sweeper.WrapDialContext(context.Background())
-	release()
-	assert.NoError(t, ctx.Err(), "nil sweeper must not cancel the dial context")
+	wrapped, err := dial.WrapConn(conn)
+	require.NoError(t, err)
+	assert.Equal(t, conn, wrapped, "nil sweeper must return the conn unchanged")
+	assert.NoError(t, dial.Ctx().Err(), "nil sweeper must not cancel the dial context")
+	assert.Equal(t, 0, sweeper.Sweep(), "nil sweeper closes nothing")
+}
+
+// wrap registers conn with the sweeper through a completed dial.
+func wrap(t *testing.T, sweeper *Sweeper, conn net.Conn) net.Conn {
+	t.Helper()
+
+	dial := sweeper.StartDial(context.Background())
+	defer dial.Release()
+
+	wrapped, err := dial.WrapConn(conn)
+	require.NoError(t, err)
+	return wrapped
 }
 
 // connPair dials a loopback TCP connection and keeps the accepted peer open