From 81b217d102f38528b1d873a828ec7c48a5c14699 Mon Sep 17 00:00:00 2001
From: Aji <prabowo.aji@gmail.com>
Date: Sun, 28 Mar 2021 17:47:37 +0700
Subject: [PATCH] Adding handler to stop loadtest when no worker is connected

---
 pkg/server/server.go | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/pkg/server/server.go b/pkg/server/server.go
index f6e2ee6..c0306a5 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -176,6 +176,7 @@ func (s *Server) acceptWorkerConn(responseWriter http.ResponseWriter, req *http.
 	defer logger.Infow("Worker removed", "name", name)
 	defer s.workerService.RemoveWorker(conn)
 	defer conn.Close()
+	defer s.stopLoadTestIfNoWorkerRemaining()
 
 	for {
 		_, message, err := conn.ReadMessage()
@@ -187,6 +188,14 @@ func (s *Server) acceptWorkerConn(responseWriter http.ResponseWriter, req *http.
 	}
 }
 
+func (s *Server) stopLoadTestIfNoWorkerRemaining(){
+	val := len(s.workerService.workers) - 1
+	logger.Infow("Number of remaining workers: ", "", val)
+	if val == 0 {
+		s.loadTestState = messages.ServerStateStopped
+	}
+}
+
 func (s *Server) acceptNotificationConn(responseWriter http.ResponseWriter, req *http.Request, _ httprouter.Params) {
 	conn, err := s.upgrader.Upgrade(responseWriter, req, nil)
 	if err != nil {