Improve emqx_mqtt_caps:get_caps/1 (#3198)

src/emqx_mqtt_caps.erl

@@ -29,6 +29,8 @@
         , get_caps/3
         ]).
 
+-export([default_caps/0]).
+
 -export([default/0]).
 
 -export_type([caps/0]).
@@ -116,23 +118,42 @@ do_check_sub(#{is_shared := true}, #{shared_subscription := false}) ->
     {error, ?RC_SHARED_SUBSCRIPTIONS_NOT_SUPPORTED};
 do_check_sub(_Flags, _Caps) -> ok.
 
--spec(get_caps(emqx_zone:zone()) -> caps()).
-get_caps(Zone) ->
-    maps:map(fun(Cap, Def) -> emqx_zone:get_env(Zone, Cap, Def) end, ?DEFAULT_CAPS).
+default_caps() ->
+    ?DEFAULT_CAPS.
 
--spec(get_caps(emqx_zone:zone(), publish|subscribe) -> caps()).
-get_caps(Zone, publish) ->
-    filter_caps(?PUBCAP_KEYS, get_caps(Zone));
-get_caps(Zone, subscribe) ->
-    filter_caps(?SUBCAP_KEYS, get_caps(Zone)).
-
--spec(get_caps(emqx_zone:zone(), atom(), term()) -> term()).
 get_caps(Zone, Cap, Def) ->
     emqx_zone:get_env(Zone, Cap, Def).
 
+get_caps(Zone, publish) ->
+    with_env(Zone, '$mqtt_pub_caps',
+             fun() ->
+                 filter_caps(?PUBCAP_KEYS, get_caps(Zone))
+             end);
+
+get_caps(Zone, subscribe) ->
+    with_env(Zone, '$mqtt_sub_caps',
+             fun() ->
+                 filter_caps(?SUBCAP_KEYS, get_caps(Zone))
+             end).
+
+get_caps(Zone) ->
+    with_env(Zone, '$mqtt_caps',
+             fun() ->
+                maps:map(fun(Cap, Def) ->
+                    emqx_zone:get_env(Zone, Cap, Def)
+                end, ?DEFAULT_CAPS)
+             end).
+
 filter_caps(Keys, Caps) ->
     maps:filter(fun(Key, _Val) -> lists:member(Key, Keys) end, Caps).
 
 -spec(default() -> caps()).
 default() -> ?DEFAULT_CAPS.
 
+with_env(Zone, Key, InitFun) ->
+    case emqx_zone:get_env(Zone, Key) of
+        undefined -> Caps = InitFun(),
+                     ok = emqx_zone:set_env(Zone, Key, Caps),
+                     Caps;
+        ZoneCaps  -> ZoneCaps
+    end.

test/emqx_mqtt_caps_SUITE.erl

@@ -28,9 +28,8 @@ t_check_pub(_) ->
     PubCaps = #{max_qos_allowed => ?QOS_1,
                 retain_available => false
                },
-    lists:foreach(fun({Key, Val}) ->
-        ok = emqx_zone:set_env(zone, Key, Val)
-    end, maps:to_list(PubCaps)),
+    emqx_zone:set_env(zone, '$mqtt_pub_caps', PubCaps),
+    timer:sleep(50),
     ok = emqx_mqtt_caps:check_pub(zone, #{qos => ?QOS_1,
                                           retain => false}),
     PubFlags1 = #{qos => ?QOS_2, retain => false},
@@ -39,9 +38,7 @@ t_check_pub(_) ->
     PubFlags2 = #{qos => ?QOS_1, retain => true},
     ?assertEqual({error, ?RC_RETAIN_NOT_SUPPORTED},
                  emqx_mqtt_caps:check_pub(zone, PubFlags2)),
-    lists:foreach(fun({Key, _Val}) ->
-        true = emqx_zone:unset_env(zone, Key)
-    end, maps:to_list(PubCaps)).
+    emqx_zone:unset_env(zone, '$mqtt_pub_caps').
 
 t_check_sub(_) ->
     SubOpts = #{rh  => 0,
@@ -54,9 +51,8 @@ t_check_sub(_) ->
                 shared_subscription => false,
                 wildcard_subscription => false
                },
-    lists:foreach(fun({Key, Val}) ->
-        ok = emqx_zone:set_env(zone, Key, Val)
-    end, maps:to_list(SubCaps)),
+    emqx_zone:set_env(zone, '$mqtt_sub_caps', SubCaps),
+    timer:sleep(50),
     ok = emqx_mqtt_caps:check_sub(zone, <<"topic">>, SubOpts),
     ?assertEqual({error, ?RC_TOPIC_FILTER_INVALID},
                  emqx_mqtt_caps:check_sub(zone, <<"a/b/c/d">>, SubOpts)),
@@ -64,6 +60,4 @@ t_check_sub(_) ->
                  emqx_mqtt_caps:check_sub(zone, <<"+/#">>, SubOpts)),
     ?assertEqual({error, ?RC_SHARED_SUBSCRIPTIONS_NOT_SUPPORTED},
                  emqx_mqtt_caps:check_sub(zone, <<"topic">>, SubOpts#{share => true})),
-    lists:foreach(fun({Key, _Val}) ->
-        true = emqx_zone:unset_env(zone, Key)
-    end, maps:to_list(SubCaps)).
+    emqx_zone:unset_env(zone, '$mqtt_pub_caps').