From 6a90fe18ee397404d25427e64b75c4eb48739885 Mon Sep 17 00:00:00 2001
From: Toby <tobyxdd@gmail.com>
Date: Sun, 30 Jun 2024 20:16:55 -0700
Subject: [PATCH] feat: graceful speed test shutdown

---
 app/cmd/speedtest.go | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/app/cmd/speedtest.go b/app/cmd/speedtest.go
index 12d3a83718..f5c37fb81a 100644
--- a/app/cmd/speedtest.go
+++ b/app/cmd/speedtest.go
@@ -3,6 +3,9 @@ package cmd
 import (
 	"errors"
 	"fmt"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/spf13/cobra"
@@ -68,11 +71,26 @@ func runSpeedtest(cmd *cobra.Command, args []string) {
 		zap.Bool("udpEnabled", info.UDPEnabled),
 		zap.Uint64("tx", info.Tx))
 
-	if !skipDownload {
-		runDownloadTest(c)
-	}
-	if !skipUpload {
-		runUploadTest(c)
+	signalChan := make(chan os.Signal, 1)
+	signal.Notify(signalChan, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(signalChan)
+
+	runChan := make(chan struct{}, 1)
+	go func() {
+		if !skipDownload {
+			runDownloadTest(c)
+		}
+		if !skipUpload {
+			runUploadTest(c)
+		}
+		runChan <- struct{}{}
+	}()
+
+	select {
+	case <-signalChan:
+		logger.Info("received signal, shutting down gracefully")
+	case <-runChan:
+		logger.Info("speed test complete")
 	}
 }