Wait for readers to exit before closing their fds on linux (#1979)

interface.go

@@ -124,6 +124,8 @@ type Interface struct {
 	// reused (every udp.Conn ListenOut guarantees that ordering).
 	batchers []*batch.MultiCoalescer
 	wg       sync.WaitGroup
+	// readers covers just the listenIn and listenOut goroutines, wg also holds the construction token Close releases
+	readers sync.WaitGroup
 
 	udpTxErrLimiter      *ratelimit.Limiter // udp direct transmit error log limiter
 	udpRelayTxErrLimiter *ratelimit.Limiter // udp relay transmit error log limiter
@@ -341,14 +343,18 @@ func (f *Interface) activate() error {
 func (f *Interface) run() {
 	// Launch n queues to read packets from udp
 	for i := 0; i < f.routines; i++ {
+		f.readers.Add(1)
 		f.wg.Go(func() {
+			defer f.readers.Done()
 			f.listenOut(i)
 		})
 	}
 
 	// Launch n queues to read packets from tun dev
 	for i := 0; i < f.routines; i++ {
+		f.readers.Add(1)
 		f.wg.Go(func() {
+			defer f.readers.Done()
 			f.listenIn(f.queues[i], i)
 		})
 	}
@@ -683,6 +689,14 @@ func (f *Interface) GetCertState() *CertState {
 	return f.pki.getCertState()
 }
 
+// shutdowner is a udp.Conn or overlay.Device whose readers can be woken without releasing their fds
+type shutdowner interface {
+	Shutdown() error
+}
+
+// readerShutdownTimeout bounds how long Close waits on woken readers before releasing their fds anyway
+const readerShutdownTimeout = 5 * time.Second
+
 // Close releases the interface's resources: the udp sockets and the tun device.
 // It is idempotent and safe to call at any point in the lifecycle, including on an interface that never activated,
 // calls after the first return nil without doing anything.
@@ -692,20 +706,55 @@ func (f *Interface) Close() error {
 	}
 
 	var errs []error
-
-	// Release the udp readers
-	for i, u := range f.writers {
-		err := u.Close()
-		if err != nil {
+	closeUdp := func(i int, u udp.Conn) {
+		if err := u.Close(); err != nil {
 			f.l.Error("Error while closing udp socket", "error", err, "writer", i)
 			errs = append(errs, err)
 		}
 	}
+	closeTun := func() {
+		if err := f.inside.Close(); err != nil {
+			errs = append(errs, err)
+		}
+	}
 
-	// Release the tun device (closing the tun also closes all readers)
-	closeErr := f.inside.Close()
-	if closeErr != nil {
-		errs = append(errs, closeErr)
+	// Wake every reader before releasing any fd, one that loops back into a syscall on a reused fd number never wakes.
+	// Anything without Shutdown is closed right away, as before.
+	var deferred []func()
+	for i, u := range f.writers {
+		if s, ok := u.(shutdowner); ok {
+			if err := s.Shutdown(); err != nil {
+				f.l.Error("Error while shutting down udp socket", "error", err, "writer", i)
+				errs = append(errs, err)
+			}
+			deferred = append(deferred, func() { closeUdp(i, u) })
+		} else {
+			closeUdp(i, u)
+		}
+	}
+	if s, ok := f.inside.(shutdowner); ok {
+		if err := s.Shutdown(); err != nil {
+			errs = append(errs, err)
+		}
+		deferred = append(deferred, closeTun)
+	} else {
+		closeTun()
+	}
+
+	if len(deferred) > 0 {
+		readersDone := make(chan struct{})
+		go func() {
+			f.readers.Wait()
+			close(readersDone)
+		}()
+		select {
+		case <-readersDone:
+		case <-time.After(readerShutdownTimeout):
+			f.l.Warn("Readers did not stop, closing their fds anyway", "timeout", readerShutdownTimeout)
+		}
+		for _, closeFn := range deferred {
+			closeFn()
+		}
 	}
 
 	// Release the construction token so waiters know the resources are gone

overlay/tio/queueset_gso_linux.go

@@ -20,6 +20,7 @@ type offloadQueueSet struct {
 	// usoEnabled is true when newTun successfully negotiated TUN_F_USO4|6 with the kernel.
 	// Queues created by Add inherit this and surface it via Offload.USOSupported so coalescers can gate USO emission.
 	usoEnabled bool
+	shutdown   atomic.Bool
 	closed     atomic.Bool
 	// l is handed to each queue for its bad-vnet-header drop logging.
 	l *slog.Logger
@@ -69,6 +70,17 @@ func (c *offloadQueueSet) wakeForShutdown() error {
 	return err
 }
 
+// Shutdown wakes every reader without releasing any fds, they see POLLIN on the eventfd and return os.ErrClosed
+func (c *offloadQueueSet) Shutdown() error {
+	if c.shutdown.Swap(true) {
+		return nil
+	}
+	for _, x := range c.pq {
+		x.shutdown.Store(true)
+	}
+	return c.wakeForShutdown()
+}
+
 func (c *offloadQueueSet) Close() error {
 	if c.closed.Swap(true) {
 		return nil
@@ -76,9 +88,7 @@ func (c *offloadQueueSet) Close() error {
 
 	errs := []error{}
 
-	// Signal all readers blocked in poll to wake up and exit.
-	// They observe POLLIN on the shutdown eventfd and return os.ErrClosed.
-	if err := c.wakeForShutdown(); err != nil {
+	if err := c.Shutdown(); err != nil {
 		errs = append(errs, err)
 	}
 

overlay/tio/queueset_poll_linux.go

@@ -16,6 +16,7 @@ type pollQueueSet struct {
 	// pqi is exactly the same as pq, but stored as the interface type
 	pqi        []Queue
 	shutdownFd int
+	shutdown   atomic.Bool
 	closed     atomic.Bool
 }
 
@@ -59,6 +60,17 @@ func (c *pollQueueSet) wakeForShutdown() error {
 	return err
 }
 
+// Shutdown wakes every reader without releasing any fds, they see POLLIN on the eventfd and return os.ErrClosed
+func (c *pollQueueSet) Shutdown() error {
+	if c.shutdown.Swap(true) {
+		return nil
+	}
+	for _, x := range c.pq {
+		x.shutdown.Store(true)
+	}
+	return c.wakeForShutdown()
+}
+
 func (c *pollQueueSet) Close() error {
 	if c.closed.Swap(true) {
 		return nil
@@ -66,9 +78,7 @@ func (c *pollQueueSet) Close() error {
 
 	errs := []error{}
 
-	// Signal all readers blocked in poll to wake up and exit.
-	// They observe POLLIN on the shutdown eventfd and return os.ErrClosed.
-	if err := c.wakeForShutdown(); err != nil {
+	if err := c.Shutdown(); err != nil {
 		errs = append(errs, err)
 	}
 

overlay/tio/tio.go

@@ -7,6 +7,8 @@ import (
 // QueueSet holds one or many Queue objects and helps close them in an orderly way.
 type QueueSet interface {
 	io.Closer
+	// Shutdown wakes every reader without releasing any fds, Close only once they have exited
+	Shutdown() error
 	Queues() []Queue
 
 	// Add takes a tun fd, adds it to the set, and prepares it for use as a Queue.

overlay/tio/tio_gso_linux.go

@@ -66,7 +66,9 @@ type Offload struct {
 	// usoEnabled records whether the kernel agreed to TUN_F_USO* on this FD,
 	// so writers can decide whether emitting GSO_UDP_L4 superpackets is safe.
 	usoEnabled bool
-	closed     atomic.Bool
+	// shutdown is set before shutdownFd is signaled, a busy reader never polls to see that
+	shutdown atomic.Bool
+	closed   atomic.Bool
 
 	// gsoHdrBuf is a per-queue 10-byte scratch for the virtio_net_hdr emitted
 	// by WriteGSO. Kept separate from the read-only package-level validVnetHdr
@@ -176,6 +178,9 @@ func (r *Offload) readPacket(block bool) (int, error) {
 // This amortizes the poll wake over bursts of small packets (e.g. TCP ACKs).
 // Packet.Bytes slices point into the Offload's internal buffer and are only valid until the next Read or Close on this Queue.
 func (r *Offload) Read() ([]Packet, error) {
+	if r.shutdown.Load() {
+		return nil, os.ErrClosed
+	}
 	r.pending = r.pending[:0]
 	r.rxOff = 0
 

overlay/tio/tio_poll_linux.go

@@ -13,7 +13,9 @@ import (
 type Poll struct {
 	fd         int
 	shutdownFd int
-	closed     atomic.Bool
+	// shutdown is set before shutdownFd is signaled, a busy reader never polls to see that
+	shutdown atomic.Bool
+	closed   atomic.Bool
 
 	readBuf  []byte
 	batchRet [1]Packet
@@ -52,6 +54,9 @@ func (t *Poll) blockOnWrite() error {
 // config currently lose that batching. blockOn and the EAGAIN plumbing are
 // already shared; kept one-packet-per-Read for now to preserve behavior.
 func (t *Poll) Read() ([]Packet, error) {
+	if t.shutdown.Load() {
+		return nil, os.ErrClosed
+	}
 	n, err := t.readOne(t.readBuf)
 	if err != nil {
 		return nil, err

overlay/tio/tun_file_linux_test.go

@@ -223,3 +223,30 @@ func TestOffloadQueueSet_Close_ClosesShutdownFd(t *testing.T) {
 	// Second Close must not touch fds (shutdownFd is now -1) and must return nil.
 	require.NoError(t, qs.Close())
 }
+
+// A busy reader never polls the shutdown eventfd, so Read has to stop on its own with data still waiting
+func TestQueueSet_ShutdownStopsReaderWithDataWaiting(t *testing.T) {
+	sets := map[string]func() (QueueSet, error){
+		"poll":    NewPollQueueSet,
+		"offload": func() (QueueSet, error) { return NewOffloadQueueSet(false, slog.New(slog.DiscardHandler)) },
+	}
+	for name, newSet := range sets {
+		t.Run(name, func(t *testing.T) {
+			var fds [2]int
+			require.NoError(t, unix.Pipe2(fds[:], unix.O_CLOEXEC))
+			t.Cleanup(func() { _ = unix.Close(fds[1]) })
+
+			qs, err := newSet()
+			require.NoError(t, err)
+			require.NoError(t, qs.Add(fds[0]))
+			t.Cleanup(func() { _ = qs.Close() })
+
+			_, err = unix.Write(fds[1], make([]byte, 64))
+			require.NoError(t, err)
+			require.NoError(t, qs.Shutdown())
+
+			_, err = qs.Queues()[0].Read()
+			require.ErrorIs(t, err, os.ErrClosed)
+		})
+	}
+}

overlay/tun_linux.go

@@ -832,6 +832,11 @@ func (t *tun) updateRoutes(r netlink.RouteUpdate) {
 	t.routeTree.Store(newTree)
 }
 
+// Shutdown wakes the tun readers without releasing their fds, Close releases them
+func (t *tun) Shutdown() error {
+	return t.readers.Shutdown()
+}
+
 func (t *tun) Close() error {
 	t.closeLock.Lock()
 	defer t.closeLock.Unlock()

udp/udp_linux.go

@@ -255,6 +255,10 @@ func (u *StdConn) ListenOut(r EncReader, flush func()) error {
 	msgs, buffers, names, _ := prepareRawMessages(u.batch, bufSize, cmsgSpace)
 
 	for {
+		// A shut down socket still queues and returns datagrams, a busy reader would never see an empty read
+		if u.closed.Load() {
+			return net.ErrClosed
+		}
 		if cmsgSpace > 0 {
 			for i := range msgs {
 				setMsgControllen(&msgs[i].Hdr, cmsgSpace)
@@ -446,12 +450,16 @@ func (u *StdConn) getMemInfo(meminfo *[unix.SK_MEMINFO_VARS]uint32) error {
 	return nil
 }
 
-func (u *StdConn) Close() error {
+// Shutdown wakes the reader without releasing the fd, Close only once it has exited
+func (u *StdConn) Shutdown() error {
 	u.closed.Store(true)
-	// Wake the reader parked in recvmmsg. shutdown(2) on an unconnected socket
-	// returns ENOTCONN but still wakes it, so ignore the error.
-	// The reader then sees closed and stops touching the fd, making the Close below safe.
+	// shutdown(2) on an unconnected socket returns ENOTCONN but still wakes the reader, so ignore the error
 	_ = unix.Shutdown(u.sysFd, unix.SHUT_RDWR)
+	return nil
+}
+
+func (u *StdConn) Close() error {
+	_ = u.Shutdown()
 	return unix.Close(u.sysFd)
 }
 

udp/udp_linux_test.go

@@ -76,7 +76,7 @@ func TestShutdownWakesAfterRx_Mechanism(t *testing.T) {
 }
 
 // TestListenOutTeardown_TrafficPatterns reproduces the field report: a blocking
-// reader must tear down cleanly on Close() regardless of what the socket has
+// reader must tear down cleanly on Shutdown() regardless of what the socket has
 // carried. The three cases the report called out:
 //
 //	no traffic ever      -> works (shutdown wakes recvmmsg with n==0)
@@ -162,11 +162,12 @@ func runTeardownCase(t *testing.T, batch int, name string, traffic func(send net
 	// empty queue with the socket already having received a packet.
 	time.Sleep(500 * time.Millisecond)
 
+	// Tear down the way Interface.Close does, the fd stays open until the reader is gone
 	start := time.Now()
-	if err := sc.Close(); err != nil {
-		t.Fatalf("Close: %v", err)
+	if err := sc.Shutdown(); err != nil {
+		t.Fatalf("Shutdown: %v", err)
 	}
-	close(stop)
+	defer func() { _ = sc.Close() }()
 
 	select {
 	case err := <-loopDone:
@@ -179,7 +180,75 @@ func runTeardownCase(t *testing.T, batch int, name string, traffic func(send net
 	case <-time.After(3 * time.Second):
 		buf := make([]byte, 1<<20)
 		n := runtime.Stack(buf, true)
-		t.Fatalf("%s: HANG, ListenOut did not return within 3s of Close\n%s", name, buf[:n])
+		t.Fatalf("%s: HANG, ListenOut did not return within 3s of Shutdown\n%s", name, buf[:n])
 	}
+	close(stop)
 	<-trafficDone
 }
+
+// A shut down socket still returns queued datagrams, so a reader that is behind has to stop on its own within the
+// batch it was working on
+func TestListenOutShutdown_BusyReader(t *testing.T) {
+	const batch = 64
+	udpSettings := Settings{Listen: netip.MustParseAddrPort("127.0.0.1:0"), Multi: true, Batch: batch}
+	c, err := NewListener(testLogger(), udpSettings)
+	if err != nil {
+		t.Fatalf("NewListener: %v", err)
+	}
+	sc := c.(*StdConn)
+	defer sc.Close()
+	addr, err := sc.LocalAddr()
+	if err != nil {
+		t.Fatalf("LocalAddr: %v", err)
+	}
+
+	stop := make(chan struct{})
+	defer close(stop)
+	for range 8 {
+		send, err := net.Dial("udp", addr.String())
+		if err != nil {
+			t.Fatalf("dial: %v", err)
+		}
+		defer send.Close()
+		go func() {
+			b := make([]byte, 1200)
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+					_, _ = send.Write(b)
+				}
+			}
+		}()
+	}
+
+	var afterShutdown atomic.Int64
+	loopDone := make(chan error, 1)
+	go func() {
+		loopDone <- sc.ListenOut(func(netip.AddrPort, []byte) {
+			if sc.closed.Load() {
+				afterShutdown.Add(1)
+			}
+			// Slower than the senders so the queue never drains
+			for start := time.Now(); time.Since(start) < 200*time.Microsecond; {
+			}
+		}, func() {})
+	}()
+
+	time.Sleep(200 * time.Millisecond)
+	if err := sc.Shutdown(); err != nil {
+		t.Fatalf("Shutdown: %v", err)
+	}
+	select {
+	case err := <-loopDone:
+		if !errors.Is(err, net.ErrClosed) {
+			t.Fatalf("ListenOut returned %v, want net.ErrClosed", err)
+		}
+	case <-time.After(3 * time.Second):
+		t.Fatalf("ListenOut still running 3s after Shutdown, delivered %d packets since", afterShutdown.Load())
+	}
+	if n := afterShutdown.Load(); n > batch {
+		t.Fatalf("delivered %d packets after Shutdown, more than the one batch in flight", n)
+	}
+}