Merge pull request #15388 from amazon-mq/fix/aws-peer-discovery-instance-states

AWS peer discovery: filter EC2 instances by state

Author: michaelklishin

deps/rabbitmq_peer_discovery_aws/priv/schema/rabbitmq_peer_discovery_aws.schema

@@ -108,3 +108,17 @@ fun(Conf) ->
         Value     -> rabbit_peer_discovery_util:as_list(Value)
     end
 end}.
+
+%% ec2_instance_states
+
+{mapping, "cluster_formation.aws.ec2_instance_states.$state", "rabbit.cluster_formation.peer_discovery_aws.aws_ec2_instance_states", [
+    {datatype, string}
+]}.
+
+{translation, "rabbit.cluster_formation.peer_discovery_aws.aws_ec2_instance_states",
+fun(Conf) ->
+    case cuttlefish_variable:filter_by_prefix("cluster_formation.aws.ec2_instance_states", Conf) of
+        [] -> ["running", "pending"];
+        Settings -> [V || {_K, V} <- Settings]
+    end
+end}.

deps/rabbitmq_peer_discovery_aws/src/rabbit_peer_discovery_aws.erl

@@ -30,6 +30,8 @@
 
 -define(BACKEND_CONFIG_KEY, peer_discovery_aws).
 
+-define(VALID_EC2_INSTANCE_STATES, ["pending", "running", "shutting-down", "terminated", "stopping", "stopped"]).
+
 -define(CONFIG_MAPPING,
          #{
           aws_autoscaling                    => #peer_discovery_config_entry_meta{
@@ -66,6 +68,11 @@
                                                    type          = atom,
                                                    env_variable  = "AWS_USE_PRIVATE_IP",
                                                    default_value = false
+                                                  },
+          aws_ec2_instance_states            => #peer_discovery_config_entry_meta{
+                                                   type          = list,
+                                                   env_variable  = "AWS_EC2_INSTANCE_STATES",
+                                                   default_value = ["running", "pending"]
                                                   }
          }).
 
@@ -280,8 +287,10 @@ get_hostname_by_instance_ids(Instances, Tag) ->
     QArgs = build_instance_list_qargs(Instances,
                                       [{"Action", "DescribeInstances"},
                                        {"Version", "2015-10-01"}]),
-    QArgs2 = lists:keysort(1, maybe_add_tag_filters(Tag, QArgs, 1)),
-    Path = "/?" ++ rabbitmq_aws_urilib:build_query_string(QArgs2),
+    QArgs1 = maybe_add_tag_filters(Tag, QArgs, 1),
+    QArgs2 = maybe_add_instance_state_filters(QArgs1, length(QArgs1) + 1),
+    QArgs3 = lists:keysort(1, QArgs2),
+    Path = "/?" ++ rabbitmq_aws_urilib:build_query_string(QArgs3),
     get_hostname_names(Path).
 
 -spec build_instance_list_qargs(Instances :: list(), Accum :: list()) -> list().
@@ -306,6 +315,53 @@ maybe_add_tag_filters(Tags, QArgs, Num) ->
                   end, {QArgs, Num}, Tags),
     Filters.
 
+-spec maybe_add_instance_state_filters(filters(), integer()) -> filters().
+maybe_add_instance_state_filters(QArgs, Num) ->
+    M = ?CONFIG_MODULE:config_map(?BACKEND_CONFIG_KEY),
+    States = get_config_key(aws_ec2_instance_states, M),
+    case States of
+        [] ->
+            QArgs;
+        [_|_] ->
+            ValidStates = validate_instance_states(States),
+            add_instance_state_filters(ValidStates, QArgs, Num)
+    end.
+
+-spec validate_instance_states(list()) -> list().
+validate_instance_states(States) ->
+    NormalizedStates = [normalize_state(State) || State <- States],
+    {Valid, Invalid} = lists:partition(
+        fun(State) -> lists:member(State, ?VALID_EC2_INSTANCE_STATES) end,
+        NormalizedStates),
+    case Invalid of
+        [] ->
+            ok;
+        [_|_] ->
+            rabbit_log:warning(
+                "Ignoring invalid EC2 instance states in configuration: ~tp. "
+                "Valid states are: ~tp", [Invalid, ?VALID_EC2_INSTANCE_STATES])
+    end,
+    Valid.
+
+-spec normalize_state(atom() | string()) -> string().
+normalize_state(State) when is_atom(State) ->
+    atom_to_list(State);
+normalize_state(State) when is_list(State) ->
+    State.
+
+-spec add_instance_state_filters(list(), filters(), integer()) -> filters().
+add_instance_state_filters(States, QArgs, Num) ->
+    FilterName = {"Filter." ++ integer_to_list(Num) ++ ".Name", "instance-state-name"},
+    StateFilters = lists:foldl(
+        fun(State, {Acc, Index}) ->
+            Filter = {"Filter." ++ integer_to_list(Num) ++ ".Value." ++ integer_to_list(Index), State},
+            {[Filter | Acc], Index + 1}
+        end,
+        {[], 1},
+        States),
+    {StateValues, _} = StateFilters,
+    [FilterName | StateValues] ++ QArgs.
+
 -spec get_node_list_from_tags(tags()) -> {ok, {[node()], disc}}.
 get_node_list_from_tags(M) when map_size(M) =:= 0 ->
     ?LOG_WARNING("Cannot discover any nodes because AWS tags are not configured!", []),
