diff --git a/opencloud/pkg/command/benchmark.go b/opencloud/pkg/command/benchmark.go index ab28ea95a..334b4089b 100644 --- a/opencloud/pkg/command/benchmark.go +++ b/opencloud/pkg/command/benchmark.go @@ -11,8 +11,11 @@ import ( "net/http" "os" "os/exec" + "os/signal" "strconv" "strings" + "sync" + "syscall" "time" "github.com/olekukonko/tablewriter" @@ -115,6 +118,18 @@ func BenchmarkClientCommand(cfg *config.Config) *cli.Command { }, Category: "benchmark", Action: func(c *cli.Context) error { + // Set up signal handling for Ctrl+C + ctx, cancel := context.WithCancel(c.Context) + defer cancel() + + sigChan := make(chan os.Signal, 1) + signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM) + go func() { + <-sigChan + fmt.Println("\nReceived interrupt signal, shutting down...") + cancel() + }() + opt := clientOptions{ url: c.Args().First(), insecure: c.Bool("insecure"), @@ -258,7 +273,7 @@ func BenchmarkClientCommand(cfg *config.Config) *cli.Command { defer opt.ticker.Stop() } - return client(c.Context, opt) + return client(ctx, opt) }, } @@ -283,8 +298,12 @@ func client(ctx context.Context, o clientOptions) error { status int } stats := make(chan stat) + var wg sync.WaitGroup + for i := 0; i < o.jobs; i++ { + wg.Add(1) go func(i int) { + defer wg.Done() tr := &http.Transport{ TLSClientConfig: &tls.Config{ MinVersion: tls.VersionTLS12, @@ -362,6 +381,7 @@ func client(ctx context.Context, o clientOptions) error { fmt.Printf("req %d took %v and returned status %d\n", numRequests, stat.duration, stat.status) case <-ctx.Done(): fmt.Println("\nShutting down...") + wg.Wait() return nil } } @@ -384,6 +404,7 @@ func client(ctx context.Context, o clientOptions) error { fmt.Printf("\n%d req at %v/req\n", numRequests, duration/time.Duration(numRequests)) } fmt.Println("Shutting down...") + wg.Wait() return nil } }