Abort on SIGTERM after waiting for the HTTP requests to finish

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
 		}
 	}