From d4496240f394ed8d6df50c8c66ce2895ecdeab6b Mon Sep 17 00:00:00 2001
From: Raghav Muddur <r@nmvk.com>
Date: Thu, 30 Sep 2021 13:31:10 -0700
Subject: [PATCH] Scale up/down and Scale in/out memory considerations

When spec is updated to perform Scale up and Scale in
we would need to perform Scale up first followed by
Scale in.

Similarly when we perform Scale down and Scale out
we would need to perform Scale out first.
---
 .../replication_group/custom_update_api.go    | 13 +++++
 .../custom_update_api_test.go                 | 55 +++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/pkg/resource/replication_group/custom_update_api.go b/pkg/resource/replication_group/custom_update_api.go
index 0f33e241..7f49bb02 100644
--- a/pkg/resource/replication_group/custom_update_api.go
+++ b/pkg/resource/replication_group/custom_update_api.go
@@ -101,6 +101,19 @@ func (rm *resourceManager) CustomModifyReplicationGroup(
 		return rm.decreaseReplicaCount(ctx, desired, latest)
 	}
 
+	// If there is a scale up modification, then we would prioritize it
+	// over increase/decrease shards. This is important since performing
+	// scale in without scale up might fail due to insufficient memory.
+	if delta.DifferentAt("Spec.CacheNodeType") && desired.ko.Status.AllowedScaleUpModifications != nil {
+		if desired.ko.Spec.CacheNodeType != nil {
+			for _, scaleUpInstance := range desired.ko.Status.AllowedScaleUpModifications {
+				if *scaleUpInstance == *desired.ko.Spec.CacheNodeType {
+					return nil, nil
+				}
+			}
+		}
+	}
+
 	// increase/decrease shards
 	if rm.shardConfigurationsDiffer(desired, latest) {
 		return rm.updateShardConfiguration(ctx, desired, latest)
diff --git a/pkg/resource/replication_group/custom_update_api_test.go b/pkg/resource/replication_group/custom_update_api_test.go
index 665c041e..6a08852b 100644
--- a/pkg/resource/replication_group/custom_update_api_test.go
+++ b/pkg/resource/replication_group/custom_update_api_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/aws-controllers-k8s/elasticache-controller/pkg/testutil"
 	ackmetrics "github.com/aws-controllers-k8s/runtime/pkg/metrics"
 	"github.com/aws-controllers-k8s/runtime/pkg/requeue"
+	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/mock"
 	"go.uber.org/zap/zapcore"
@@ -391,6 +392,60 @@ func TestCustomModifyReplicationGroup_NodeGroup_available(t *testing.T) {
 	})
 }
 
+func TestCustomModifyReplicationGroup_ScaleUpAndDown_And_Resharding(t *testing.T)  {
+	assert := assert.New(t)
+
+	// Tests
+	t.Run("ScaleInAndScaleUp=Diff", func(t *testing.T) {
+		desired := provideResource()
+		latest := provideResource()
+		rgId := "RGID"
+		desired.ko.Spec.ReplicationGroupID = &rgId
+		latest.ko.Spec.ReplicationGroupID = &rgId
+
+		mocksdkapi := &mocksvcsdkapi.ElastiCacheAPI{}
+		rm := provideResourceManagerWithMockSDKAPI(mocksdkapi)
+
+		var delta ackcompare.Delta
+		delta.Add("Spec.CacheNodeType", "cache.t3.small", "cache.m5.large")
+		delta.Add("Spec.NumNodeGroups", 3, 4)
+
+		var ctx context.Context
+		res, err := rm.CustomModifyReplicationGroup(ctx, desired, latest, &delta)
+		assert.Nil(res)
+		assert.Nil(err)
+		assert.Empty(mocksdkapi.Calls)
+	})
+
+	t.Run("ScaleOutAndScaleDown=Diff", func(t *testing.T) {
+		desired := provideResource()
+		latest := provideResource()
+		rgId := "RGID"
+		desired.ko.Spec.ReplicationGroupID = &rgId
+		latest.ko.Spec.ReplicationGroupID = &rgId
+
+		mocksdkapi := &mocksvcsdkapi.ElastiCacheAPI{}
+		rm := provideResourceManagerWithMockSDKAPI(mocksdkapi)
+
+		var delta ackcompare.Delta
+		delta.Add("Spec.CacheNodeType", "cache.t3.small", "cache.t3.micro")
+		oldshardCount := int64(4)
+		newShardCount := int64(10)
+		delta.Add("Spec.NumNodeGroups", oldshardCount, newShardCount)
+		desired.ko.Spec.NumNodeGroups = &newShardCount
+		latest.ko.Spec.NumNodeGroups = &oldshardCount
+		mocksdkapi.On("ModifyReplicationGroupShardConfigurationWithContext", mock.Anything, mock.Anything).Return(nil,
+			awserr.New("Invalid", "Invalid error", nil))
+		var ctx context.Context
+		res, err := rm.CustomModifyReplicationGroup(ctx, desired, latest, &delta)
+		assert.Nil(res)
+		assert.NotNil(err)
+		assert.NotEmpty(mocksdkapi.Calls)
+		assert.Equal("ModifyReplicationGroupShardConfigurationWithContext", mocksdkapi.Calls[0].Method)
+	})
+
+}
+
 // TestReplicaCountDifference tests scenarios to check if desired, latest replica count
 // configurations differ
 func TestReplicaCountDifference(t *testing.T) {