EMQX消息存储MongoDB源码分析(3)

-module(emqx_backend_mongo_cli).

-behaviour(ecpool_worker).
-include("../include/emqx.hrl").
-include("../include/mongo_protocol.hrl").

-export([client_connected/2,
  subscribe_lookup/2,
  client_disconnected/2,
  message_fetch/3,
  lookup_retain/2,
  acked_delete/2,
  message_publish/3,
  message_store/3,
  message_retain/3,
  delete_retain/2,
  message_acked/2]).

-export([mongo_insert/3,
  mongo_query/3,
  mongo_insert_update/4,
  mongo_delete/3,
  record_to_msg/1,
  maps_to_binary/1]).

-export([connect/1]).

-vsn("4.2.1").

client_connected(Pool, Msg) ->
  Selector = [<<"clientid">>],
  Docs = [<<"state">>, 1, <<"node">>, node(), <<"online_at">>, erlang:system_time(millisecond), <<"offline_at">>, null],
  mongo_insert_update(Pool, <<"mqtt_client">>, feed_var(Selector, Msg), Docs),
  ok.

subscribe_lookup(Pool, Msg) ->
  Selector = [<<"clientid">>],
  case mongo_query(Pool,
    <<"mqtt_sub">>,
    feed_var(Selector, Msg))
  of
    {ok, []} -> [];
    {ok, Rows} ->
      [{maps:get(<<"topic">>, Map),
        #{qos => trunc(maps:get(<<"qos">>, Map, 0))}}
        || Map <- Rows];
    {error, Error} ->
      logger:error("Lookup subscription error: ~p", [Error]),
      []
  end.

client_disconnected(Pool, Msg) ->
  Selector = [<<"clientid">>],
  Docs = [<<"state">>,
    0,
    <<"offline_at">>,
    erlang:system_time(millisecond)],
  mongo_insert_update(Pool,
    <<"mqtt_client">>,
    feed_var(Selector, Msg),
    Docs),
  ok.

message_fetch(Pool, Msg, Opts) ->
  Topic = proplists:get_value(topic, Msg),
  ClientId = proplists:get_value(clientid, Msg),
  Selector0 = [<<"clientid">>, <<"topic">>],
  case mongo_query_for_field(Pool,
    <<"mqtt_acked">>,
    feed_var(Selector0, Msg),
    <<"mongo_id">>)
  of
    undefined ->
      Selector1 = #{<<"$query">> => {<<"topic">>, Topic},
        <<"$orderby">> => #{<<"id">> => -1}},
      Mid = case ecpool:with_client(Pool,
        fun (C) ->
          mongo_api:find(C,
            <<"mqtt_msg">>,
            Selector1,
            #{},
            0,
            1)
        end)
            of
              [] -> 0;
              {ok, RPid} when is_pid(RPid) ->
                [Result] = mc_cursor:next_batch(RPid),
                mc_cursor:close(RPid),
                maps:get(<<"id">>, Result);
              _ -> 0
            end,
      mongo_insert(Pool,
        <<"mqtt_acked">>,
        [<<"clientid">>,
          ClientId,
          <<"topic">>,
          Topic,
          <<"mongo_id">>,
          Mid]),
      [];
    AckId ->
      Selector2 = case proplists:get_value(time_range, Opts)
                  of
                    undefined ->
                      [<<"id">>,
                        {<<"$gt">>, AckId},
                        <<"topic">>,
                        Topic];
                    TimeRange ->
                      Time = (erlang:system_time(seconds) - TimeRange)
                        * 1000,
                      [<<"id">>,
                        {<<"$gt">>, AckId},