Start TCP dispatchers after endpoints are completely loaded during restore.

- During restore if endpoints' states are not loaded before starting
dispatchers, then the incoming packets will be queued and the dispatcher will
try to send the packets, which will result in a panic as the endpoint might
still be in an invalid state (eg: initial).
- To avoid this, wait for all the endpoints to be loaded completely before
starting the dispatchers by calling AsycnLoading.Wait() in stack.Restore() and
remove this from kernel.LoadFrom()
- Added a test to check save restore of multiple(100) listening sockets.

PiperOrigin-RevId: 762010964

Author: nybidari

images/basic/integrationtest/Dockerfile

@@ -15,6 +15,7 @@ RUN gcc -O2 -o test_sticky test_sticky.c
 RUN gcc -O2 -o host_fd host_fd.c
 RUN gcc -O2 -o host_connect host_connect.c
 RUN gcc -O2 -o tcp_server tcp_server.c
+RUN gcc -O2 -o tcp_stress_server tcp_stress_server.c -pthread
 
 # Add nonprivileged regular user named "nonroot".
 RUN groupadd --gid 1337 nonroot && \

images/basic/integrationtest/tcp_stress_server.c

@@ -0,0 +1,132 @@
+// Copyright 2025 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// This test validates the checkpoint/restore functionality for approximately
+// 100 listening sockets operating in network=sandbox mode. It will create the
+// sockets, accept all incoming connections, and then a checkpoint will be
+// initiated from the other end. After restore, the test verifies that these
+// listening sockets are correctly re-established and new client connections can
+// be made.
+
+#include <errno.h>
+#include <netinet/in.h>
+#include <pthread.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/socket.h>
+#include <unistd.h>
+
+#define NUM_SOCKETS 100
+#define START_PORT 9000
+#define BACKLOG 10
+
+// Structure to pass arguments to the listener thread
+typedef struct {
+  int port;
+  int listen_fd;
+} listener_arg_t;
+
+void *listener_thread(void *arg) {
+  listener_arg_t *listener_arg = (listener_arg_t *)arg;
+  int port = listener_arg->port;
+  int listen_fd = listener_arg->listen_fd;
+
+  while (1) {
+    printf("Listener on port %d started. Waiting for connection...\n", port);
+    int client_fd = accept(listen_fd, NULL, NULL);
+    if (errno == EINTR) {
+      continue;
+    }
+    if (client_fd < 0) {
+      perror("accept");
+      exit(EXIT_FAILURE);
+    }
+    printf("Accepted connection on port %d\n", port);
+    close(client_fd);
+  }
+}
+
+int main() {
+  struct sockaddr_in server_addr;
+  int listen_fds[NUM_SOCKETS];
+  pthread_t listener_threads[NUM_SOCKETS];
+  int active_listeners = 0;
+  int i;
+
+  // Initialize listen_fds with -1 to indicate no socket yet
+  for (i = 0; i < NUM_SOCKETS; ++i) {
+    listen_fds[i] = -1;
+  }
+
+  printf("Attempting to create %d listening sockets starting from port %d...\n",
+         NUM_SOCKETS, START_PORT);
+
+  for (i = 0; i < NUM_SOCKETS; ++i) {
+    int port = START_PORT + i;
+    int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
+    if (listen_fd < 0) {
+      perror("socket failed");
+      exit(EXIT_FAILURE);
+    }
+
+    memset(&server_addr, 0, sizeof(server_addr));
+    server_addr.sin_family = AF_INET;
+    server_addr.sin_addr.s_addr = INADDR_ANY;
+    server_addr.sin_port = htons(port);
+
+    if (bind(listen_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) <
+        0) {
+      perror("bind failed");
+      exit(EXIT_FAILURE);
+    }
+    if (listen(listen_fd, BACKLOG) < 0) {
+      perror("listen failed");
+      exit(EXIT_FAILURE);
+    }
+    printf("Successfully listening on port %d (FD: %d)\n", port, listen_fd);
+
+    listen_fds[active_listeners] = listen_fd;
+    active_listeners++;
+
+    // Allocate memory for the argument to pass to the thread
+    listener_arg_t *arg = (listener_arg_t *)malloc(sizeof(listener_arg_t));
+    if (arg == NULL) {
+      close(listen_fd);
+      active_listeners--;
+      exit(EXIT_FAILURE);
+    }
+    arg->port = port;
+    arg->listen_fd = listen_fd;
+
+    if (pthread_create(&listener_threads[i], NULL, listener_thread,
+                       (void *)arg) != 0) {
+      perror("pthread_create failed");
+      close(listen_fd);
+      free(arg);
+      active_listeners--;
+      exit(EXIT_FAILURE);
+    }
+  }
+
+  // Wait for all threads to finish
+  for (int i = 0; i < NUM_SOCKETS; i++) {
+    if (pthread_join(listener_threads[i], NULL) != 0) {
+      perror("failed to join thread");
+      exit(EXIT_FAILURE);
+    }
+  }
+  printf("Program finished.\n");
+  return 0;
+}

