func TestDeleteImageOtherRemoveImageErrors(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := NewMockDockerClient(ctrl)
imageManager := &dockerImageManager{client: client, state: dockerstate.NewDockerTaskEngineState()}
imageManager.SetSaver(statemanager.NewNoopStateManager())
container := &api.Container{
Name: "testContainer",
Image: "testContainerImage",
}
imageInspected := &docker.Image{
ID: "sha256:qwerty",
}
client.EXPECT().InspectImage(container.Image).Return(imageInspected, nil).AnyTimes()
err := imageManager.AddContainerReferenceToImageState(container)
if err != nil {
t.Error("Error in adding container to an existing image state")
}
imageState, _ := imageManager.getImageState(imageInspected.ID)
client.EXPECT().RemoveImage(container.Image, removeImageTimeout).Return(errors.New("container for this image exists"))
imageManager.deleteImage(container.Image, imageState)
if len(imageState.Image.Names) == 0 {
t.Error("Incorrectly removed Image name from image state")
}
if len(imageManager.getAllImageStates()) == 0 {
t.Error("Incorrecting removed image state from image manager before deletion")
}
}
// TestHandlePayloadMessageWithNoMessageId tests that agent doesn't ack payload messages
// that do not contain message ids
func TestHandlePayloadMessageWithNoMessageId(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
ecsClient := mock_api.NewMockECSClient(ctrl)
stateManager := statemanager.NewNoopStateManager()
credentialsManager := credentials.NewManager()
ctx := context.Background()
buffer := newPayloadRequestHandler(ctx, taskEngine, ecsClient, clusterName, containerInstanceArn, nil, stateManager, refreshCredentialsHandler{}, credentialsManager)
// test adding a payload message without the MessageId field
payloadMessage := &ecsacs.PayloadMessage{
Tasks: []*ecsacs.Task{
&ecsacs.Task{
Arn: aws.String("t1"),
},
},
}
err := buffer.handleSingleMessage(payloadMessage)
if err == nil {
t.Error("Expected error while adding a task with no message id")
}
// test adding a payload message with blank MessageId
payloadMessage.MessageId = aws.String("")
err = buffer.handleSingleMessage(payloadMessage)
if err == nil {
t.Error("Expected error while adding a task with no message id")
}
}
func TestRemoveUnusedImagesNoImages(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := NewMockDockerClient(ctrl)
imageManager := &dockerImageManager{client: client, state: dockerstate.NewDockerTaskEngineState()}
imageManager.SetSaver(statemanager.NewNoopStateManager())
imageManager.removeUnusedImages()
}
func TestDeleteImageIDNull(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := NewMockDockerClient(ctrl)
imageManager := &dockerImageManager{client: client, state: dockerstate.NewDockerTaskEngineState()}
imageManager.SetSaver(statemanager.NewNoopStateManager())
imageManager.deleteImage("", nil)
}
// TestAddPayloadTaskAddsNonStoppedTasksAfterStoppedTasks tests if tasks with desired status
// 'RUNNING' are added after tasks with desired status 'STOPPED'
func TestAddPayloadTaskAddsNonStoppedTasksAfterStoppedTasks(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
ecsClient := mock_api.NewMockECSClient(ctrl)
taskEngine := engine.NewMockTaskEngine(ctrl)
credentialsManager := credentials.NewManager()
var tasksAddedToEngine []*api.Task
taskEngine.EXPECT().AddTask(gomock.Any()).Do(func(task *api.Task) {
tasksAddedToEngine = append(tasksAddedToEngine, task)
}).Times(2)
stoppedTaskArn := "stoppedTask"
runningTaskArn := "runningTask"
payloadMessage := &ecsacs.PayloadMessage{
Tasks: []*ecsacs.Task{
&ecsacs.Task{
Arn: aws.String(runningTaskArn),
DesiredStatus: aws.String("RUNNING"),
},
&ecsacs.Task{
Arn: aws.String(stoppedTaskArn),
DesiredStatus: aws.String("STOPPED"),
},
},
MessageId: aws.String(payloadMessageId),
}
ctx := context.Background()
stateManager := statemanager.NewNoopStateManager()
buffer := newPayloadRequestHandler(ctx, taskEngine, ecsClient, clusterName, containerInstanceArn, nil, stateManager, refreshCredentialsHandler{}, credentialsManager)
_, ok := buffer.addPayloadTasks(payloadMessage)
if !ok {
t.Error("addPayloadTasks returned false")
}
if len(tasksAddedToEngine) != 2 {
t.Errorf("Incorrect number of tasks added to the engine. Expected: %d, got: %d", 2, len(tasksAddedToEngine))
}
// Verify if stopped task is added before running task
firstTaskAdded := tasksAddedToEngine[0]
if firstTaskAdded.Arn != stoppedTaskArn {
t.Errorf("Expected first task arn: %s, got: %s", stoppedTaskArn, firstTaskAdded.Arn)
}
if firstTaskAdded.DesiredStatus != api.TaskStopped {
t.Errorf("Expected first task state be be: %s , got: %s", "STOPPED", firstTaskAdded.DesiredStatus.String())
}
secondTaskAdded := tasksAddedToEngine[1]
if secondTaskAdded.Arn != runningTaskArn {
t.Errorf("Expected second task arn: %s, got: %s", runningTaskArn, secondTaskAdded.Arn)
}
if secondTaskAdded.DesiredStatus != api.TaskRunning {
t.Errorf("Expected second task state be be: %s , got: %s", "RUNNNING", secondTaskAdded.DesiredStatus.String())
}
}
func TestHandlerReconnects(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
ecsclient := mock_api.NewMockECSClient(ctrl)
statemanager := statemanager.NewNoopStateManager()
closeWS := make(chan bool)
server, serverIn, requests, errs, err := startMockAcsServer(t, closeWS)
if err != nil {
t.Fatal(err)
}
go func() {
for {
select {
case <-requests:
case <-errs:
}
}
}()
ecsclient.EXPECT().DiscoverPollEndpoint("myArn").Return(server.URL, nil).Times(10)
taskEngine.EXPECT().Version().Return("Docker: 1.5.0", nil).AnyTimes()
ctx, cancel := context.WithCancel(context.Background())
ended := make(chan bool, 1)
go func() {
handler.StartSession(ctx, handler.StartSessionArguments{
ContainerInstanceArn: "myArn",
CredentialProvider: credentials.AnonymousCredentials,
Config: &config.Config{Cluster: "someCluster"},
TaskEngine: taskEngine,
ECSClient: ecsclient,
StateManager: statemanager,
AcceptInvalidCert: true,
})
// This should never return
ended <- true
}()
start := time.Now()
for i := 0; i < 10; i++ {
serverIn <- `{"type":"HeartbeatMessage","message":{"healthy":true}}`
closeWS <- true
}
if time.Since(start) > 2*time.Second {
t.Error("Test took longer than expected; backoff should not have occured for EOF")
}
select {
case <-ended:
t.Fatal("Should not have stopped session")
default:
}
cancel()
<-ended
}
// TestHandlerReconnectsCorrectlySetsSendCredentialsURLParameter tests if
// the 'sendCredentials' URL parameter is set correctly for successive
// invocations of startACSSession
func TestHandlerReconnectsCorrectlySetsSendCredentialsURLParameter(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
ecsClient := mock_api.NewMockECSClient(ctrl)
statemanager := statemanager.NewNoopStateManager()
ctx, cancel := context.WithCancel(context.Background())
mockWsClient := mock_wsclient.NewMockClientServer(ctrl)
args := StartSessionArguments{
ContainerInstanceArn: "myArn",
CredentialProvider: credentials.AnonymousCredentials,
Config: &config.Config{Cluster: "someCluster"},
TaskEngine: taskEngine,
ECSClient: ecsClient,
StateManager: statemanager,
AcceptInvalidCert: true,
_heartbeatTimeout: 20 * time.Millisecond,
_heartbeatJitter: 10 * time.Millisecond,
}
session := newSessionResources(args)
mockWsClient.EXPECT().SetAnyRequestHandler(gomock.Any()).AnyTimes()
mockWsClient.EXPECT().AddRequestHandler(gomock.Any()).AnyTimes()
mockWsClient.EXPECT().Close().Return(nil).AnyTimes()
mockWsClient.EXPECT().Serve().Return(io.EOF).AnyTimes()
gomock.InOrder(
// When the websocket client connects to ACS for the first
// time, 'sendCredentials' should be set to true
mockWsClient.EXPECT().Connect().Do(func() {
validateSendCredentialsInSession(t, session, "true")
}).Return(nil),
// For all subsequent connections to ACS, 'sendCredentials'
// should be set to false
mockWsClient.EXPECT().Connect().Do(func() {
validateSendCredentialsInSession(t, session, "false")
}).Return(nil).AnyTimes(),
)
backoff := utils.NewSimpleBackoff(connectionBackoffMin, connectionBackoffMax, connectionBackoffJitter, connectionBackoffMultiplier)
timer := newDisconnectionTimer(mockWsClient, args.time(), args.heartbeatTimeout(), args.heartbeatJitter())
defer timer.Stop()
go func() {
for i := 0; i < 10; i++ {
startACSSession(ctx, mockWsClient, timer, args, backoff, session)
}
cancel()
}()
// Wait for context to be cancelled
select {
case <-ctx.Done():
}
}
// TestHandlePayloadMessageAckedWhenTaskAdded tests if the handler generates an ack
// after processing a payload message.
func TestHandlePayloadMessageAckedWhenTaskAdded(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
ecsClient := mock_api.NewMockECSClient(ctrl)
stateManager := statemanager.NewNoopStateManager()
credentialsManager := credentials.NewManager()
ctx, cancel := context.WithCancel(context.Background())
taskEngine := engine.NewMockTaskEngine(ctrl)
var addedTask *api.Task
taskEngine.EXPECT().AddTask(gomock.Any()).Do(func(task *api.Task) {
addedTask = task
}).Times(1)
var ackRequested *ecsacs.AckRequest
mockWsClient := mock_wsclient.NewMockClientServer(ctrl)
mockWsClient.EXPECT().MakeRequest(gomock.Any()).Do(func(ackRequest *ecsacs.AckRequest) {
ackRequested = ackRequest
cancel()
}).Times(1)
buffer := newPayloadRequestHandler(ctx, taskEngine, ecsClient, clusterName, containerInstanceArn, mockWsClient, stateManager, refreshCredentialsHandler{}, credentialsManager)
go buffer.start()
// Send a payload message
payloadMessage := &ecsacs.PayloadMessage{
Tasks: []*ecsacs.Task{
&ecsacs.Task{
Arn: aws.String("t1"),
},
},
MessageId: aws.String(payloadMessageId),
}
err := buffer.handleSingleMessage(payloadMessage)
if err != nil {
t.Errorf("Error handling payload message: %v", err)
}
// Wait till we get an ack from the ackBuffer
select {
case <-ctx.Done():
}
// Verify the message id acked
if aws.StringValue(ackRequested.MessageId) != payloadMessageId {
t.Errorf("Message Id mismatch. Expected: %s, got: %s", payloadMessageId, aws.StringValue(ackRequested.MessageId))
}
// Verify if task added == expected task
expectedTask := &api.Task{
Arn: "t1",
}
if !reflect.DeepEqual(addedTask, expectedTask) {
t.Errorf("Mismatch between expected and added tasks, expected: %v, added: %v", expectedTask, addedTask)
}
}
func TestRemoveLeastRecentlyUsedImageNoImage(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := NewMockDockerClient(ctrl)
imageManager := &dockerImageManager{client: client, state: dockerstate.NewDockerTaskEngineState()}
imageManager.SetSaver(statemanager.NewNoopStateManager())
err := imageManager.removeLeastRecentlyUsedImage()
if err == nil {
t.Error("Expected Error for no LRU image to remove")
}
}
func TestImageCleanupHappyPath(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := NewMockDockerClient(ctrl)
imageManager := &dockerImageManager{
client: client,
state: dockerstate.NewDockerTaskEngineState(),
minimumAgeBeforeDeletion: 1 * time.Millisecond,
numImagesToDelete: config.DefaultNumImagesToDeletePerCycle,
imageCleanupTimeInterval: config.DefaultImageCleanupTimeInterval,
}
imageManager.SetSaver(statemanager.NewNoopStateManager())
container := &api.Container{
Name: "testContainer",
Image: "testContainerImage",
}
imageInspected := &docker.Image{
ID: "sha256:qwerty",
}
client.EXPECT().InspectImage(container.Image).Return(imageInspected, nil).AnyTimes()
err := imageManager.AddContainerReferenceToImageState(container)
if err != nil {
t.Error("Error in adding container to an existing image state")
}
err = imageManager.RemoveContainerReferenceFromImageState(container)
if err != nil {
t.Error("Error removing container reference from image state")
}
imageState, _ := imageManager.getImageState(imageInspected.ID)
imageState.PulledAt = time.Now().AddDate(0, -2, 0)
imageState.LastUsedAt = time.Now().AddDate(0, -2, 0)
imageState.AddImageName("anotherImage")
client.EXPECT().RemoveImage(container.Image, removeImageTimeout).Return(nil)
client.EXPECT().RemoveImage("anotherImage", removeImageTimeout).Return(nil)
parent := context.Background()
ctx, cancel := context.WithCancel(parent)
go imageManager.performPeriodicImageCleanup(ctx, 2*time.Millisecond)
time.Sleep(1 * time.Second)
cancel()
if len(imageState.Image.Names) != 0 {
t.Error("Error removing image name from state after the image is removed")
}
if len(imageManager.imageStates) != 0 {
t.Error("Error removing image state after the image is removed")
}
}
// TestHandlerReconnectsOnServeErrors tests if the handler retries to
// to establish the session with ACS when ClientServer.Connect() returns errors
func TestHandlerReconnectsOnServeErrors(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
taskEngine.EXPECT().Version().Return("Docker: 1.5.0", nil).AnyTimes()
ecsClient := mock_api.NewMockECSClient(ctrl)
ecsClient.EXPECT().DiscoverPollEndpoint(gomock.Any()).Return(acsURL, nil).AnyTimes()
statemanager := statemanager.NewNoopStateManager()
ctx, cancel := context.WithCancel(context.Background())
mockWsClient := mock_wsclient.NewMockClientServer(ctrl)
mockWsClient.EXPECT().SetAnyRequestHandler(gomock.Any()).AnyTimes()
mockWsClient.EXPECT().AddRequestHandler(gomock.Any()).AnyTimes()
mockWsClient.EXPECT().Connect().Return(nil).AnyTimes()
mockWsClient.EXPECT().Close().Return(nil).AnyTimes()
gomock.InOrder(
// Serve fails 10 times
mockWsClient.EXPECT().Serve().Return(io.EOF).Times(10),
// Cancel trying to Serve ACS requests on the 11th attempt
// Failure to retry on Serve() errors should cause the
// test to time out as the context is never cancelled
mockWsClient.EXPECT().Serve().Do(func() {
cancel()
}).Return(io.EOF),
)
session := &mockSession{mockWsClient}
args := StartSessionArguments{
ContainerInstanceArn: "myArn",
CredentialProvider: credentials.AnonymousCredentials,
Config: &config.Config{Cluster: "someCluster"},
TaskEngine: taskEngine,
ECSClient: ecsClient,
StateManager: statemanager,
AcceptInvalidCert: true,
_heartbeatTimeout: 20 * time.Millisecond,
_heartbeatJitter: 10 * time.Millisecond,
}
backoff := utils.NewSimpleBackoff(connectionBackoffMin, connectionBackoffMax, connectionBackoffJitter, connectionBackoffMultiplier)
go func() {
startSession(ctx, args, backoff, session)
}()
// Wait for context to be cancelled
select {
case <-ctx.Done():
}
}
func TestImageCleanupCannotRemoveImage(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := NewMockDockerClient(ctrl)
imageManager := &dockerImageManager{
client: client,
state: dockerstate.NewDockerTaskEngineState(),
minimumAgeBeforeDeletion: config.DefaultImageDeletionAge,
numImagesToDelete: config.DefaultNumImagesToDeletePerCycle,
imageCleanupTimeInterval: config.DefaultImageCleanupTimeInterval,
}
imageManager.SetSaver(statemanager.NewNoopStateManager())
container := &api.Container{
Name: "testContainer",
Image: "testContainerImage",
}
sourceImage := &image.Image{
ImageID: "sha256:qwerty",
}
sourceImage.Names = append(sourceImage.Names, container.Image)
imageInspected := &docker.Image{
ID: "sha256:qwerty",
}
client.EXPECT().InspectImage(container.Image).Return(imageInspected, nil).AnyTimes()
err := imageManager.AddContainerReferenceToImageState(container)
if err != nil {
t.Error("Error in adding container to an existing image state")
}
err = imageManager.RemoveContainerReferenceFromImageState(container)
if err != nil {
t.Error("Error removing container reference from image state")
}
imageState, _ := imageManager.getImageState(imageInspected.ID)
imageState.PulledAt = time.Now().AddDate(0, -2, 0)
imageState.LastUsedAt = time.Now().AddDate(0, -2, 0)
client.EXPECT().RemoveImage(container.Image, removeImageTimeout).Return(errors.New("error removing image")).AnyTimes()
imageManager.removeUnusedImages()
if len(imageState.Image.Names) == 0 {
t.Error("Error: image name should not be removed")
}
if len(imageManager.imageStates) == 0 {
t.Error("Error: image state should not be removed")
}
}
func initializeStateManager(cfg *config.Config, taskEngine engine.TaskEngine, cluster, containerInstanceArn, savedInstanceID *string, sequenceNumber *utilatomic.IncreasingInt64) (statemanager.StateManager, error) {
if !cfg.Checkpoint {
return statemanager.NewNoopStateManager(), nil
}
stateManager, err := statemanager.NewStateManager(cfg,
statemanager.AddSaveable("TaskEngine", taskEngine),
statemanager.AddSaveable("ContainerInstanceArn", containerInstanceArn),
statemanager.AddSaveable("Cluster", cluster),
statemanager.AddSaveable("EC2InstanceID", savedInstanceID),
statemanager.AddSaveable("ACSSeqNum", sequenceNumber),
)
if err != nil {
return nil, err
}
return stateManager, nil
}
// NewDockerTaskEngine returns a created, but uninitialized, DockerTaskEngine.
// The distinction between created and initialized is that when created it may
// be serialized/deserialized, but it will not communicate with docker until it
// is also initialized.
func NewDockerTaskEngine(cfg *config.Config) *DockerTaskEngine {
dockerTaskEngine := &DockerTaskEngine{
client: nil,
saver: statemanager.NewNoopStateManager(),
state: dockerstate.NewDockerTaskEngineState(),
managedTasks: make(map[string]*managedTask),
taskStopGroup: utilsync.NewSequentialWaitGroup(),
containerEvents: make(chan api.ContainerStateChange),
taskEvents: make(chan api.TaskStateChange),
}
dockerauth.SetConfig(cfg)
return dockerTaskEngine
}
func TestUndownloadedUpdate(t *testing.T) {
u, ctrl, cfg, _, mockacs, _ := mocks(t, nil)
defer ctrl.Finish()
mockacs.EXPECT().MakeRequest(&nackRequestMatcher{&ecsacs.NackRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("mid").(*string),
}})
u.performUpdateHandler(statemanager.NewNoopStateManager(), engine.NewTaskEngine(cfg))(&ecsacs.PerformUpdateMessage{
ClusterArn: ptr("cluster").(*string),
ContainerInstanceArn: ptr("containerInstance").(*string),
MessageId: ptr("mid").(*string),
})
}
// TestConnectionIsClosedOnIdle tests if the connection to ACS is closed
// when the channel is idle
func TestConnectionIsClosedOnIdle(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
taskEngine.EXPECT().Version().Return("Docker: 1.5.0", nil).AnyTimes()
ecsClient := mock_api.NewMockECSClient(ctrl)
statemanager := statemanager.NewNoopStateManager()
mockWsClient := mock_wsclient.NewMockClientServer(ctrl)
mockWsClient.EXPECT().SetAnyRequestHandler(gomock.Any()).Do(func(v interface{}) {}).AnyTimes()
mockWsClient.EXPECT().AddRequestHandler(gomock.Any()).Do(func(v interface{}) {}).AnyTimes()
mockWsClient.EXPECT().Connect().Return(nil)
mockWsClient.EXPECT().Serve().Do(func() {
// Pretend as if the maximum heartbeatTimeout duration has
// been breached while Serving requests
time.Sleep(30 * time.Millisecond)
}).Return(io.EOF)
connectionClosed := make(chan bool)
mockWsClient.EXPECT().Close().Do(func() {
// Record connection closed
connectionClosed <- true
}).Return(nil)
ctx := context.Background()
backoff := utils.NewSimpleBackoff(connectionBackoffMin, connectionBackoffMax, connectionBackoffJitter, connectionBackoffMultiplier)
args := StartSessionArguments{
ContainerInstanceArn: "myArn",
CredentialProvider: credentials.AnonymousCredentials,
Config: &config.Config{Cluster: "someCluster"},
TaskEngine: taskEngine,
ECSClient: ecsClient,
StateManager: statemanager,
AcceptInvalidCert: true,
_heartbeatTimeout: 20 * time.Millisecond,
_heartbeatJitter: 10 * time.Millisecond,
}
go func() {
timer := newDisconnectionTimer(mockWsClient, args.time(), args.heartbeatTimeout(), args.heartbeatJitter())
defer timer.Stop()
startACSSession(ctx, mockWsClient, timer, args, backoff, &mockSession{})
}()
// Wait for connection to be closed. If the connection is not closed
// due to inactivity, the test will time out
<-connectionClosed
}
// NewDockerTaskEngine returns a created, but uninitialized, DockerTaskEngine.
// The distinction between created and initialized is that when created it may
// be serialized/deserialized, but it will not communicate with docker until it
// is also initialized.
func NewDockerTaskEngine(cfg *config.Config, acceptInsecureCert bool) *DockerTaskEngine {
dockerTaskEngine := &DockerTaskEngine{
cfg: cfg,
acceptInsecureCert: acceptInsecureCert,
client: nil,
saver: statemanager.NewNoopStateManager(),
state: dockerstate.NewDockerTaskEngineState(),
managedTasks: make(map[string]*managedTask),
taskStopGroup: utilsync.NewSequentialWaitGroup(),
containerEvents: make(chan api.ContainerStateChange),
taskEvents: make(chan api.TaskStateChange),
}
return dockerTaskEngine
}
func TestFullUpdateFlow(t *testing.T) {
u, ctrl, cfg, mockfs, mockacs, mockhttp := mocks(t, nil)
defer ctrl.Finish()
var writtenFile bytes.Buffer
gomock.InOrder(
mockhttp.EXPECT().RoundTrip(mock_http.NewHTTPSimpleMatcher("GET", "https://s3.amazonaws.com/amazon-ecs-agent/update.tar")).Return(mock_http.SuccessResponse("update-tar-data"), nil),
mockfs.EXPECT().Create(gomock.Any()).Return(mock_os.NopReadWriteCloser(&writtenFile), nil),
mockfs.EXPECT().WriteFile("/tmp/test/desired-image", gomock.Any(), gomock.Any()).Return(nil),
mockacs.EXPECT().MakeRequest(gomock.Eq(&ecsacs.AckRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("mid").(*string),
})),
mockacs.EXPECT().MakeRequest(gomock.Eq(&ecsacs.AckRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("mid2").(*string),
})),
mockfs.EXPECT().Exit(exitcodes.ExitUpdate),
)
u.stageUpdateHandler()(&ecsacs.StageUpdateMessage{
ClusterArn: ptr("cluster").(*string),
ContainerInstanceArn: ptr("containerInstance").(*string),
MessageId: ptr("mid").(*string),
UpdateInfo: &ecsacs.UpdateInfo{
Location: ptr("https://s3.amazonaws.com/amazon-ecs-agent/update.tar").(*string),
Signature: ptr("6caeef375a080e3241781725b357890758d94b15d7ce63f6b2ff1cb5589f2007").(*string),
},
})
if writtenFile.String() != "update-tar-data" {
t.Error("Incorrect data written")
}
u.performUpdateHandler(statemanager.NewNoopStateManager(), engine.NewTaskEngine(cfg))(&ecsacs.PerformUpdateMessage{
ClusterArn: ptr("cluster").(*string),
ContainerInstanceArn: ptr("containerInstance").(*string),
MessageId: ptr("mid2").(*string),
UpdateInfo: &ecsacs.UpdateInfo{
Location: ptr("https://s3.amazonaws.com/amazon-ecs-agent/update.tar").(*string),
Signature: ptr("c54518806ff4d14b680c35784113e1e7478491fe").(*string),
},
})
}
// NewDockerTaskEngine returns a created, but uninitialized, DockerTaskEngine.
// The distinction between created and initialized is that when created it may
// be serialized/deserialized, but it will not communicate with docker until it
// is also initialized.
func NewDockerTaskEngine(cfg *config.Config, client DockerClient, credentialsManager credentials.Manager) *DockerTaskEngine {
dockerTaskEngine := &DockerTaskEngine{
cfg: cfg,
client: client,
saver: statemanager.NewNoopStateManager(),
state: dockerstate.NewDockerTaskEngineState(),
managedTasks: make(map[string]*managedTask),
taskStopGroup: utilsync.NewSequentialWaitGroup(),
containerEvents: make(chan api.ContainerStateChange),
taskEvents: make(chan api.TaskStateChange),
credentialsManager: credentialsManager,
}
return dockerTaskEngine
}
func initializeStateManager(cfg *config.Config, taskEngine engine.TaskEngine, cluster, containerInstanceArn, savedInstanceID *string) (statemanager.StateManager, error) {
if !cfg.Checkpoint {
return statemanager.NewNoopStateManager(), nil
}
stateManager, err := statemanager.NewStateManager(cfg,
statemanager.AddSaveable("TaskEngine", taskEngine),
statemanager.AddSaveable("ContainerInstanceArn", containerInstanceArn),
statemanager.AddSaveable("Cluster", cluster),
statemanager.AddSaveable("EC2InstanceID", savedInstanceID),
//The ACSSeqNum field is retained for compatibility with statemanager.EcsDataVersion 4 and
//can be removed in the future with a version bump.
statemanager.AddSaveable("ACSSeqNum", 1),
)
if err != nil {
return nil, err
}
return stateManager, nil
}
// TestHandlePayloadMessageAddTaskError tests that agent does not ack payload messages
// when task engine fails to add tasks
func TestHandlePayloadMessageAddTaskError(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
ecsClient := mock_api.NewMockECSClient(ctrl)
stateManager := statemanager.NewNoopStateManager()
credentialsManager := credentials.NewManager()
// Return error from AddTask
taskEngine.EXPECT().AddTask(gomock.Any()).Return(fmt.Errorf("oops")).Times(2)
ctx := context.Background()
buffer := newPayloadRequestHandler(ctx, taskEngine, ecsClient, clusterName, containerInstanceArn, nil, stateManager, refreshCredentialsHandler{}, credentialsManager)
// Test AddTask error with RUNNING task
payloadMessage := &ecsacs.PayloadMessage{
Tasks: []*ecsacs.Task{
&ecsacs.Task{
Arn: aws.String("t1"),
DesiredStatus: aws.String("RUNNING"),
},
},
MessageId: aws.String(payloadMessageId),
}
err := buffer.handleSingleMessage(payloadMessage)
if err == nil {
t.Error("Expected error while adding the task")
}
payloadMessage = &ecsacs.PayloadMessage{
Tasks: []*ecsacs.Task{
&ecsacs.Task{
Arn: aws.String("t1"),
DesiredStatus: aws.String("STOPPED"),
},
},
MessageId: aws.String(payloadMessageId),
}
// Test AddTask error with STOPPED task
err = buffer.handleSingleMessage(payloadMessage)
if err == nil {
t.Error("Expected error while adding the task")
}
}
func TestPerformUpdateWithUpdatesDisabled(t *testing.T) {
u, ctrl, cfg, _, mockacs, _ := mocks(t, &config.Config{
UpdatesEnabled: false,
})
defer ctrl.Finish()
mockacs.EXPECT().MakeRequest(&nackRequestMatcher{&ecsacs.NackRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("mid").(*string),
Reason: ptr("Updates are disabled").(*string),
}})
u.performUpdateHandler(statemanager.NewNoopStateManager(), engine.NewTaskEngine(cfg, nil, nil, nil, nil, nil))(&ecsacs.PerformUpdateMessage{
ClusterArn: ptr("cluster").(*string),
ContainerInstanceArn: ptr("containerInstance").(*string),
MessageId: ptr("mid").(*string),
UpdateInfo: &ecsacs.UpdateInfo{
Location: ptr("https://s3.amazonaws.com/amazon-ecs-agent/update.tar").(*string),
Signature: ptr("c54518806ff4d14b680c35784113e1e7478491fe").(*string),
},
})
}
func TestGetImageStateFromImageNameNoImageState(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
client := NewMockDockerClient(ctrl)
imageManager := &dockerImageManager{client: client, state: dockerstate.NewDockerTaskEngineState()}
imageManager.SetSaver(statemanager.NewNoopStateManager())
container := &api.Container{
Name: "testContainer",
Image: "testContainerImage",
}
imageInspected := &docker.Image{
ID: "sha256:qwerty",
}
client.EXPECT().InspectImage(container.Image).Return(imageInspected, nil).AnyTimes()
err := imageManager.AddContainerReferenceToImageState(container)
if err != nil {
t.Error("Error in adding container to an existing image state")
}
imageState := imageManager.GetImageStateFromImageName("noSuchImage")
if imageState != nil {
t.Error("Incorrect image state retrieved by image name")
}
}
func TestHandlerReconnects(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := mock_engine.NewMockTaskEngine(ctrl)
ecsclient := mock_api.NewMockECSClient(ctrl)
statemanager := statemanager.NewNoopStateManager()
closeWS := make(chan bool)
server, serverIn, _, _, err := startMockAcsServer(t, closeWS)
if err != nil {
t.Fatal(err)
}
ecsclient.EXPECT().DiscoverPollEndpoint("myArn").Return(server.URL, nil).AnyTimes()
taskEngine.EXPECT().Version().Return("Docker: 1.5.0", nil).AnyTimes()
ended := make(chan bool, 1)
go func() {
handler.StartSession("myArn", credentials.NewCredentialProvider("", ""), &config.Config{Cluster: "someCluster"}, taskEngine, ecsclient, statemanager, true)
// This should never return
ended <- true
}()
start := time.Now()
for i := 0; i < 10; i++ {
serverIn <- `{"type":"HeartbeatMessage","message":{"healthy":true}}`
closeWS <- true
}
if time.Since(start) > 2*time.Second {
t.Error("Test took longer than expected; backoff should not have occured for EOF")
}
select {
case <-ended:
t.Fatal("Should never stop session")
default:
}
}
// TestPayloadBufferHandlerWithCredentials tests if the async payloadBufferHandler routine
// acks the payload message and credentials after adding tasks
func TestPayloadBufferHandlerWithCredentials(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
ecsClient := mock_api.NewMockECSClient(ctrl)
stateManager := statemanager.NewNoopStateManager()
ctx, cancel := context.WithCancel(context.Background())
credentialsManager := credentials.NewManager()
// The payload message in the test consists of two tasks, record both of them in
// the order in which they were added
taskEngine := engine.NewMockTaskEngine(ctrl)
var firstAddedTask *api.Task
var secondAddedTask *api.Task
gomock.InOrder(
taskEngine.EXPECT().AddTask(gomock.Any()).Do(func(task *api.Task) {
firstAddedTask = task
}),
taskEngine.EXPECT().AddTask(gomock.Any()).Do(func(task *api.Task) {
secondAddedTask = task
}),
)
// The payload message in the test consists of two tasks, with credentials set
// for both. Record the credentials' ack and the payload message ack
var payloadAckRequested *ecsacs.AckRequest
var firstTaskCredentialsAckRequested *ecsacs.IAMRoleCredentialsAckRequest
var secondTaskCredentialsAckRequested *ecsacs.IAMRoleCredentialsAckRequest
mockWsClient := mock_wsclient.NewMockClientServer(ctrl)
gomock.InOrder(
mockWsClient.EXPECT().MakeRequest(gomock.Any()).Do(func(ackRequest *ecsacs.IAMRoleCredentialsAckRequest) {
firstTaskCredentialsAckRequested = ackRequest
}),
mockWsClient.EXPECT().MakeRequest(gomock.Any()).Do(func(ackRequest *ecsacs.IAMRoleCredentialsAckRequest) {
secondTaskCredentialsAckRequested = ackRequest
}),
mockWsClient.EXPECT().MakeRequest(gomock.Any()).Do(func(ackRequest *ecsacs.AckRequest) {
payloadAckRequested = ackRequest
// Cancel the context when the ack for the payload message is received
// This signals a successful workflow in the test
cancel()
}),
)
refreshCredsHandler := newRefreshCredentialsHandler(ctx, clusterName, containerInstanceArn, mockWsClient, credentialsManager, taskEngine)
defer refreshCredsHandler.clearAcks()
refreshCredsHandler.start()
payloadHandler := newPayloadRequestHandler(ctx, taskEngine, ecsClient, clusterName, containerInstanceArn, mockWsClient, stateManager, refreshCredsHandler, credentialsManager)
go payloadHandler.start()
firstTaskArn := "t1"
firstTaskCredentialsExpiration := "expiration"
firstTaskCredentialsRoleArn := "r1"
firstTaskCredentialsAccessKey := "akid"
firstTaskCredentialsSecretKey := "skid"
firstTaskCredentialsSessionToken := "token"
firstTaskCredentialsId := "credsid1"
secondTaskArn := "t2"
secondTaskCredentialsExpiration := "expirationSecond"
secondTaskCredentialsRoleArn := "r2"
secondTaskCredentialsAccessKey := "akid2"
secondTaskCredentialsSecretKey := "skid2"
secondTaskCredentialsSessionToken := "token2"
secondTaskCredentialsId := "credsid2"
// Send a payload message to the payloadBufferChannel
payloadHandler.messageBuffer <- &ecsacs.PayloadMessage{
Tasks: []*ecsacs.Task{
&ecsacs.Task{
Arn: aws.String(firstTaskArn),
RoleCredentials: &ecsacs.IAMRoleCredentials{
AccessKeyId: aws.String(firstTaskCredentialsAccessKey),
Expiration: aws.String(firstTaskCredentialsExpiration),
RoleArn: aws.String(firstTaskCredentialsRoleArn),
SecretAccessKey: aws.String(firstTaskCredentialsSecretKey),
SessionToken: aws.String(firstTaskCredentialsSessionToken),
CredentialsId: aws.String(firstTaskCredentialsId),
},
},
&ecsacs.Task{
Arn: aws.String(secondTaskArn),
RoleCredentials: &ecsacs.IAMRoleCredentials{
AccessKeyId: aws.String(secondTaskCredentialsAccessKey),
Expiration: aws.String(secondTaskCredentialsExpiration),
RoleArn: aws.String(secondTaskCredentialsRoleArn),
SecretAccessKey: aws.String(secondTaskCredentialsSecretKey),
SessionToken: aws.String(secondTaskCredentialsSessionToken),
CredentialsId: aws.String(secondTaskCredentialsId),
},
},
},
MessageId: aws.String(payloadMessageId),
ClusterArn: aws.String(cluster),
ContainerInstanceArn: aws.String(containerInstance),
}
// Wait till we get an ack
select {
case <-ctx.Done():
}
// Verify if payloadMessageId read from the ack buffer is correct
if aws.StringValue(payloadAckRequested.MessageId) != payloadMessageId {
t.Errorf("Message Id mismatch. Expected: %s, got: %s", payloadMessageId, aws.StringValue(payloadAckRequested.MessageId))
}
// Verify the correctness of the first task added to the engine and the
// credentials ack generated for it
expectedCredentialsAckForFirstTask := &ecsacs.IAMRoleCredentialsAckRequest{
MessageId: aws.String(payloadMessageId),
Expiration: aws.String(firstTaskCredentialsExpiration),
CredentialsId: aws.String(firstTaskCredentialsId),
}
expectedCredentialsForFirstTask := credentials.IAMRoleCredentials{
AccessKeyId: firstTaskCredentialsAccessKey,
Expiration: firstTaskCredentialsExpiration,
RoleArn: firstTaskCredentialsRoleArn,
SecretAccessKey: firstTaskCredentialsSecretKey,
SessionToken: firstTaskCredentialsSessionToken,
CredentialsId: firstTaskCredentialsId,
}
err := validateTaskAndCredentials(firstTaskCredentialsAckRequested, expectedCredentialsAckForFirstTask, firstAddedTask, firstTaskArn, expectedCredentialsForFirstTask)
if err != nil {
t.Errorf("Error validating added task or credentials ack for the same: %v", err)
}
// Verify the correctness of the second task added to the engine and the
// credentials ack generated for it
expectedCredentialsAckForSecondTask := &ecsacs.IAMRoleCredentialsAckRequest{
MessageId: aws.String(payloadMessageId),
Expiration: aws.String(secondTaskCredentialsExpiration),
CredentialsId: aws.String(secondTaskCredentialsId),
}
expectedCredentialsForSecondTask := credentials.IAMRoleCredentials{
AccessKeyId: secondTaskCredentialsAccessKey,
Expiration: secondTaskCredentialsExpiration,
RoleArn: secondTaskCredentialsRoleArn,
SecretAccessKey: secondTaskCredentialsSecretKey,
SessionToken: secondTaskCredentialsSessionToken,
CredentialsId: secondTaskCredentialsId,
}
err = validateTaskAndCredentials(secondTaskCredentialsAckRequested, expectedCredentialsAckForSecondTask, secondAddedTask, secondTaskArn, expectedCredentialsForSecondTask)
if err != nil {
t.Errorf("Error validating added task or credentials ack for the same: %v", err)
}
}
// permissions and limitations under the License.
package eventhandler
import (
"github.com/aws/amazon-ecs-agent/agent/api"
"github.com/aws/amazon-ecs-agent/agent/engine"
"github.com/aws/amazon-ecs-agent/agent/logger"
"github.com/aws/amazon-ecs-agent/agent/statemanager"
)
var log = logger.ForModule("eventhandler")
// statesaver is a package-wise statemanager which may be used to save any
// changes to a task or container's SentStatus
var statesaver statemanager.Saver = statemanager.NewNoopStateManager()
func HandleEngineEvents(taskEngine engine.TaskEngine, client api.ECSClient, saver statemanager.Saver) {
statesaver = saver
for {
taskEvents, containerEvents := taskEngine.TaskEvents()
for taskEvents != nil && containerEvents != nil {
select {
case event, open := <-containerEvents:
if !open {
containerEvents = nil
log.Error("Container events closed")
break
}
// Deletion of images in the order of LRU time: Happy path
// a. This includes starting up agent, pull images, start containers,
// account them in image manager, stop containers, remove containers, account this in image manager,
// b. Simulate the pulled time (so that it passes the minimum age criteria
// for getting chosen for deletion )
// c. Start image cleanup , ensure that ONLY the top 2 eligible LRU images
// are removed from the instance, and those deleted images’ image states are removed from image manager.
// d. Ensure images that do not pass the ‘minimumAgeForDeletion’ criteria are not removed.
// e. Image has not passed the ‘hasNoAssociatedContainers’ criteria.
// f. Ensure that that if not eligible, image is not deleted from the instance and image reference in ImageManager is not removed.
func TestIntegImageCleanupHappyCase(t *testing.T) {
cfg := defaultTestConfigIntegTest()
cfg.TaskCleanupWaitDuration = 5 * time.Second
// Set low values so this test can complete in a sane amout of time
cfg.MinimumImageDeletionAge = 1 * time.Second
cfg.NumImagesToDeletePerCycle = 2
// start agent
taskEngine, done, _ := setup(cfg, t)
imageManager := taskEngine.(*DockerTaskEngine).imageManager.(*dockerImageManager)
imageManager.SetSaver(statemanager.NewNoopStateManager())
defer func() {
done()
// Force cleanup all test images and containers
cleanupImages(imageManager)
}()
taskEvents, containerEvents := taskEngine.TaskEvents()
defer discardEvents(containerEvents)()
// Create test Task
taskName := "imgClean"
testTask := createImageCleanupTestTask(taskName)
go taskEngine.AddTask(testTask)
// Verify that Task is running
err := verifyTaskIsRunning(taskEvents, testTask)
if err != nil {
t.Fatal(err)
}
imageState1 := imageManager.GetImageStateFromImageName(testImage1Name)
if imageState1 == nil {
t.Fatalf("Could not find image state for %s", testImage1Name)
} else {
t.Logf("Found image state for %s", testImage1Name)
}
imageState2 := imageManager.GetImageStateFromImageName(testImage2Name)
if imageState2 == nil {
t.Fatalf("Could not find image state for %s", testImage2Name)
} else {
t.Logf("Found image state for %s", testImage2Name)
}
imageState3 := imageManager.GetImageStateFromImageName(testImage3Name)
if imageState3 == nil {
t.Fatalf("Could not find image state for %s", testImage3Name)
} else {
t.Logf("Found image state for %s", testImage3Name)
}
imageState1ImageId := imageState1.Image.ImageID
imageState2ImageId := imageState2.Image.ImageID
imageState3ImageId := imageState3.Image.ImageID
// Set the ImageState.LastUsedAt to a value far in the past to ensure the test images are deleted.
// This will make these test images the LRU images.
imageState1.LastUsedAt = imageState1.LastUsedAt.Add(-99995 * time.Hour)
imageState2.LastUsedAt = imageState2.LastUsedAt.Add(-99994 * time.Hour)
imageState3.LastUsedAt = imageState3.LastUsedAt.Add(-99993 * time.Hour)
// Verify Task is stopped.
verifyTaskIsStopped(taskEvents, testTask)
// Allow Task cleanup to occur
time.Sleep(5 * time.Second)
// Verify Task is cleaned up
err = verifyTaskIsCleanedUp(taskName, taskEngine)
if err != nil {
t.Fatal(err)
}
// Call Image removal
imageManager.removeUnusedImages()
// Verify top 2 LRU images are deleted from image manager
err = verifyImagesAreRemoved(imageManager, imageState1ImageId, imageState2ImageId)
if err != nil {
t.Fatal(err)
}
// Verify 3rd LRU image is not removed
err = verifyImagesAreNotRemoved(imageManager, imageState3ImageId)
if err != nil {
t.Fatal(err)
}
// Verify top 2 LRU images are removed from docker
_, err = taskEngine.(*DockerTaskEngine).client.InspectImage(imageState1ImageId)
if err != docker.ErrNoSuchImage {
t.Fatalf("Image was not removed successfully")
}
_, err = taskEngine.(*DockerTaskEngine).client.InspectImage(imageState2ImageId)
if err != docker.ErrNoSuchImage {
t.Fatalf("Image was not removed successfully")
}
// Verify 3rd LRU image has not been removed from Docker
_, err = taskEngine.(*DockerTaskEngine).client.InspectImage(imageState3ImageId)
if err != nil {
t.Fatalf("Image should not have been removed from Docker")
}
}
func TestNewerUpdateMessages(t *testing.T) {
u, ctrl, cfg, mockfs, mockacs, mockhttp := mocks(t, nil)
defer ctrl.Finish()
var writtenFile bytes.Buffer
gomock.InOrder(
mockhttp.EXPECT().RoundTrip(mock_http.NewHTTPSimpleMatcher("GET", "https://s3.amazonaws.com/amazon-ecs-agent/update.tar")).Return(mock_http.SuccessResponse("update-tar-data"), nil),
mockfs.EXPECT().Create(gomock.Any()).Return(mock_os.NopReadWriteCloser(&writtenFile), nil),
mockfs.EXPECT().WriteFile("/tmp/test/desired-image", gomock.Any(), gomock.Any()).Return(nil),
mockacs.EXPECT().MakeRequest(gomock.Eq(&ecsacs.AckRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("StageMID").(*string),
})),
mockacs.EXPECT().MakeRequest(&nackRequestMatcher{&ecsacs.NackRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("StageMID").(*string),
Reason: ptr("New update arrived: StageMIDNew").(*string),
}}),
mockhttp.EXPECT().RoundTrip(mock_http.NewHTTPSimpleMatcher("GET", "https://s3.amazonaws.com/amazon-ecs-agent/new.tar")).Return(mock_http.SuccessResponse("newer-update-tar-data"), nil),
mockfs.EXPECT().Create(gomock.Any()).Return(mock_os.NopReadWriteCloser(&writtenFile), nil),
mockfs.EXPECT().WriteFile("/tmp/test/desired-image", gomock.Any(), gomock.Any()).Return(nil),
mockacs.EXPECT().MakeRequest(gomock.Eq(&ecsacs.AckRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("StageMIDNew").(*string),
})),
mockacs.EXPECT().MakeRequest(gomock.Eq(&ecsacs.AckRequest{
Cluster: ptr("cluster").(*string),
ContainerInstance: ptr("containerInstance").(*string),
MessageId: ptr("mid2").(*string),
})),
mockfs.EXPECT().Exit(exitcodes.ExitUpdate),
)
u.stageUpdateHandler()(&ecsacs.StageUpdateMessage{
ClusterArn: ptr("cluster").(*string),
ContainerInstanceArn: ptr("containerInstance").(*string),
MessageId: ptr("StageMID").(*string),
UpdateInfo: &ecsacs.UpdateInfo{
Location: ptr("https://s3.amazonaws.com/amazon-ecs-agent/update.tar").(*string),
Signature: ptr("6caeef375a080e3241781725b357890758d94b15d7ce63f6b2ff1cb5589f2007").(*string),
},
})
if writtenFile.String() != "update-tar-data" {
t.Error("Incorrect data written")
}
writtenFile.Reset()
// Never perform, make sure a new hash results in a new stage
u.stageUpdateHandler()(&ecsacs.StageUpdateMessage{
ClusterArn: ptr("cluster").(*string),
ContainerInstanceArn: ptr("containerInstance").(*string),
MessageId: ptr("StageMIDNew").(*string),
UpdateInfo: &ecsacs.UpdateInfo{
Location: ptr("https://s3.amazonaws.com/amazon-ecs-agent/new.tar").(*string),
Signature: ptr("9c6ea7bd7d49f95b6d516517e453b965897109bf8a1d6ff3a6e57287049eb2de").(*string),
},
})
if writtenFile.String() != "newer-update-tar-data" {
t.Error("Incorrect data written")
}
u.performUpdateHandler(statemanager.NewNoopStateManager(), engine.NewTaskEngine(cfg))(&ecsacs.PerformUpdateMessage{
ClusterArn: ptr("cluster").(*string),
ContainerInstanceArn: ptr("containerInstance").(*string),
MessageId: ptr("mid2").(*string),
UpdateInfo: &ecsacs.UpdateInfo{
Location: ptr("https://s3.amazonaws.com/amazon-ecs-agent/update.tar").(*string),
Signature: ptr("c54518806ff4d14b680c35784113e1e7478491fe").(*string),
},
})
}
func TestHeartbeatOnlyWhenIdle(t *testing.T) {
testTime := ttime.NewTestTime()
ttime.SetTime(testTime)
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
ecsclient := mock_api.NewMockECSClient(ctrl)
statemanager := statemanager.NewNoopStateManager()
closeWS := make(chan bool)
server, serverIn, requestsChan, errChan, err := startMockAcsServer(t, closeWS)
defer close(serverIn)
go func() {
for {
<-requestsChan
}
}()
if err != nil {
t.Fatal(err)
}
// We're testing that it does not reconnect here; must be the case
ecsclient.EXPECT().DiscoverPollEndpoint("myArn").Return(server.URL, nil).Times(1)
taskEngine.EXPECT().Version().Return("Docker: 1.5.0", nil).AnyTimes()
ctx, cancel := context.WithCancel(context.Background())
ended := make(chan bool, 1)
go func() {
handler.StartSession(ctx, handler.StartSessionArguments{
ContainerInstanceArn: "myArn",
CredentialProvider: credentials.AnonymousCredentials,
Config: &config.Config{Cluster: "someCluster"},
TaskEngine: taskEngine,
ECSClient: ecsclient,
StateManager: statemanager,
AcceptInvalidCert: true,
})
ended <- true
}()
taskAdded := make(chan bool)
taskEngine.EXPECT().AddTask(gomock.Any()).Do(func(interface{}) {
taskAdded <- true
}).Times(10)
for i := 0; i < 10; i++ {
serverIn <- samplePayloadMessage
testTime.Warp(1 * time.Minute)
<-taskAdded
}
select {
case <-ended:
t.Fatal("Should not have stop session")
case err := <-errChan:
t.Fatal("Error should not have been returned from server", err)
default:
}
go server.Close()
cancel()
<-ended
}
func TestHandlerDoesntLeakGouroutines(t *testing.T) {
ctrl := gomock.NewController(t)
defer ctrl.Finish()
taskEngine := engine.NewMockTaskEngine(ctrl)
ecsclient := mock_api.NewMockECSClient(ctrl)
statemanager := statemanager.NewNoopStateManager()
testTime := ttime.NewTestTime()
ttime.SetTime(testTime)
closeWS := make(chan bool)
server, serverIn, requests, errs, err := startMockAcsServer(t, closeWS)
if err != nil {
t.Fatal(err)
}
go func() {
for {
select {
case <-requests:
case <-errs:
}
}
}()
timesConnected := 0
ecsclient.EXPECT().DiscoverPollEndpoint("myArn").Return(server.URL, nil).AnyTimes().Do(func(_ interface{}) {
timesConnected++
})
taskEngine.EXPECT().Version().Return("Docker: 1.5.0", nil).AnyTimes()
taskEngine.EXPECT().AddTask(gomock.Any()).AnyTimes()
ctx, cancel := context.WithCancel(context.Background())
ended := make(chan bool, 1)
go func() {
handler.StartSession(ctx, handler.StartSessionArguments{"myArn", credentials.AnonymousCredentials, &config.Config{Cluster: "someCluster"}, taskEngine, ecsclient, statemanager, true})
ended <- true
}()
// Warm it up
serverIn <- `{"type":"HeartbeatMessage","message":{"healthy":true}}`
serverIn <- samplePayloadMessage
beforeGoroutines := runtime.NumGoroutine()
for i := 0; i < 100; i++ {
serverIn <- `{"type":"HeartbeatMessage","message":{"healthy":true}}`
serverIn <- samplePayloadMessage
closeWS <- true
}
cancel()
testTime.Cancel()
<-ended
afterGoroutines := runtime.NumGoroutine()
t.Logf("Gorutines after 1 and after 100 acs messages: %v and %v", beforeGoroutines, afterGoroutines)
if timesConnected < 50 {
t.Fatal("Expected times connected to be a large number, was ", timesConnected)
}
if afterGoroutines > beforeGoroutines+5 {
t.Error("Goroutine leak, oh no!")
pprof.Lookup("goroutine").WriteTo(os.Stdout, 1)
}
}