Add the loop to stop the clustering engine

Author: fabricekabongo

clustering/engine.go

@@ -1,11 +1,16 @@
 package clustering
 
-import "github.com/fabricekabongo/loggerhead/query"
+import (
+	"context"
+
+	"github.com/fabricekabongo/loggerhead/query"
+)
 
 type EngineDecorator struct {
 	cluster     *Cluster
 	engine      *query.Engine
 	commandChan chan string
+	ctx         context.Context
 }
 
 func (e EngineDecorator) ExecuteQuery(query string) string {
@@ -15,11 +20,12 @@ func (e EngineDecorator) ExecuteQuery(query string) string {
 	return e.engine.ExecuteQuery(query)
 }
 
-func NewEngineDecorator(cluster *Cluster, engine *query.Engine) query.EngineInterface {
+func NewEngineDecorator(ctx context.Context, cluster *Cluster, engine *query.Engine) query.EngineInterface {
 	eng := &EngineDecorator{
 		cluster:     cluster,
 		engine:      engine,
 		commandChan: make(chan string),
+		ctx:         ctx,
 	}
 
 	go eng.commandLoop()
@@ -28,12 +34,12 @@ func NewEngineDecorator(cluster *Cluster, engine *query.Engine) query.EngineInte
 }
 
 func (e EngineDecorator) commandLoop() {
-
 	for {
 		select {
+		case <-e.ctx.Done():
+			return
 		case command := <-e.commandChan:
-			e.cluster.Broadcasts().QueueBroadcast(NewLocationBroadcast(command)) // At least broadcast the query to the cluster in case we go down before executing it
-
+			e.cluster.Broadcasts().QueueBroadcast(NewLocationBroadcast(command))
 		}
 	}
 }

main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"log"
@@ -18,6 +19,8 @@ import (
 )
 
 func main() {
+
+	ctx := context.Background()
 	start := time.Now()
 	cfg := config.GetConfig()
 
@@ -43,7 +46,8 @@ func main() {
 		}
 	}(cluster)
 
-	clusterEngine := clustering.NewEngineDecorator(cluster, writeEngine)
+	ClusterCtx, concel := context.WithCancel(ctx)
+	clusterEngine := clustering.NewEngineDecorator(ClusterCtx, cluster, writeEngine)
 
 	opsServer := admin.NewOpsServer(cluster, cfg)
 	go opsServer.Start()
@@ -58,22 +62,21 @@ func main() {
 
 	printWelcomeMessage(cfg, cluster)
 	sigc := make(chan os.Signal, 1)
-	signal.Notify(sigc,
-		syscall.SIGHUP,
-		syscall.SIGINT,
-		syscall.SIGTERM,
-		syscall.SIGQUIT)
-	go func() {
+	signal.Notify(sigc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
+
+	go func(cancel context.CancelFunc) {
 		s := <-sigc
 		log.Println("Received signal: ", s)
 		svr.Stop()
 		close(sigc)
+		cancel()
+
 		err := cluster.Close(0)
 		if err != nil {
 			return
 		}
 		os.Exit(0)
-	}()
+	}(concel)
 
 	end := time.Now()
 	fmt.Println("Startup time: ", end.Sub(start))