-
Notifications
You must be signed in to change notification settings - Fork 619
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
ECS agent to acknowledge server heartbeat messages
- Loading branch information
1 parent
683c258
commit dab3be2
Showing
7 changed files
with
276 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,118 @@ | ||
| // Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"). You may | ||
| // not use this file except in compliance with the License. A copy of the | ||
| // License is located at | ||
| // | ||
| // http://aws.amazon.com/apache2.0/ | ||
| // | ||
| // or in the "license" file accompanying this file. This file is distributed | ||
| // on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either | ||
| // express or implied. See the License for the specific language governing | ||
| // permissions and limitations under the License. | ||
|
|
||
| package handler | ||
|
|
||
| import ( | ||
| "context" | ||
|
|
||
| "github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs" | ||
| "github.com/aws/amazon-ecs-agent/agent/wsclient" | ||
| "github.com/aws/aws-sdk-go/aws" | ||
| "github.com/cihub/seelog" | ||
| ) | ||
|
|
||
| // heartbeatHandler handles heartbeat messages from ACS | ||
| type heartbeatHandler struct { | ||
| heartbeatMessageBuffer chan *ecsacs.HeartbeatMessage | ||
| heartbeatAckMessageBuffer chan *ecsacs.HeartbeatAckRequest | ||
| ctx context.Context | ||
| cancel context.CancelFunc | ||
| acsClient wsclient.ClientServer | ||
| } | ||
|
|
||
| // newHeartbeatHandler returns an instance of the heartbeatHandler struct | ||
| func newHeartbeatHandler(ctx context.Context, | ||
| acsClient wsclient.ClientServer) heartbeatHandler { | ||
|
|
||
| // Create a cancelable context from the parent context | ||
| derivedContext, cancel := context.WithCancel(ctx) | ||
| return heartbeatHandler{ | ||
| heartbeatMessageBuffer: make(chan *ecsacs.HeartbeatMessage), | ||
| heartbeatAckMessageBuffer: make(chan *ecsacs.HeartbeatAckRequest), | ||
| ctx: derivedContext, | ||
| cancel: cancel, | ||
| acsClient: acsClient, | ||
| } | ||
| } | ||
|
|
||
| // handlerFunc returns a function to enqueue requests onto the buffer | ||
| func (heartbeatHandler *heartbeatHandler) handlerFunc() func(message *ecsacs.HeartbeatMessage) { | ||
| return func(message *ecsacs.HeartbeatMessage) { | ||
| heartbeatHandler.heartbeatMessageBuffer <- message | ||
| } | ||
| } | ||
|
|
||
| // start() invokes go routines to handle receive and respond to heartbeats | ||
| func (heartbeatHandler *heartbeatHandler) start() { | ||
| go heartbeatHandler.handleHeartbeatMessage() | ||
| go heartbeatHandler.sendHeartbeatAck() | ||
| } | ||
|
|
||
| func (heartbeatHandler *heartbeatHandler) handleHeartbeatMessage() { | ||
| for { | ||
| select { | ||
| case message := <-heartbeatHandler.heartbeatMessageBuffer: | ||
| if err := heartbeatHandler.handleSingleHeartbeatMessage(message); err != nil { | ||
| seelog.Warnf("Unable to handle heartbeat message [%s]: %s", message.String(), err) | ||
| } | ||
| case <-heartbeatHandler.ctx.Done(): | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (heartbeatHandler *heartbeatHandler) handleSingleHeartbeatMessage(message *ecsacs.HeartbeatMessage) error { | ||
| // Agent currently has no other action hooked to heartbeat messages, except simple ack | ||
| go func() { | ||
| response := &ecsacs.HeartbeatAckRequest{ | ||
| MessageId: message.MessageId, | ||
| } | ||
| heartbeatHandler.heartbeatAckMessageBuffer <- response | ||
| }() | ||
| return nil | ||
| } | ||
|
|
||
| func (heartbeatHandler *heartbeatHandler) sendHeartbeatAck() { | ||
| for { | ||
| select { | ||
| case ack := <-heartbeatHandler.heartbeatAckMessageBuffer: | ||
| heartbeatHandler.sendSingleHeartbeatAck(ack) | ||
| case <-heartbeatHandler.ctx.Done(): | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (heartbeatHandler *heartbeatHandler) sendSingleHeartbeatAck(ack *ecsacs.HeartbeatAckRequest) { | ||
| err := heartbeatHandler.acsClient.MakeRequest(ack) | ||
| if err != nil { | ||
| seelog.Warnf("Error acknowledging server heartbeat, message id: %s, error: %s", aws.StringValue(ack.MessageId), err) | ||
| } | ||
| } | ||
|
|
||
| // stop() cancels the context being used by this handler, which stops the go routines started by 'start()' | ||
| func (heartbeatHandler *heartbeatHandler) stop() { | ||
| heartbeatHandler.cancel() | ||
| } | ||
|
|
||
| // clearAcks drains the ack request channel | ||
| func (heartbeatHandler *heartbeatHandler) clearAcks() { | ||
| for { | ||
| select { | ||
| case <-heartbeatHandler.heartbeatAckMessageBuffer: | ||
| default: | ||
| return | ||
| } | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,104 @@ | ||
| // +build unit | ||
|
|
||
| // Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"). You may | ||
| // not use this file except in compliance with the License. A copy of the | ||
| // License is located at | ||
| // | ||
| // http://aws.amazon.com/apache2.0/ | ||
| // | ||
| // or in the "license" file accompanying this file. This file is distributed | ||
| // on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either | ||
| // express or implied. See the License for the specific language governing | ||
| // permissions and limitations under the License. | ||
|
|
||
| package handler | ||
|
|
||
| import ( | ||
| "context" | ||
| "testing" | ||
|
|
||
| "github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs" | ||
| mock_wsclient "github.com/aws/amazon-ecs-agent/agent/wsclient/mock" | ||
| "github.com/aws/aws-sdk-go/aws" | ||
| "github.com/golang/mock/gomock" | ||
| "github.com/stretchr/testify/require" | ||
| ) | ||
|
|
||
| const ( | ||
| heartbeatMessageId = "heartbeatMessageId" | ||
| ) | ||
|
|
||
| func TestAckHeartbeatMessage(t *testing.T) { | ||
| heartbeatReceived := &ecsacs.HeartbeatMessage{ | ||
| MessageId: aws.String(heartbeatMessageId), | ||
| Healthy: aws.Bool(true), | ||
| } | ||
|
|
||
| heartbeatAckExpected := &ecsacs.HeartbeatAckRequest{ | ||
| MessageId: aws.String(heartbeatMessageId), | ||
| } | ||
|
|
||
| validateHeartbeatAck(t, heartbeatReceived, heartbeatAckExpected) | ||
| } | ||
|
|
||
| func TestAckHeartbeatMessageNotHealthy(t *testing.T) { | ||
| heartbeatReceived := &ecsacs.HeartbeatMessage{ | ||
| MessageId: aws.String(heartbeatMessageId), | ||
| // ECS Agent currently ignores this field so we expect no behavior change | ||
| Healthy: aws.Bool(false), | ||
| } | ||
|
|
||
| heartbeatAckExpected := &ecsacs.HeartbeatAckRequest{ | ||
| MessageId: aws.String(heartbeatMessageId), | ||
| } | ||
|
|
||
| validateHeartbeatAck(t, heartbeatReceived, heartbeatAckExpected) | ||
| } | ||
|
|
||
| func TestAckHeartbeatMessageWithoutMessageId(t *testing.T) { | ||
| heartbeatReceived := &ecsacs.HeartbeatMessage{ | ||
| Healthy: aws.Bool(true), | ||
| } | ||
|
|
||
| heartbeatAckExpected := &ecsacs.HeartbeatAckRequest{ | ||
| MessageId: nil, | ||
| } | ||
|
|
||
| validateHeartbeatAck(t, heartbeatReceived, heartbeatAckExpected) | ||
| } | ||
|
|
||
| func TestAckHeartbeatMessageEmpty(t *testing.T) { | ||
| heartbeatReceived := &ecsacs.HeartbeatMessage{} | ||
|
|
||
| heartbeatAckExpected := &ecsacs.HeartbeatAckRequest{ | ||
| MessageId: nil, | ||
| } | ||
|
|
||
| validateHeartbeatAck(t, heartbeatReceived, heartbeatAckExpected) | ||
| } | ||
|
|
||
| func validateHeartbeatAck(t *testing.T, heartbeatReceived *ecsacs.HeartbeatMessage, heartbeatAckExpected *ecsacs.HeartbeatAckRequest) { | ||
| ctrl := gomock.NewController(t) | ||
| defer ctrl.Finish() | ||
|
|
||
| ctx, cancel := context.WithCancel(context.Background()) | ||
| var heartbeatAckSent *ecsacs.HeartbeatAckRequest | ||
|
|
||
| mockWsClient := mock_wsclient.NewMockClientServer(ctrl) | ||
| mockWsClient.EXPECT().MakeRequest(gomock.Any()).Do(func(message *ecsacs.HeartbeatAckRequest) { | ||
| heartbeatAckSent = message | ||
| cancel() | ||
| }).Times(1) | ||
|
|
||
| handler := newHeartbeatHandler(ctx, mockWsClient) | ||
| go handler.sendHeartbeatAck() | ||
|
|
||
| handler.handleSingleHeartbeatMessage(heartbeatReceived) | ||
|
|
||
| // wait till we get an ack from heartbeatAckMessageBuffer | ||
| <-ctx.Done() | ||
|
|
||
| require.Equal(t, heartbeatAckExpected, heartbeatAckSent) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
Oops, something went wrong.