|
| 1 | +// Copyright Amazon.com Inc. or its affiliates. All Rights Reserved. |
| 2 | +// |
| 3 | +// Licensed under the Apache License, Version 2.0 (the "License"). You may |
| 4 | +// not use this file except in compliance with the License. A copy of the |
| 5 | +// License is located at |
| 6 | +// |
| 7 | +// http://aws.amazon.com/apache2.0/ |
| 8 | +// |
| 9 | +// or in the "license" file accompanying this file. This file is distributed |
| 10 | +// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either |
| 11 | +// express or implied. See the License for the specific language governing |
| 12 | +// permissions and limitations under the License. |
| 13 | + |
| 14 | +package handler |
| 15 | + |
| 16 | +import ( |
| 17 | + "context" |
| 18 | + |
| 19 | + "github.com/aws/amazon-ecs-agent/agent/acs/model/ecsacs" |
| 20 | + "github.com/aws/amazon-ecs-agent/agent/wsclient" |
| 21 | + "github.com/aws/aws-sdk-go/aws" |
| 22 | + "github.com/cihub/seelog" |
| 23 | +) |
| 24 | + |
| 25 | +// heartbeatHandler handles heartbeat messages from ACS |
| 26 | +type heartbeatHandler struct { |
| 27 | + heartbeatMessageBuffer chan *ecsacs.HeartbeatMessage |
| 28 | + heartbeatAckMessageBuffer chan *ecsacs.HeartbeatAckRequest |
| 29 | + ctx context.Context |
| 30 | + cancel context.CancelFunc |
| 31 | + acsClient wsclient.ClientServer |
| 32 | +} |
| 33 | + |
| 34 | +// newHeartbeatHandler returns an instance of the heartbeatHandler struct |
| 35 | +func newHeartbeatHandler(ctx context.Context, |
| 36 | + acsClient wsclient.ClientServer) heartbeatHandler { |
| 37 | + |
| 38 | + // Create a cancelable context from the parent context |
| 39 | + derivedContext, cancel := context.WithCancel(ctx) |
| 40 | + return heartbeatHandler{ |
| 41 | + heartbeatMessageBuffer: make(chan *ecsacs.HeartbeatMessage), |
| 42 | + heartbeatAckMessageBuffer: make(chan *ecsacs.HeartbeatAckRequest), |
| 43 | + ctx: derivedContext, |
| 44 | + cancel: cancel, |
| 45 | + acsClient: acsClient, |
| 46 | + } |
| 47 | +} |
| 48 | + |
| 49 | +// handlerFunc returns a function to enqueue requests onto the buffer |
| 50 | +func (heartbeatHandler *heartbeatHandler) handlerFunc() func(message *ecsacs.HeartbeatMessage) { |
| 51 | + return func(message *ecsacs.HeartbeatMessage) { |
| 52 | + heartbeatHandler.heartbeatMessageBuffer <- message |
| 53 | + } |
| 54 | +} |
| 55 | + |
| 56 | +// start() invokes go routines to handle receive and respond to heartbeats |
| 57 | +func (heartbeatHandler *heartbeatHandler) start() { |
| 58 | + go heartbeatHandler.handleHeartbeatMessage() |
| 59 | + go heartbeatHandler.sendHeartbeatAck() |
| 60 | +} |
| 61 | + |
| 62 | +func (heartbeatHandler *heartbeatHandler) handleHeartbeatMessage() { |
| 63 | + for { |
| 64 | + select { |
| 65 | + case message := <-heartbeatHandler.heartbeatMessageBuffer: |
| 66 | + if err := heartbeatHandler.handleSingleHeartbeatMessage(message); err != nil { |
| 67 | + seelog.Warnf("Unable to handle heartbeat message [%s]: %s", message.String(), err) |
| 68 | + } |
| 69 | + case <-heartbeatHandler.ctx.Done(): |
| 70 | + return |
| 71 | + } |
| 72 | + } |
| 73 | +} |
| 74 | + |
| 75 | +func (heartbeatHandler *heartbeatHandler) handleSingleHeartbeatMessage(message *ecsacs.HeartbeatMessage) error { |
| 76 | + // Agent currently has no other action hooked to heartbeat messages, except simple ack |
| 77 | + go func() { |
| 78 | + response := &ecsacs.HeartbeatAckRequest{ |
| 79 | + MessageId: message.MessageId, |
| 80 | + } |
| 81 | + heartbeatHandler.heartbeatAckMessageBuffer <- response |
| 82 | + }() |
| 83 | + return nil |
| 84 | +} |
| 85 | + |
| 86 | +func (heartbeatHandler *heartbeatHandler) sendHeartbeatAck() { |
| 87 | + for { |
| 88 | + select { |
| 89 | + case ack := <-heartbeatHandler.heartbeatAckMessageBuffer: |
| 90 | + heartbeatHandler.sendSingleHeartbeatAck(ack) |
| 91 | + case <-heartbeatHandler.ctx.Done(): |
| 92 | + return |
| 93 | + } |
| 94 | + } |
| 95 | +} |
| 96 | + |
| 97 | +func (heartbeatHandler *heartbeatHandler) sendSingleHeartbeatAck(ack *ecsacs.HeartbeatAckRequest) { |
| 98 | + err := heartbeatHandler.acsClient.MakeRequest(ack) |
| 99 | + if err != nil { |
| 100 | + seelog.Warnf("Error acknowledging server heartbeat, message id: %s, error: %s", aws.StringValue(ack.MessageId), err) |
| 101 | + } |
| 102 | +} |
| 103 | + |
| 104 | +// stop() cancels the context being used by this handler, which stops the go routines started by 'start()' |
| 105 | +func (heartbeatHandler *heartbeatHandler) stop() { |
| 106 | + heartbeatHandler.cancel() |
| 107 | +} |
| 108 | + |
| 109 | +// clearAcks drains the ack request channel |
| 110 | +func (heartbeatHandler *heartbeatHandler) clearAcks() { |
| 111 | + for { |
| 112 | + select { |
| 113 | + case <-heartbeatHandler.heartbeatAckMessageBuffer: |
| 114 | + default: |
| 115 | + return |
| 116 | + } |
| 117 | + } |
| 118 | +} |
0 commit comments