From e64e5368eeca63345d5500680d9282fc579312a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gu=C3=B0laugur=20Stef=C3=A1n=20Egilsson?=
 <gudlaugur.egilsson@gmail.com>
Date: Wed, 5 Mar 2025 08:23:46 +0000
Subject: [PATCH] fix(aws-sd): service instances registration and
 deregistration (#5135)

* Only de-register removed targets

* Use maps for current targets lookup.

* Use camelCase, not _
---
 provider/awssd/aws_sd.go      | 22 ++++++++-----
 provider/awssd/aws_sd_test.go | 58 +++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/provider/awssd/aws_sd.go b/provider/awssd/aws_sd.go
index babe48593..95ecda406 100644
--- a/provider/awssd/aws_sd.go
+++ b/provider/awssd/aws_sd.go
@@ -222,13 +222,6 @@ func (p *AWSSDProvider) ApplyChanges(ctx context.Context, changes *plan.Changes)
 		return err
 	}
 
-	// Deletes must be executed first to support update case.
-	// When just list of targets is updated `[1.2.3.4] -> [1.2.3.4, 1.2.3.5]` it is translated to:
-	// ```
-	// deletes = [1.2.3.4]
-	// creates = [1.2.3.4, 1.2.3.5]
-	// ```
-	// then when deletes are executed after creates it will miss the `1.2.3.4` instance.
 	err = p.submitDeletes(ctx, namespaces, changes.Delete)
 	if err != nil {
 		return err
@@ -252,7 +245,20 @@ func (p *AWSSDProvider) updatesToCreates(changes *plan.Changes) (creates []*endp
 		current := updateNewMap[old.DNSName]
 
 		if !old.Targets.Same(current.Targets) {
-			// when targets differ the old instances need to be de-registered first
+			currentTargetsMap := make(map[string]struct{}, len(current.Targets))
+			for _, newTarget := range current.Targets {
+				currentTargetsMap[newTarget] = struct{}{}
+			}
+
+			// If targets changed, only deregister removed targets (i.e. in `UpdateOld` but not in `UpdateNew`)
+			targetsToRemove := make(endpoint.Targets, 0)
+			for _, oldTarget := range old.Targets {
+				if _, found := currentTargetsMap[oldTarget]; !found {
+					targetsToRemove = append(targetsToRemove, oldTarget)
+				}
+			}
+
+			old.Targets = targetsToRemove
 			deletes = append(deletes, old)
 		}
 
diff --git a/provider/awssd/aws_sd_test.go b/provider/awssd/aws_sd_test.go
index 38cc0711b..dd851c455 100644
--- a/provider/awssd/aws_sd_test.go
+++ b/provider/awssd/aws_sd_test.go
@@ -52,6 +52,9 @@ type AWSSDClientStub struct {
 
 	// map[service_id] => map[inst_id]instance
 	instances map[string]map[string]*sdtypes.Instance
+
+	// []inst_id
+	deregistered []string
 }
 
 func (s *AWSSDClientStub) CreateService(ctx context.Context, input *sd.CreateServiceInput, optFns ...func(*sd.Options)) (*sd.CreateServiceOutput, error) {
@@ -79,6 +82,7 @@ func (s *AWSSDClientStub) CreateService(ctx context.Context, input *sd.CreateSer
 func (s *AWSSDClientStub) DeregisterInstance(ctx context.Context, input *sd.DeregisterInstanceInput, optFns ...func(options *sd.Options)) (*sd.DeregisterInstanceOutput, error) {
 	serviceInstances := s.instances[*input.ServiceId]
 	delete(serviceInstances, *input.InstanceId)
+	s.deregistered = append(s.deregistered, *input.InstanceId)
 
 	return &sd.DeregisterInstanceOutput{}, nil
 }
@@ -436,6 +440,60 @@ func TestAWSSDProvider_ApplyChanges(t *testing.T) {
 	assert.Empty(t, endpoints)
 }
 
+func TestAWSSDProvider_ApplyChanges_Update(t *testing.T) {
+	namespaces := map[string]*sdtypes.Namespace{
+		"private": {
+			Id:   aws.String("private"),
+			Name: aws.String("private.com"),
+			Type: sdtypes.NamespaceTypeDnsPrivate,
+		},
+	}
+
+	api := &AWSSDClientStub{
+		namespaces: namespaces,
+		services:   make(map[string]map[string]*sdtypes.Service),
+		instances:  make(map[string]map[string]*sdtypes.Instance),
+	}
+
+	oldEndpoints := []*endpoint.Endpoint{
+		{DNSName: "service1.private.com", Targets: endpoint.Targets{"1.2.3.4", "1.2.3.5"}, RecordType: endpoint.RecordTypeA, RecordTTL: 60},
+	}
+
+	newEndpoints := []*endpoint.Endpoint{
+		{DNSName: "service1.private.com", Targets: endpoint.Targets{"1.2.3.4", "1.2.3.6"}, RecordType: endpoint.RecordTypeA, RecordTTL: 60},
+	}
+
+	provider := newTestAWSSDProvider(api, endpoint.NewDomainFilter([]string{}), "", "")
+
+	ctx := context.Background()
+
+	// apply creates
+	provider.ApplyChanges(ctx, &plan.Changes{
+		Create: oldEndpoints,
+	})
+
+	ctx = context.Background()
+
+	// apply update
+	provider.ApplyChanges(ctx, &plan.Changes{
+		UpdateOld: oldEndpoints,
+		UpdateNew: newEndpoints,
+	})
+
+	// make sure services were created
+	assert.Len(t, api.services["private"], 1)
+	existingServices, _ := provider.ListServicesByNamespaceID(ctx, namespaces["private"].Id)
+	assert.NotNil(t, existingServices["service1"])
+
+	// make sure instances were registered
+	endpoints, _ := provider.Records(ctx)
+	assert.True(t, testutils.SameEndpoints(newEndpoints, endpoints), "expected and actual endpoints don't match, expected=%v, actual=%v", newEndpoints, endpoints)
+
+	// make sure only one instance is de-registered
+	assert.Len(t, api.deregistered, 1)
+	assert.Equal(t, api.deregistered[0], "1.2.3.5", "wrong target de-registered")
+}
+
 func TestAWSSDProvider_ListNamespaces(t *testing.T) {
 	namespaces := map[string]*sdtypes.Namespace{
 		"private": {