EMQXメッセージストレージMongoDBソースコード分析(4)

-module(emqx_backend_mongo).

-include("../include/emqx_backend_mongo.hrl").
-include("../include/emqx.hrl").


-export([pool_name/1]).

-export([register_metrics/0, load/0, unload/0]).

-export([on_client_connected/3,
  on_subscribe_lookup/3,
  on_client_disconnected/4,
  on_message_fetch/4,
  on_retain_lookup/4,
  on_acked_delete/4,
  on_message_publish/2,
  on_message_store/2,
  on_message_retain/2,
  on_retain_delete/2,
  on_message_acked/3]).


pool_name(Pool) ->
  list_to_atom(lists:concat([emqx_backend_mongo, '_', Pool])).

register_metrics() ->
  [emqx_metrics:new(MetricName)
    || MetricName
    <- ['backend.mongodb.client_connected',
      'backend.mongodb.subscribe_lookup',
      'backend.mongodb.client_disconnected',
      'backend.mongodb.message_fetch',
      'backend.mongodb.retain_lookup',
      'backend.mongodb.acked_delete',
      'backend.mongodb.message_publish',
      'backend.mongodb.message_store',
      'backend.mongodb.message_retain',
      'backend.mongodb.retain_delete',
      'backend.mongodb.message_acked']].

load() ->
  HookList =
    parse_hook(application:get_env(emqx_backend_mongo,hooks,[])),
  lists:foreach(fun ({Hook,
    Action,
    Pool,
    Filter,
    PayloadFormat,
    OfflineOpts}) ->
    case proplists:get_value(<<"function">>, Action) of
      undefined -> ok;
      Fun ->
        load_(Hook,
          b2a(Fun),
          Filter,
          OfflineOpts,
          {Filter,
            Pool,
            PayloadFormat,
            undefined})
    end
                end,
    HookList),
  io:format("~s is loaded.~n", [emqx_backend_mongo]),
  ok.

load_(Hook, Fun, Filter, OfflineOpts,
    {Filter, Pool, PayloadFormat, undefined}) ->
  load_(Hook,
    Fun,
    OfflineOpts,
    {Filter, Pool, PayloadFormat}).

load_(Hook, Fun, OfflineOpts, Params) ->
  case Hook of
    'client.connected' ->
      emqx:hook(Hook, fun emqx_backend_mongo:Fun/3, [Params]);
    'client.disconnected' ->
      emqx:hook(Hook, fun emqx_backend_mongo:Fun/4, [Params]);
    'session.subscribed' ->
      emqx:hook(Hook,
        fun emqx_backend_mongo:Fun/4,
        [erlang:append_element(Params, OfflineOpts)]);
    'session.unsubscribed' ->
      emqx:hook(Hook, fun emqx_backend_mongo:Fun/4, [Params]);
    'message.publish' ->
      emqx:hook(Hook, fun emqx_backend_mongo:Fun/2, [Params]);
    'message.acked' ->
      emqx:hook(Hook, fun emqx_backend_mongo:Fun/3, [Params])
  end.

unload() ->
  HookList =
    parse_hook(application:get_env(emqx_backend_mongo,
      hooks,
      [])),
  lists:foreach(fun ({Hook,
    Action,
    _Pool,
    _Filter,
    _PayloadFormat,
    _OfflineOpts}) ->
    case proplists:get_value(<<"function">>, Action) of
      undefined -> ok;
      Fun -> unload_(Hook, b2a(Fun))
    end
                end,
    HookList),
  io:format("~s is unloaded.~n", [emqx_backend_mongo]),
  ok.

unload_(Hook, Fun) ->
  case Hook of
    'client.connected' ->
      emqx:unhook(Hook, fun emqx_backend_mongo:Fun/3);
    'client.disconnected' ->
      emqx:unhook(Hook, fun emqx_backend_mongo:Fun/4);
    'session.subscribed' ->
      emqx:unhook(Hook, fun emqx_backend_mongo:Fun/4);
    'session.unsubscribed' ->
      emqx:unhook(Hook, fun emqx_backend_mongo:Fun/4);
    'message.publish' ->
      emqx:unhook(Hook, fun emqx_backend_mongo:Fun/2);
    'message.acked' ->
      emqx:unhook(Hook, fun emqx_backend_mongo:Fun/3)
  end.

