Merge 76ac8c4 into 3f934ba

agent/agent.go

@@ -63,18 +63,17 @@ type (
 		appDieChan         chan bool           // app die channel
 		chDie              chan struct{}       // wait for close
 		chSend             chan pendingMessage // push message queue
-		chStopHeartbeat    chan struct{}       // stop heartbeats
-		chStopWrite        chan struct{}       // stop writing messages
+		closeMutex         sync.Mutex
 		conn               net.Conn            // low-level conn fd
 		decoder            codec.PacketDecoder // binary decoder
 		encoder            codec.PacketEncoder // binary encoder
 		heartbeatTimeout   time.Duration
-		lastAt             int64                // last heartbeat unix time stamp
-		messagesBufferSize int                  // size of the pending messages buffer
-		serializer         serialize.Serializer // message serializer
-		state              int32                // current agent state
+		lastAt             int64 // last heartbeat unix time stamp
 		messageEncoder     message.Encoder
+		messagesBufferSize int // size of the pending messages buffer
 		metricsReporters   []metrics.Reporter
+		serializer         serialize.Serializer // message serializer
+		state              int32                // current agent state
 	}
 
 	pendingMessage struct {
@@ -109,8 +108,6 @@ func NewAgent(
 		chDie:              make(chan struct{}),
 		chSend:             make(chan pendingMessage, messagesBufferSize),
 		messagesBufferSize: messagesBufferSize,
-		chStopHeartbeat:    make(chan struct{}),
-		chStopWrite:        make(chan struct{}),
 		conn:               conn,
 		decoder:            packetDecoder,
 		encoder:            packetEncoder,
@@ -193,6 +190,8 @@ func (a *Agent) ResponseMID(ctx context.Context, mid uint, v interface{}, isErro
 // Close closes the agent, cleans inner state and closes low-level connection.
 // Any blocked Read or Write operations will be unblocked and return errors.
 func (a *Agent) Close() error {
+	a.closeMutex.Lock()
+	defer a.closeMutex.Unlock()
 	if a.GetStatus() == constants.StatusClosed {
 		return constants.ErrCloseClosedSession
 	}
@@ -206,8 +205,6 @@ func (a *Agent) Close() error {
 	case <-a.chDie:
 		// expect
 	default:
-		close(a.chStopWrite)
-		close(a.chStopHeartbeat)
 		close(a.chDie)
 		onSessionClosed(a.Session)
 	}
@@ -307,8 +304,6 @@ func (a *Agent) heartbeat() {
 			}
 		case <-a.chDie:
 			return
-		case <-a.chStopHeartbeat:
-			return
 		}
 	}
 }
@@ -401,7 +396,7 @@ func (a *Agent) write() {
 			if data.typ == message.Response {
 				metrics.ReportTimingFromCtx(data.ctx, a.metricsReporters, handlerType, m.Err)
 			}
-		case <-a.chStopWrite:
+		case <-a.chDie:
 			return
 		}
 	}

agent/agent_test.go

@@ -97,8 +97,6 @@ func TestNewAgent(t *testing.T) {
 	assert.NotNil(t, ag)
 	assert.IsType(t, make(chan struct{}), ag.chDie)
 	assert.IsType(t, make(chan pendingMessage), ag.chSend)
-	assert.IsType(t, make(chan struct{}), ag.chStopHeartbeat)
-	assert.IsType(t, make(chan struct{}), ag.chStopWrite)
 	assert.Equal(t, dieChan, ag.appDieChan)
 	assert.Equal(t, 10, ag.messagesBufferSize)
 	assert.Equal(t, mockConn, ag.conn)
@@ -425,16 +423,10 @@ func TestAgentClose(t *testing.T) {
 	assert.NoError(t, err)
 
 	// validate channels are closed
-	stopWrite := false
-	stopHeartbeat := false
 	die := false
 	go func() {
 		for {
 			select {
-			case <-ag.chStopWrite:
-				stopWrite = true
-			case <-ag.chStopHeartbeat:
-				stopHeartbeat = true
 			case <-ag.chDie:
 				die = true
 			}
@@ -448,8 +440,8 @@ func TestAgentClose(t *testing.T) {
 	assert.Equal(t, ag.state, constants.StatusClosed)
 	assert.True(t, expected)
 	helpers.ShouldEventuallyReturn(
-		t, func() bool { return stopWrite && stopHeartbeat && die },
-		true, 50*time.Millisecond, 500*time.Millisecond)
+		t, func() bool { return die },
+		true, 20*time.Millisecond, 500*time.Millisecond)
 }
 
 func TestAgentRemoteAddr(t *testing.T) {