From d7e56d285e3c3c1778fdffbd7be6daf0c352ecd0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BE=90=E8=81=AA?= <xucong.053@bytedance.com>
Date: Fri, 15 Jul 2022 17:47:19 +0800
Subject: [PATCH] fix: worker missing state

---
 hrp/internal/boomer/boomer.go |  2 +-
 hrp/internal/boomer/runner.go | 16 +++++++++++-----
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/hrp/internal/boomer/boomer.go b/hrp/internal/boomer/boomer.go
index 69b782b4..b4fe7256 100644
--- a/hrp/internal/boomer/boomer.go
+++ b/hrp/internal/boomer/boomer.go
@@ -499,7 +499,7 @@ func (b *Boomer) GetWorkersInfo() []WorkerNode {
 // GetMasterInfo gets master information
 func (b *Boomer) GetMasterInfo() map[string]interface{} {
 	masterInfo := make(map[string]interface{})
-	masterInfo["state"] = getStateName(b.masterRunner.getState())
+	masterInfo["state"] = b.masterRunner.getState()
 	masterInfo["workers"] = b.masterRunner.server.getClientsLength()
 	masterInfo["target_users"] = b.masterRunner.getSpawnCount()
 	return masterInfo
diff --git a/hrp/internal/boomer/runner.go b/hrp/internal/boomer/runner.go
index 9a2f8794..cf6ceb11 100644
--- a/hrp/internal/boomer/runner.go
+++ b/hrp/internal/boomer/runner.go
@@ -874,8 +874,11 @@ func (r *workerRunner) run() {
 				}
 				r.updateState(StateInit)
 			}
-			if atomic.LoadInt32(&r.client.failCount) > 2 {
-				r.updateState(StateMissing)
+			if atomic.LoadInt32(&r.client.failCount) > 3 {
+				go r.stop()
+				if !r.isStarting() && !r.isStopping() {
+					r.updateState(StateMissing)
+				}
 			}
 			CPUUsage := GetCurrentCPUPercent()
 			MemoryUsage := GetCurrentMemoryPercent()
@@ -1007,12 +1010,15 @@ func (r *masterRunner) heartbeatWorker() {
 				if !ok {
 					log.Error().Msg("failed to get worker information")
 				}
-				if atomic.LoadInt32(&workerInfo.Heartbeat) <= 0 {
+				if atomic.LoadInt32(&workerInfo.Heartbeat) < 0 {
+					if workerInfo.getState() == StateQuitting {
+						return true
+					}
 					if workerInfo.getState() != StateMissing {
 						workerInfo.setState(StateMissing)
 					}
-					if r.getState() == StateRunning {
-						// all running workers missed, stopping runner
+					if r.isStopping() {
+						// all running workers missed, setting state to stopped
 						if r.server.getClientsLength() <= 0 {
 							r.updateState(StateStopped)
 						}