on_client_connected(#{clientid := ClientId}, _ConnInfo,
    {Filter, Pool, _PayloadFormat}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.client_connected'),
    emqx_backend_mongo_cli:client_connected(Pool,
      [{clientid,
        ClientId}]),
    ok
              end,
    undefined,
    Filter).

on_subscribe_lookup(#{clientid := ClientId}, _ConnInfo,
    {Filter, Pool, _PayloadFormat}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.subscribe_lookup'),
    case emqx_backend_mongo_cli:subscribe_lookup(Pool,
      [{clientid,
        ClientId}])
    of
      [] -> ok;
      TopicTable ->
        self() ! {subscribe, TopicTable},
        ok
    end
              end,
    undefined,
    Filter).

on_client_disconnected(#{clientid := ClientId}, _Reason,
    _ConnInfo, {Filter, Pool, _PayloadFormat}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.client_disconnected'),
    emqx_backend_mongo_cli:client_disconnected(Pool,
      [{clientid,
        ClientId}])
              end,
    undefined,
    Filter).

on_message_fetch(#{clientid := ClientId}, Topic, Opts,
    {Filter, Pool, _PayloadFormat, OfflineOpts}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.message_fetch'),
    case maps:get(qos, Opts, 0) > 0 andalso
      maps:get(is_new, Opts, true)
    of
      true ->
        MsgList =
          emqx_backend_mongo_cli:message_fetch(Pool,
            [{clientid,
              ClientId},
              {topic,
                Topic}],
            OfflineOpts),
        [self() ! {deliver, Topic, Msg}
          || Msg <- MsgList];
      false -> ok
    end
              end,
    Topic,
    Filter).

on_retain_lookup(_Client, Topic, _Opts,
    {Filter, Pool, _PayloadFormat, _OfflineOpts}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.retain_lookup'),
    MsgList = emqx_backend_mongo_cli:lookup_retain(Pool,
      [{topic,
        Topic}]),
    [self() !
      {deliver,
        Topic,
        emqx_message:set_header(retained, true, Msg)}
      || Msg <- MsgList]
              end,
    Topic,
    Filter).

on_acked_delete(#{clientid := ClientId}, Topic, _Opts,
    {Filter, Pool, _Payload_Format}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.acked_delete'),
    emqx_backend_mongo_cli:acked_delete(Pool,
      [{clientid,
        ClientId},
        {topic, Topic}])
              end,
    Topic,
    Filter).

on_message_publish(Msg = #message{flags =
#{retain := true},
  payload = <<>>},
    _Rule) ->
  {ok, Msg};
on_message_publish(Msg = #message{qos = Qos},
    {_Filter, _Pool, _PayloadFormat})
  when Qos =:= 0 ->
  {ok, Msg};
on_message_publish(Msg0 = #message{topic = Topic},
    {Filter, Pool, PayloadFormat}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.message_publish'),
    Msg = emqx_backend_mongo_cli:message_publish(Pool,
      Msg0,
      PayloadFormat),
    {ok, Msg}
              end,
    Msg0,
    Topic,
    Filter).

on_message_store(Msg = #message{flags =
#{retain := true},
  payload = <<>>},
    _Rule) ->
  {ok, Msg};
on_message_store(Msg0 = #message{topic = Topic},
    {Filter, Pool, PayloadFormat}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.message_store'),
    Msg = emqx_backend_mongo_cli:message_store(Pool,
      Msg0,
      PayloadFormat),
    {ok, Msg}
              end,
    Msg0,
    Topic,
    Filter).

on_message_retain(Msg = #message{flags =
#{retain := false}},
    _Rule) ->
  {ok, Msg};
on_message_retain(Msg = #message{flags =
#{retain := true},
  payload = <<>>},
    _Rule) ->
  {ok, Msg};
