-
Notifications
You must be signed in to change notification settings - Fork 612
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
76bcf85
commit 1279313
Showing
7 changed files
with
268 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,119 @@ | ||
| // 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]: %v", message.String(), err) | ||
| } | ||
| case <-heartbeatHandler.ctx.Done(): | ||
| return | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func (heartbeatHandler *heartbeatHandler) handleSingleHeartbeatMessage(message *ecsacs.HeartbeatMessage) error { | ||
| seelog.Tracef("Received server heartbeat message: %s", message.MessageId) | ||
| 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: %v", aws.StringValue(ack.MessageId), err) | ||
| } | ||
| seelog.Tracef("Acknowledging server heartbeat message: %s", ack.MessageId) | ||
| } | ||
|
|
||
| // 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,95 @@ | ||
| // +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" | ||
| "reflect" | ||
| "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" | ||
| ) | ||
|
|
||
| const ( | ||
| heartbeatMessageId = "heartbeatMessageId" | ||
| heartbeatHealthy = true | ||
| ) | ||
|
|
||
| func TestAckHeartbeatMessage(t *testing.T) { | ||
| heartbeatReceived := &ecsacs.HeartbeatMessage{ | ||
| MessageId: aws.String(heartbeatMessageId), | ||
| Healthy: aws.Bool(heartbeatHealthy), | ||
| } | ||
|
|
||
| heartbeatAckExpected := &ecsacs.HeartbeatAckRequest{ | ||
| MessageId: aws.String(heartbeatMessageId), | ||
| } | ||
|
|
||
| validateHeartbeatAck(t, heartbeatReceived, heartbeatAckExpected) | ||
| } | ||
|
|
||
| func TestAckHeartbeatMessageWithoutMessageId(t *testing.T) { | ||
| heartbeatReceived := &ecsacs.HeartbeatMessage{ | ||
| Healthy: aws.Bool(heartbeatHealthy), | ||
| } | ||
|
|
||
| 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 | ||
| select { | ||
| case <-ctx.Done(): | ||
| } | ||
|
|
||
| if !reflect.DeepEqual(heartbeatAckExpected, heartbeatAckSent) { | ||
| t.Errorf("Message mismatch between expected and sent ack, expected: %v, sent: %v", 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.