@@ -336,8 +392,10 @@ get_hostname_names(Path) ->
 
 get_hostname_by_tags(Tags) ->
     QArgs = [{"Action", "DescribeInstances"}, {"Version", "2015-10-01"}],
-    QArgs2 = lists:keysort(1, maybe_add_tag_filters(Tags, QArgs, 1)),
-    Path = "/?" ++ rabbitmq_aws_urilib:build_query_string(QArgs2),
+    QArgs1 = maybe_add_tag_filters(Tags, QArgs, 1),
+    QArgs2 = maybe_add_instance_state_filters(QArgs1, length(QArgs1) + 1),
+    QArgs3 = lists:keysort(1, QArgs2),
+    Path = "/?" ++ rabbitmq_aws_urilib:build_query_string(QArgs3),
     case get_hostname_names(Path) of
         error ->
             ?LOG_WARNING("Cannot discover any nodes because AWS "

deps/rabbitmq_peer_discovery_aws/test/config_schema_SUITE_data/rabbitmq_peer_discovery_aws.snippets

@@ -103,5 +103,30 @@
                 ]}
             ]}
         ], [rabbitmq_peer_discovery_aws]
+    },
+    {aws_ec2_instance_states,
+        "cluster_formation.aws.ec2_instance_states.1 = running
+         cluster_formation.aws.ec2_instance_states.2 = pending",
+        [
+            {rabbit, [
+                {cluster_formation, [
+                    {peer_discovery_aws, [
+                        {aws_ec2_instance_states, ["running", "pending"]}
+                    ]}
+                ]}
+            ]}
+        ], [rabbitmq_peer_discovery_aws]
+    },
+    {aws_ec2_instance_states_single,
+        "cluster_formation.aws.ec2_instance_states.1 = running",
+        [
+            {rabbit, [
+                {cluster_formation, [
+                    {peer_discovery_aws, [
+                        {aws_ec2_instance_states, ["running"]}
+                    ]}
+                ]}
+            ]}
+        ], [rabbitmq_peer_discovery_aws]
     }
 ].

deps/rabbitmq_peer_discovery_aws/test/unit_SUITE.erl

