minor quality-of-life improvements to tcp_benchmark

- Get rid of `os.Exit()`, which didn't actually run any of the defers.
- Warn users that putting profiles in /tmp/ will lead to hours of perplexed
  frustration.

PiperOrigin-RevId: 534534483

test/benchmarks/tcp/tcp_benchmark.sh

@@ -57,6 +57,13 @@ if [[ "$?" != "0" ]]; then
   echo "warning: sch_netem may not be installed." >&2
 fi
 
+function checktmp() {
+  if [[ "$1" =~ ^/tmp/ ]]; then
+    echo "Don't use /tmp for output files ('$1') -- tcp_benchmark mounts over /tmp and your file will never make it to the root /tmp."
+    exit 1
+  fi
+}
+
 while [[ $# -gt 0 ]]; do
   case "$1" in
     --client)
@@ -135,22 +142,27 @@ while [[ $# -gt 0 ]]; do
     --cpuprofile)
       shift
       netstack_opts="${netstack_opts} -cpuprofile=$1"
+      checktmp "$1"
       ;;
     --memprofile)
       shift
       netstack_opts="${netstack_opts} -memprofile=$1"
+      checktmp "$1"
       ;;
     --blockprofile)
       shift
       netstack_opts="${netstack_opts} -blockprofile=$1"
+      checktmp "$1"
       ;;
     --mutexprofile)
       shift
       netstack_opts="${netstack_opts} -mutexprofile=$1"
+      checktmp "$1"
       ;;
     --traceprofile)
       shift
       netstack_opts="${netstack_opts} -traceprofile=$1"
+      checktmp "$1"
       ;;
     --disable-linux-gso)
       disable_linux_gso=1

test/benchmarks/tcp/tcp_proxy.go

@@ -483,88 +483,88 @@ func main() {
 
 	sigs := make(chan os.Signal, 1)
 	signal.Notify(sigs, unix.SIGTERM)
+
+	// Accept connections and proxy data between them.
 	go func() {
-		<-sigs
-		if *cpuprofile != "" {
-			pprof.StopCPUProfile()
-		}
-		if *memprofile != "" {
-			f, err := os.Create(*memprofile)
+		for {
+			// Forward all connections.
+			inConn, err := listener.Accept()
 			if err != nil {
-				log.Fatal("could not create memory profile: ", err)
+				// This should not happen; we are listening
+				// successfully. Exhausted all available FDs?
+				log.Fatalf("accept error: %v", err)
 			}
-			defer func() {
-				if err := f.Close(); err != nil {
-					log.Print("error closing memory profile: ", err)
+			log.Printf("incoming connection established.")
+
+			// Copy both ways.
+			go io.Copy(inConn, next)
+			go io.Copy(next, inConn)
+
+			// Print stats every second.
+			go func() {
+				t := time.NewTicker(time.Second)
+				defer t.Stop()
+				for {
+					<-t.C
+					in.printStats()
+					out.printStats()
 				}
 			}()
-			runtime.GC() // get up-to-date statistics
-			if err := pprof.WriteHeapProfile(f); err != nil {
-				log.Fatalf("Unable to write heap profile: %v", err)
-			}
-		}
-		if *blockprofile != "" {
-			f, err := os.Create(*blockprofile)
-			if err != nil {
-				log.Fatal("could not create block profile: ", err)
-			}
-			defer func() {
-				if err := f.Close(); err != nil {
-					log.Print("error closing block profile: ", err)
+
+			for {
+				// Dial again.
+				next, err = out.dial(*forward)
+				if err == nil {
+					break
 				}
-			}()
-			if err := pprof.Lookup("block").WriteTo(f, 0); err != nil {
-				log.Fatalf("Unable to write block profile: %v", err)
 			}
 		}
-		if *mutexprofile != "" {
-			f, err := os.Create(*mutexprofile)
-			if err != nil {
-				log.Fatal("could not create mutex profile: ", err)
-			}
-			defer func() {
-				if err := f.Close(); err != nil {
-					log.Print("error closing mutex profile: ", err)
-				}
-			}()
-			if err := pprof.Lookup("mutex").WriteTo(f, 0); err != nil {
-				log.Fatalf("Unable to write mutex profile: %v", err)
-			}
-		}
-		os.Exit(0)
 	}()
 
-	for {
-		// Forward all connections.
-		inConn, err := listener.Accept()
+	// Wait for the SIGTERM notifying us to stop.
+	<-sigs
+
+	if *memprofile != "" {
+		f, err := os.Create(*memprofile)
 		if err != nil {
-			// This should not happen; we are listening
-			// successfully. Exhausted all available FDs?
-			log.Fatalf("accept error: %v", err)
+			log.Fatal("could not create memory profile: ", err)
 		}
-		log.Printf("incoming connection established.")
-
-		// Copy both ways.
-		go io.Copy(inConn, next)
-		go io.Copy(next, inConn)
-
-		// Print stats every second.
-		go func() {
-			t := time.NewTicker(time.Second)
-			defer t.Stop()
-			for {
-				<-t.C
-				in.printStats()
-				out.printStats()
+		defer func() {
+			if err := f.Close(); err != nil {
+				log.Print("error closing memory profile: ", err)
 			}
 		}()
-
-		for {
-			// Dial again.
-			next, err = out.dial(*forward)
-			if err == nil {
-				break
+		runtime.GC() // get up-to-date statistics
+		if err := pprof.WriteHeapProfile(f); err != nil {
+			log.Fatalf("Unable to write heap profile: %v", err)
+		}
+	}
+	if *blockprofile != "" {
+		f, err := os.Create(*blockprofile)
+		if err != nil {
+			log.Fatal("could not create block profile: ", err)
+		}
+		defer func() {
+			if err := f.Close(); err != nil {
+				log.Print("error closing block profile: ", err)
 			}
+		}()
+		if err := pprof.Lookup("block").WriteTo(f, 0); err != nil {
+			log.Fatalf("Unable to write block profile: %v", err)
+		}
+	}
+	if *mutexprofile != "" {
+		f, err := os.Create(*mutexprofile)
+		if err != nil {
+			log.Fatal("could not create mutex profile: ", err)
+		}
+		defer func() {
+			if err := f.Close(); err != nil {
+				log.Print("error closing mutex profile: ", err)
+			}
+		}()
+		if err := pprof.Lookup("mutex").WriteTo(f, 0); err != nil {
+			log.Fatalf("Unable to write mutex profile: %v", err)
 		}
 	}
 }