FIX: panic on dup ack #187

runner.go

@@ -297,6 +297,7 @@ type slaveRunner struct {
 	masterHost                 string
 	masterPort                 int
 	waitForAck                 sync.WaitGroup
+	ackReceived                int32
 	lastReceivedSpawnTimestamp int64
 	client                     client
 }
@@ -396,13 +397,20 @@ func (r *slaveRunner) onCustomMessage(msg *CustomMessage) {
 }
 
 func (r *slaveRunner) onAckMessage(msg *genericMessage) {
+	// Maybe we should add a state for waiting?
+	if !atomic.CompareAndSwapInt32(&r.ackReceived, 0, 1) {
+		log.Println("Receive duplicate ack message, ignored")
+		return
+	}
 	r.waitForAck.Done()
 	Events.Publish(EVENT_CONNECTED)
 }
 
 func (r *slaveRunner) sendClientReadyAndWaitForAck() {
 	r.waitForAck = sync.WaitGroup{}
 	r.waitForAck.Add(1)
+
+	atomic.StoreInt32(&r.ackReceived, 0)
 	// locust allows workers to bypass version check by sending -1 as version
 	r.client.sendChannel() <- newClientReadyMessage("client_ready", -1, r.nodeID)
 

runner_test.go

@@ -1,6 +1,7 @@
 package boomer
 
 import (
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -341,6 +342,20 @@ func TestOnQuitMessage(t *testing.T) {
 	assert.Equal(t, stateInit, runner.state)
 }
 
+func TestOnAckMessage(t *testing.T) {
+	eventCount := 0
+	Events.Subscribe(EVENT_CONNECTED, func() {
+		eventCount++
+	})
+	runner := newSlaveRunner("localhost", 5557, []*Task{}, nil)
+	runner.waitForAck = sync.WaitGroup{}
+	runner.waitForAck.Add(1)
+
+	runner.onAckMessage(nil)
+	runner.onAckMessage(nil)
+	assert.Equal(t, 1, eventCount)
+}
+
 func TestOnMessage(t *testing.T) {
 	taskA := &Task{
 		Fn: func() {