handle cancellation and print status

signal/loadtest/README.md

@@ -92,6 +92,14 @@ go build -o signal-loadtest
 ./signal-loadtest -h
 ```
 
+**Graceful Shutdown:**
+
+The load test supports graceful shutdown via Ctrl+C (SIGINT/SIGTERM):
+- Press Ctrl+C to interrupt the test at any time
+- All active clients will be closed gracefully
+- A final aggregated report will be printed showing metrics up to the point of interruption
+- Shutdown timeout: 5 seconds (after which the process will force exit)
+
 **Available Flags:**
 - `-server`: Signal server URL (http:// or https://) (default: `http://localhost:10000`)
 - `-pairs-per-sec`: Peer pairs created per second (default: 10)
@@ -233,6 +241,30 @@ The load test collects and reports:
 - **Throughput**: Pairs per second (actual)
 - **Latency Statistics**: Min, Max, Avg message exchange latency
 
+## Graceful Shutdown Example
+
+You can interrupt a long-running test at any time with Ctrl+C:
+
+```
+./signal-loadtest -server http://localhost:10000 -pairs-per-sec 10 -total-pairs 100 -exchange-duration 10m
+
+# Press Ctrl+C after some time...
+^C
+WARN[0045]
+Received interrupt signal, shutting down gracefully...
+
+=== Load Test Report ===
+Test Duration: 45.234s
+Total Pairs Sent: 75
+Successful Exchanges: 75
+Failed Exchanges: 0
+Total Messages Exchanged: 22500
+Total Errors: 0
+Throughput: 1.66 pairs/sec
+...
+========================
+```
+
 ## Test Results
 
 Example output from a 20 pairs/sec test:

signal/loadtest/cmd/signal-loadtest/main.go

@@ -1,9 +1,12 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	log "github.com/sirupsen/logrus"
@@ -90,13 +93,41 @@ func main() {
 	}
 	fmt.Println()
 
-	lt := loadtest.NewLoadTest(config)
+	// Set up signal handler for graceful shutdown
-	if err := lt.Run(); err != nil {
+	ctx, cancel := context.WithCancel(context.Background())
-		log.Errorf("Load test failed: %v", err)
+	defer cancel()
-		os.Exit(1)
+
+	sigChan := make(chan os.Signal, 1)
+	signal.Notify(sigChan, os.Interrupt, syscall.SIGTERM)
+
+	lt := loadtest.NewLoadTestWithContext(ctx, config)
+
+	// Run load test in a goroutine
+	done := make(chan error, 1)
+	go func() {
+		done <- lt.Run()
+	}()
+
+	// Wait for completion or signal
+	select {
+	case <-sigChan:
+		log.Warnf("\nReceived interrupt signal, shutting down gracefully...")
+		cancel()
+		// Wait a bit for graceful shutdown
+		select {
+		case <-done:
+		case <-time.After(5 * time.Second):
+			log.Warnf("Shutdown timeout, forcing exit")
+		}
+	case err := <-done:
+		if err != nil && err != context.Canceled {
+			log.Errorf("Load test failed: %v", err)
+			os.Exit(1)
+		}
 	}
 
 	metrics := lt.GetMetrics()
+	fmt.Println() // Add blank line before report
 	metrics.PrintReport()
 }
 

signal/loadtest/rate_loadtest.go

@@ -62,7 +62,17 @@ type LoadTest struct {
 // NewLoadTest creates a new load test instance
 func NewLoadTest(config LoadTestConfig) *LoadTest {
 	ctx, cancel := context.WithCancel(context.Background())
-	reporterCtx, reporterCancel := context.WithCancel(context.Background())
+	return newLoadTestWithContext(ctx, cancel, config)
+}
+
+// NewLoadTestWithContext creates a new load test instance with a custom context
+func NewLoadTestWithContext(ctx context.Context, config LoadTestConfig) *LoadTest {
+	ctx, cancel := context.WithCancel(ctx)
+	return newLoadTestWithContext(ctx, cancel, config)
+}
+
+func newLoadTestWithContext(ctx context.Context, cancel context.CancelFunc, config LoadTestConfig) *LoadTest {
+	reporterCtx, reporterCancel := context.WithCancel(ctx)
 	config.IDPrefix = fmt.Sprintf("%d-", time.Now().UnixNano())
 	return &LoadTest{
 		config:         config,