From ef72cba8634c960116e03643aa06dd1e99a5453f Mon Sep 17 00:00:00 2001 From: Kevin Krakauer Date: Tue, 23 May 2023 13:16:15 -0700 Subject: [PATCH] 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 | 12 +++ test/benchmarks/tcp/tcp_proxy.go | 134 +++++++++++++-------------- 2 files changed, 79 insertions(+), 67 deletions(-) diff --git a/test/benchmarks/tcp/tcp_benchmark.sh b/test/benchmarks/tcp/tcp_benchmark.sh index 719c9135a..910154246 100755 --- a/test/benchmarks/tcp/tcp_benchmark.sh +++ b/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 diff --git a/test/benchmarks/tcp/tcp_proxy.go b/test/benchmarks/tcp/tcp_proxy.go index d184aec7b..28eacd9a7 100644 --- a/test/benchmarks/tcp/tcp_proxy.go +++ b/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) } } }