on_message_retain(Msg0 = #message{flags =
#{retain := true},
  topic = Topic, headers = Headers0},
    {Filter, Pool, PayloadFormat}) ->
  Headers = case erlang:is_map(Headers0) of
              true -> Headers0;
              false -> #{}
            end,
  case maps:find(retained, Headers) of
    {ok, true} -> {ok, Msg0};
    _ ->
      with_filter(fun () ->
        emqx_metrics:inc('backend.mongodb.message_retain'),
        Msg =
          emqx_backend_mongo_cli:message_retain(Pool,
            Msg0,
            PayloadFormat),
        {ok, Msg}
                  end,
        Msg0,
        Topic,
        Filter)
  end;
on_message_retain(Msg, _Rule) -> {ok, Msg}.

on_retain_delete(Msg0 = #message{flags =
#{retain := true},
  topic = Topic, payload = <<>>},
    {Filter, Pool, _PayloadFormat}) ->
  with_filter(fun () ->
    emqx_metrics:inc('backend.mongodb.retain_delete'),
    Msg = emqx_backend_mongo_cli:delete_retain(Pool, Msg0),
    {ok, Msg}
              end,
    Msg0,
    Topic,
    Filter);
on_retain_delete(Msg, _Rule) -> {ok, Msg}.

on_message_acked(#{clientid := ClientId},
    #message{topic = Topic, headers = Headers0},
    {Filter, Pool, _PayloadFormat}) ->
  Headers = case erlang:is_map(Headers0) of
              true -> Headers0;
              false -> #{}
            end,
  case maps:get(mongo_id, Headers, undefined) of
    undefined -> ok;
    Id ->
      with_filter(fun () ->
        emqx_metrics:inc('backend.mongodb.message_acked'),
        emqx_backend_mongo_cli:message_acked(Pool,
          [{clientid,
            ClientId},
            {topic,
              Topic},
            {mongo_id,
              Id}])
                  end,
        Topic,
        Filter)
  end.

parse_hook(Hooks) -> parse_hook(Hooks, []).

parse_hook([], Acc) -> Acc;
parse_hook([{Hook, Item} | Hooks], Acc) ->
  Params = emqx_json:decode(Item),
  Action = proplists:get_value(<<"action">>, Params),
  Pool = proplists:get_value(<<"pool">>, Params),
  Filter = proplists:get_value(<<"topic">>, Params),
  OfflineOpts =
    parse_offline_opts(proplists:get_value(<<"offline_opts">>,
      Params,
      [])),
  PayloadFormat =
    b2a(proplists:get_value(<<"payload_format">>,
      Params,
      <<"plain_text">>)),
  parse_hook(Hooks,
    [{l2a(Hook),
      Action,
      pool_name(b2a(Pool)),
      Filter,
      PayloadFormat,
      OfflineOpts}
      | Acc]).

parse_offline_opts(OfflineOpts) ->
  parse_offline_opts(OfflineOpts, []).

parse_offline_opts([], Acc) -> Acc;
parse_offline_opts([{<<"time_range">>, TimeRange}
  | Opts],
    Acc)
  when is_binary(TimeRange) ->
  parse_offline_opts(Opts,
    [{time_range, parse_time(TimeRange)} | Acc]);
parse_offline_opts([{<<"max_returned_count">>, MaxCount}
  | Opts],
    Acc)
  when is_integer(MaxCount) ->
  parse_offline_opts(Opts,
    [{max_returned_count, MaxCount} | Acc]);
parse_offline_opts([_ | Opts], Acc) ->
  parse_offline_opts(Opts, Acc).

parse_time(TimeRange) when is_binary(TimeRange) ->
  cuttlefish_duration:parse(b2l(TimeRange), s).

with_filter(Fun, _, undefined) ->
  Fun(),
  ok;
with_filter(Fun, Topic, Filter) ->
  case emqx_topic:match(Topic, Filter) of
    true ->
      Fun(),
      ok;
    false -> ok
  end.

with_filter(Fun, _, _, undefined) -> Fun();
with_filter(Fun, Msg, Topic, Filter) ->
  case emqx_topic:match(Topic, Filter) of
    true -> Fun();
    false -> {ok, Msg}
  end.

l2a(L) -> erlang:list_to_atom(L).

b2a(B) -> erlang:binary_to_atom(B, utf8).

b2l(B) -> erlang:binary_to_list(B).

 

おすすめ

転載: blog.csdn.net/qq513036862/article/details/112962895