proper exit/handling with immediate cancel() instead of sleep

mstore/dnsWatch.go

@@ -3,6 +3,7 @@ package mstore
 import (
 	"context"
 	"net"
+	"sync"
 	"time"
 
 	"fortio.org/dflag"
@@ -14,40 +15,46 @@ var (
 	peers             sets.Set[string]
 	cancel            context.CancelFunc
 	DNSWatchSleepTime = dflag.New(15*time.Second, "Sleep time between DNS resolution")
+	wg                sync.WaitGroup
 )
 
+// Resolve the service name to a list of IPs.
+func checkDNS(serviceName string) {
+	ips, err := net.LookupHost(serviceName)
+	if err != nil {
+		log.Errf("Error resolving service %q: %v", serviceName, err)
+		return
+	}
+	newIPs := sets.FromSlice(ips)
+	log.LogVf("Resolved %q to %v", serviceName, newIPs)
+	// If the list changes, update the peers list
+	if newIPs.Equals(peers) {
+		log.LogVf("No change in peers: %v", peers)
+		return
+	}
+	peers = newIPs.Clone()
+	log.Infof("Updated peers: %v", peers)
+	_ = Peers.SetV(peers)
+}
+
 func dnsWatcher(ctx context.Context, serviceName string) {
+	defer wg.Done()
 	for {
+		checkDNS(serviceName) // first time, without waiting or cancel check
 		select {
 		case <-ctx.Done():
 			log.Warnf("DNS Watcher for %q exiting", serviceName)
 			return
-		default:
-			// Resolve the service name to a list of IPs
-			ips, err := net.LookupHost(serviceName)
-			if err != nil {
-				log.Errf("Error resolving service %q: %v", serviceName, err)
-				time.Sleep(DNSWatchSleepTime.Get()) // Sleep for a minute before resolving again
-				continue
-			}
-			newIPs := sets.FromSlice(ips)
-			log.LogVf("Resolved %q to %v", serviceName, newIPs)
-			// If the list changes, update the peers list
-			if !newIPs.Equals(peers) {
-				peers = newIPs.Clone()
-				log.Infof("Updated peers: %v", peers)
-				_ = Peers.SetV(peers)
-			} else {
-				log.LogVf("No change in peers: %v", peers)
-			}
-			time.Sleep(DNSWatchSleepTime.Get()) // Sleep for a minute before resolving again
+		case <-time.After(DNSWatchSleepTime.Get()):
+			checkDNS(serviceName)
 		}
 	}
 }
 
 func DNSWatcher(serviceName string) context.CancelFunc {
 	ctx := context.Background()
 	ctx, cancel = context.WithCancel(ctx)
+	wg.Add(1)
 	go dnsWatcher(ctx, serviceName)
 	return cancel
 }
@@ -58,3 +65,11 @@ func StartDNSWatch(serviceName string) {
 	}
 	cancel = DNSWatcher(serviceName)
 }
+
+func StopDNSWatch() {
+	if cancel != nil {
+		cancel()
+		cancel = nil
+		wg.Wait()
+	}
+}

mstore/mstore.go

@@ -46,3 +46,8 @@ func Start() {
 	log.Infof("memstore Start()")
 	// peerChange does get call for even initial flag value
 }
+
+func Stop() {
+	StopDNSWatch()
+	log.Infof("memstore stopped")
+}

proto.go

@@ -19,4 +19,5 @@ func main() {
 	log.Infof("Starting memstore prototype...")
 	mstore.Start()
 	scli.UntilInterrupted()
+	mstore.Stop()
 }