pkg/sentry/kernel/kernel.go

@@ -79,7 +79,6 @@ import (
 	"gvisor.dev/gvisor/pkg/sentry/vfs"
 	"gvisor.dev/gvisor/pkg/state"
 	"gvisor.dev/gvisor/pkg/sync"
-	"gvisor.dev/gvisor/pkg/tcpip"
 )
 
 // IOUringEnabled is set to true when IO_URING is enabled. Added as a global to
@@ -806,8 +805,6 @@ func (k *Kernel) LoadFrom(ctx context.Context, r io.Reader, loadMFs bool, timeRe
 		return vfs.PrependErrMsg("vfs.CompleteRestore() failed", err)
 	}
 
-	tcpip.AsyncLoading.Wait()
-
 	log.Infof("Overall load took [%s] after async work", time.Since(loadStart))
 
 	// Applications may size per-cpu structures based on k.applicationCores, so

pkg/tcpip/stack/stack.go

@@ -2009,6 +2009,11 @@ func (s *Stack) Restore() {
 	for _, e := range eps {
 		e.Restore(s)
 	}
+
+	// Make sure all the endpoints are loaded correctly before resuming the
+	// protocol level background workers.
+	tcpip.AsyncLoading.Wait()
+
 	// Now resume any protocol level background workers.
 	for _, p := range s.transportProtocols {
 		if saveRestoreEnabled {

test/e2e/integration_test.go

@@ -1223,37 +1223,64 @@ func TestCheckpointResume(t *testing.T) {
 	}
 }
 
-func testCheckpointRestoreListeningConnection(ctx context.Context, t *testing.T, d *dockerutil.Container) {
+func connectWithTCPServer(t *testing.T, serverIP string, startPort int, numConn int) {
+	const timeout = 1 * time.Minute
+	var conn net.Conn
+	var err error
+
+	for i := range numConn {
+		ip := serverIP + ":" + strconv.Itoa(startPort+i)
+		for {
+			conn, err = net.DialTimeout("tcp", ip, timeout)
+			// Retry to connect to the server if there are any errors.
+			// Connect can fail sometimes when the server is not ready
+			// and is not in listening state when the request was sent.
+			if err != nil {
+				t.Logf("Error connecting to server %s: %v, retrying...", ip, err)
+				continue
+			}
+			break
+		}
+		conn.Close()
+	}
+}
+
+func connectAndReadWrite(t *testing.T, serverIP string, port int) {
+	newserverIP := serverIP + ":" + strconv.Itoa(port)
+	newConn, err := net.DialTimeout("tcp", newserverIP, 1*time.Minute)
+	if err != nil {
+		t.Fatalf("Error connecting to server: %v", err)
+	}
+	defer newConn.Close()
+
+	readBuf := make([]byte, 32)
+	if _, err := newConn.Read(readBuf); err != nil {
+		t.Fatalf("Read failed: %v", err)
+	}
+
+	if _, err := newConn.Write([]byte("Hello!")); err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+}
+
+func testCheckpointRestoreListeningConnection(ctx context.Context, t *testing.T, d *dockerutil.Container, fName string, numConn int) {
 	defer d.CleanUp(ctx)
 
-	const port = 9000
 	opts := dockerutil.RunOpts{
 		Image: "basic/integrationtest",
 	}
-
 	// Start the tcp server.
-	if err := d.Spawn(ctx, opts, "./tcp_server"); err != nil {
+	if err := d.Spawn(ctx, opts, fName); err != nil {
 		t.Fatalf("docker run failed: %v", err)
 	}
 
-	var (
-		ip   net.IP
-		err  error
-		conn net.Conn
-	)
-	ip, err = d.FindIP(ctx, false)
+	ip, err := d.FindIP(ctx, false)
 	if err != nil {
 		t.Fatalf("docker.FindIP failed: %v", err)
 	}
-	serverIP := ip.String() + ":" + strconv.Itoa(port)
-	const timeout = 1 * time.Minute
-	for {
-		conn, err = net.DialTimeout("tcp", serverIP, timeout)
-		if err == nil {
-			break
-		}
-	}
-	conn.Close()
+
+	const port = 9000
+	connectWithTCPServer(t, ip.String(), port, numConn)
 
 	// Create a snapshot.
 	const checkpointFile = "networktest"
@@ -1265,30 +1292,20 @@ func testCheckpointRestoreListeningConnection(ctx context.Context, t *testing.T,
 	}
 	d.RestoreInTest(ctx, t, checkpointFile)
 
-	var (
-		newIP   net.IP
-		newConn net.Conn
-	)
+	var newIP net.IP
 	newIP, err = d.FindIP(ctx, false)
 	if err != nil {
 		t.Fatalf("docker.FindIP failed: %v", err)
 	}
-	newserverIP := newIP.String() + ":" + strconv.Itoa(port)
-	newConn, err = net.DialTimeout("tcp", newserverIP, timeout)
-	if err != nil {
-		t.Fatalf("Error connecting to server: %v", err)
-	}
-	defer newConn.Close()
-
-	readBuf := make([]byte, 32)
-	if _, err := newConn.Read(readBuf); err != nil {
-		t.Fatalf("Read failed: %v", err)
-	}
-
-	if _, err := newConn.Write([]byte("Hello!")); err != nil {
-		t.Fatalf("Write failed: %v", err)
+	if numConn > 1 {
+		connectWithTCPServer(t, newIP.String(), port, numConn)
+		if err := d.Kill(ctx); err != nil {
+			t.Fatalf("Wait failed: %v", err)
+		}
+		return
 	}
 
+	connectAndReadWrite(t, newIP.String(), port)
 	if err := d.Wait(ctx); err != nil {
 		t.Fatalf("Wait failed: %v", err)
 	}
@@ -1303,7 +1320,7 @@ func TestRestoreListenConn(t *testing.T) {
 
 	ctx := context.Background()
 	d := dockerutil.MakeContainer(ctx, t)
-	testCheckpointRestoreListeningConnection(ctx, t, d)
+	testCheckpointRestoreListeningConnection(ctx, t, d, "./tcp_server" /* fName */, 1 /* numConn */)
 }
 
 // Test to check restore of a TCP listening connection with netstack S/R.
@@ -1318,5 +1335,20 @@ func TestRestoreListenConnWithNetstackSR(t *testing.T) {
 
 	ctx := context.Background()
 	d := dockerutil.MakeContainerWithRuntime(ctx, t, "-TESTONLY-save-restore-netstack")
-	testCheckpointRestoreListeningConnection(ctx, t, d)
+	testCheckpointRestoreListeningConnection(ctx, t, d, "./tcp_server" /* fName */, 1 /* numConn */)
+}
+
+// Test to check restore of multiple TCP listening connections with netstack S/R.
+func TestRestoreMultipleListenConnWithNetstackSR(t *testing.T) {
+	if !testutil.IsCheckpointSupported() {
+		t.Skip("Checkpoint is not supported.")
+	}
+	if !testutil.IsRunningWithSaveRestoreNetstack() {
+		t.Skip("Netstack save restore is not supported.")
+	}
+	dockerutil.EnsureDockerExperimentalEnabled()
+
+	ctx := context.Background()
+	d := dockerutil.MakeContainerWithRuntime(ctx, t, "-TESTONLY-save-restore-netstack")
+	testCheckpointRestoreListeningConnection(ctx, t, d, "./tcp_stress_server" /* fName */, 100 /* numConn */)
 }