@@ -22,10 +22,14 @@ groups() ->
     [
      {unit, [], [
                  maybe_add_tag_filters,
+                 maybe_add_instance_state_filters,
+                 validate_instance_states,
                  get_hostname_name_from_reservation_set,
                  registration_support,
                  network_interface_sorting,
-                 private_ip_address_sorting
+                 private_ip_address_sorting,
+                 get_hostname_by_instance_ids_with_state_filter,
+                 get_hostname_by_tags_with_state_filter
                 ]},
      {lock, [], [
                  lock_single_node,
@@ -48,6 +52,27 @@ maybe_add_tag_filters(_Config) ->
     Result = lists:sort(rabbit_peer_discovery_aws:maybe_add_tag_filters(Tags, [], 1)),
     ?assertEqual(Expectation, Result).
 
+maybe_add_instance_state_filters(_Config) ->
+    application:set_env(rabbit, cluster_formation,
+                       [{peer_discovery_aws, [{aws_ec2_instance_states, ["running", "pending"]}]}]),
+    QArgs = [{"Action", "DescribeInstances"}, {"Version", "2015-10-01"}],
+    Result = rabbit_peer_discovery_aws:maybe_add_instance_state_filters(QArgs, 1),
+    Expectation = [{"Filter.1.Name", "instance-state-name"},
+                   {"Filter.1.Value.2", "pending"},
+                   {"Filter.1.Value.1", "running"},
+                   {"Action", "DescribeInstances"},
+                   {"Version", "2015-10-01"}],
+    ?assertEqual(Expectation, Result),
+    application:unset_env(rabbit, cluster_formation).
+
+validate_instance_states(_Config) ->
+    ValidStates = ["pending", "running", "shutting-down", "terminated", "stopping", "stopped"],
+    ?assertEqual(ValidStates, rabbit_peer_discovery_aws:validate_instance_states(ValidStates)),
+    ?assertEqual(["running"], rabbit_peer_discovery_aws:validate_instance_states(["running", "invalid"])),
+    ?assertEqual([], rabbit_peer_discovery_aws:validate_instance_states(["bogus", "invalid"])),
+    ?assertEqual(["running", "pending"], rabbit_peer_discovery_aws:validate_instance_states([running, pending])),
+    ?assertEqual(["running"], rabbit_peer_discovery_aws:validate_instance_states([running, invalid])).
+
 get_hostname_name_from_reservation_set(_Config) ->
     ok = eunit:test({
       foreach,
@@ -89,6 +114,44 @@ get_hostname_name_from_reservation_set(_Config) ->
         end}]
     }).
 
+get_hostname_by_instance_ids_with_state_filter(_Config) ->
+    application:set_env(rabbit, cluster_formation,
+                       [{peer_discovery_aws, [{aws_ec2_instance_states, ["running", "pending"]}]}]),
+    meck:new(rabbitmq_aws, [passthrough]),
+    meck:expect(rabbitmq_aws, api_get_request,
+        fun("ec2", Path) ->
+            ?assert(string:str(Path, "Filter") > 0),
+            ?assert(string:str(Path, "instance-state-name") > 0),
+            ?assert(string:str(Path, "Value.1=running") > 0),
+            ?assert(string:str(Path, "Value.2=pending") > 0),
+            {ok, mock_describe_instances_response()}
+        end),
+    Result = rabbit_peer_discovery_aws:get_hostname_by_instance_ids(
+        ["i-abc123", "i-def456"], #{}),
+    ?assertEqual(["ip-10-0-16-29.eu-west-1.compute.internal",
+                  "ip-10-0-16-31.eu-west-1.compute.internal"], Result),
+    meck:unload(rabbitmq_aws),
+    application:unset_env(rabbit, cluster_formation).
+
+get_hostname_by_tags_with_state_filter(_Config) ->
+    application:set_env(rabbit, cluster_formation,
+                       [{peer_discovery_aws, [{aws_ec2_instance_states, ["running"]}]}]),
+    meck:new(rabbitmq_aws, [passthrough]),
+    meck:expect(rabbitmq_aws, api_get_request,
+        fun("ec2", Path) ->
+            ?assert(string:str(Path, "Filter") > 0),
+            ?assert(string:str(Path, "instance-state-name") > 0),
+            ?assert(string:str(Path, "Value.1=running") > 0),
+            ?assert(string:str(Path, "tag%3Aservice") > 0),
+            {ok, mock_describe_instances_response()}
+        end),
+    Tags = maps:from_list([{"service", "rabbitmq"}]),
+    Result = rabbit_peer_discovery_aws:get_hostname_by_tags(Tags),
+    ?assertEqual(["ip-10-0-16-29.eu-west-1.compute.internal",
+                  "ip-10-0-16-31.eu-west-1.compute.internal"], Result),
+    meck:unload(rabbitmq_aws),
+    application:unset_env(rabbit, cluster_formation).
+
 registration_support(_Config) ->
     ?assertEqual(false, rabbit_peer_discovery_aws:supports_registration()).
 
@@ -289,3 +352,7 @@ reservation_set() ->
                         ]}
                      ]}]}]},
                    {"privateIpAddress","10.0.16.31"}]}]}]}].
+
+mock_describe_instances_response() ->
+    [{"DescribeInstancesResponse",
+      [{"reservationSet", reservation_set()}]}].