refactoring select_servers function for supporting driver sessions

Author: zookzook

lib/mongo.ex

@@ -290,13 +290,32 @@ defmodule Mongo do
 
     opts = Keyword.drop(opts, ~w(bypass_document_validation max_time projection return_document sort upsert collation)a)
 
-    with {:ok, conn, _, _} <- select_server(topology_pid, :write, opts),
-         {:ok, doc} <- exec_command(conn, cmd, opts) do
+    with {:ok, doc} <- call_command_for_type(topology_pid, cmd, :write, tops) do
         {:ok, doc["value"]}
     end
 
   end
 
+  def call_command_for_type(topology_pid, cmd, :write, opts) do
+    with {:ok, conn, _, _, session_id} <- select_server(topology_pid, :write, opts),
+         {:ok, doc} <- exec_command(conn, update_session_id(cmd, session_id), opts) do
+      {:ok, doc}
+    end
+  end
+  def call_command_for_type(topology_pid, cmd, :read, opts) do
+    with {:ok, conn, slave_ok, _, session_id} <- select_server(topology_pid, :read, opts),
+         {:ok, doc} <- exec_command(conn, update_session_id(cmd, session_id), opts) do
+      {:ok, doc}
+    end
+  end
+
+  defp update_session_id(cmd, nil) do
+    cmd
+  end
+  defp update_session_id(cmd, session_id) do
+    Keyword.merge(cmd, [lsid: session_id])
+  end
+
   @doc """
   Finds a document and replaces it.
 
@@ -1002,13 +1021,13 @@ defmodule Mongo do
     the options for the following request because you are requesting a secondary server.
   """
     def select_server(topology_pid, type, opts \\ []) do
-    with {:ok, servers, slave_ok, mongos?} <- select_servers(topology_pid, type, opts) do
+    with {:ok, servers, slave_ok, mongos?, session_id} <- select_servers(topology_pid, type, opts) do
       if Enum.empty? servers do
-        {:ok, nil, slave_ok, mongos?}
+        {:ok, nil, slave_ok, mongos?, session_id}
       else
         with {:ok, connection} <- servers |> Enum.take_random(1) |> Enum.at(0)
                                           |> get_connection(topology_pid) do
-          {:ok, connection, slave_ok, mongos?}
+          {:ok, connection, slave_ok, mongos?, session_id}
         end
       end
     end
@@ -1022,24 +1041,34 @@ defmodule Mongo do
   #       connection.
   defp select_servers(topology_pid, type, opts, start_time) do
     topology = Topology.topology(topology_pid)
-    with {:ok, servers, slave_ok, mongos?} <- TopologyDescription.select_servers(topology, type, opts) do
-      case Enum.empty?(servers) do
-        true ->
-          case Topology.wait_for_connection(topology_pid, @sel_timeout, start_time) do
-            {:ok, _servers} -> select_servers(topology_pid, type, opts, start_time)
-            {:error, :selection_timeout} = error -> error
-          end
-        false -> {:ok, servers, slave_ok, mongos?}
-      end
+
+    case TopologyDescription.select_servers(topology, type, opts) do
+
+      :empty ->
+        case Topology.wait_for_connection(topology_pid, @sel_timeout, start_time) do
+          {:ok, _servers}                      -> select_servers(topology_pid, type, opts, start_time) ##todo wait a little
+          {:error, :selection_timeout} = error -> error
+        end
+
+      {:ok, result} -> result
+
+      error         -> error
     end
+
+#    with {:ok, servers, slave_ok, mongos?, session_id} <- TopologyDescription.select_servers(topology, type, opts) do
+#      case servers do
+#        [] ->
+#          case Topology.wait_for_connection(topology_pid, @sel_timeout, start_time) do
+#            {:ok, _servers}                      -> select_servers(topology_pid, type, opts, start_time) ##todo wait a little
+#            {:error, :selection_timeout} = error -> error
+#          end
+#        _full -> {:ok, servers, slave_ok, mongos?, session_id}
+#      end
+#    end
   end
 
   defp get_connection(nil, _pid), do: {:ok, nil}
-  defp get_connection(server, pid) do
-    with {:ok, connection} <- Topology.connection_for_address(pid, server) do
-      {:ok, connection}
-    end
-  end
+  defp get_connection(server, pid), do: Topology.connection_for_address(pid, server)
 
   defp modifier_docs([{key, _}|_], type), do: key |> key_to_string |> modifier_key(type)
   defp modifier_docs(map, _type) when is_map(map) and map_size(map) == 0,  do: :ok

lib/mongo/monitor.ex

@@ -70,10 +70,6 @@ defmodule Mongo.Monitor do
            |> Keyword.put(:idle_interval, 5_000)
 
     with {:ok, pid} <- DBConnection.start_link(Mongo.MongoDBConnection, opts) do
-
-      ## we start after one second
-      Process.send_after(self(), :update, 1_000)
-
       {:ok, %{
         connection_pid: pid,                            ## our connection pid to the mongodb server
         topology_pid: topology_pid,                     ## the topology_pid to which we report
@@ -99,6 +95,7 @@ defmodule Mongo.Monitor do
   """
   def connected(_connection, me, topology_pid) do
     Topology.monitor_connected(topology_pid, me)
+    GenServer.cast(me, :update)
   end
 
   @doc """

lib/mongo/server_description.ex

@@ -22,7 +22,8 @@ defmodule Mongo.ServerDescription do
     set_version: non_neg_integer | nil,
     election_id: BSON.ObjectId.t | nil,
     primary: String.t | nil,
-    last_update_time: non_neg_integer
+    last_update_time: non_neg_integer,
+    logical_session_timeout: non_neg_integer
   }
 
   def defaults(map \\ %{}) do
@@ -44,7 +45,8 @@ defmodule Mongo.ServerDescription do
       set_version: nil,
       election_id: nil,
       primary: nil,
-      last_update_time: 0
+      last_update_time: 0,
+      logical_session_timeout: 30
     }, map)
   end
 
@@ -74,7 +76,8 @@ defmodule Mongo.ServerDescription do
       set_name: is_master_reply["setName"],
       set_version: is_master_reply["setVersion"],
       election_id: is_master_reply["electionId"],
-      primary: is_master_reply["primary"]
+      primary: is_master_reply["primary"],
+      logical_session_timeout: is_master_reply["logicalSessionTimeoutMinutes"]
     }
   end
 
@@ -84,7 +87,7 @@ defmodule Mongo.ServerDescription do
   defp determine_server_type(%{"isreplicaset" => true}), do: :rs_ghost
   defp determine_server_type(%{"setName" => set_name} = is_master_reply) when set_name != nil do
     case is_master_reply do
-      %{"ismaster" => true}    ->  :rs_primary
+      %{"ismaster" => true}    -> :rs_primary
       %{"secondary" => true}   -> :rs_secondary
       %{"arbiterOnly" => true} -> :rs_arbiter
       _                        -> :rs_other

lib/mongo/topology.ex

@@ -7,6 +7,7 @@ defmodule Mongo.Topology do
   alias Mongo.TopologyDescription
   alias Mongo.ServerDescription
   alias Mongo.Monitor
+  alias Mongo.Session.SessionPool
 
   # https://github.com/mongodb/specifications/blob/master/source/server-discovery-and-monitoring/server-discovery-and-monitoring.rst#heartbeatfrequencyms-defaults-to-10-seconds-or-60-seconds
   @heartbeat_frequency_ms 10_000
@@ -49,8 +50,8 @@ defmodule Mongo.Topology do
   defp wait_for_connection(pid, timeout) when timeout >= 0 do
     try do
       case GenServer.call(pid, :wait_for_connection, timeout) do
-        {:new_connection, server} ->  {:ok, [server]}
-        {:connected, servers} -> {:ok, servers}
+        {:new_connection, server} -> {:ok, [server]}
+        {:connected, servers}     -> {:ok, servers}
       end
     catch
       :exit, {:timeout, _} -> {:error, :selection_timeout}
@@ -100,10 +101,10 @@ defmodule Mongo.Topology do
             opts: opts,
             monitors: %{},
             connection_pools: %{},
-            ## session_pid: nil
+            session_pool: nil,
             waiting_pids: []
           }
-          |> reconcile_servers()
+          |> update_monitor()
         {:ok, state}
     end
   end
@@ -153,7 +154,7 @@ defmodule Mongo.Topology do
   end
 
   def handle_cast(:reconcile, state) do
-    new_state = reconcile_servers(state)
+    new_state = update_monitor(state)
     {:noreply, new_state}
   end
   def handle_cast({:disconnect, :monitor, host}, state) do
@@ -168,8 +169,7 @@ defmodule Mongo.Topology do
   def handle_cast({:connected, monitor_pid}, state) do
     monitor = Enum.find(state.monitors, fn {_key, value} -> value == monitor_pid end)
     new_state = case monitor do
-      nil ->
-        state
+      nil -> state
       {host, ^monitor_pid} ->
         arbiters = fetch_arbiters(state)
         if host in arbiters do
@@ -186,7 +186,7 @@ defmodule Mongo.Topology do
           Enum.each(state.waiting_pids, fn from ->
             GenServer.reply(from, {:new_connection, host})
           end)
-          %{ state | connection_pools: connection_pools, waiting_pids: [] }
+          %{ state | connection_pools: connection_pools, waiting_pids: []}
         end
     end
     {:noreply, new_state}
@@ -204,17 +204,26 @@ defmodule Mongo.Topology do
     end
   end
 
+  ##
+  # Update server description: in case of logical session the function creates a session pool for the `deployment`.
+  #
+  defp handle_server_description(state, %{:logical_session_timeout => logical_session_timeout} = server_description) do
+    state
+    |> get_and_update_in([:topology], &TopologyDescription.update(&1, server_description, length(state.seeds)))
+    |> process_events()
+    |> update_monitor()
+    |> update_session_pool(logical_session_timeout)
+  end
   defp handle_server_description(state, server_description) do
     state
     |> get_and_update_in([:topology], &TopologyDescription.update(&1, server_description, length(state.seeds)))
     |> process_events()
-    |> reconcile_servers()
+    |> update_monitor()
   end
 
   defp process_events({events, state}) do
     Enum.each(events, fn
-      {:force_check, _} = message ->
-        :ok = GenServer.cast(self(), message)
+      {:force_check, _} = message -> :ok = GenServer.cast(self(), message)
       {previous, next} ->
         if previous != next do
           :ok = Mongo.Events.notify(%ServerDescriptionChangedEvent{
@@ -230,8 +239,11 @@ defmodule Mongo.Topology do
     state
   end
 
-  ## todo: update_monitor(state)
-  defp reconcile_servers(state) do
+
+  ##
+  # update the monitor process. For new servers the function creates new monitor processes.
+  #
+  defp update_monitor(state) do
     arbiters = fetch_arbiters(state)
     old_addrs = Map.keys(state.monitors)
     # remove arbiters from connection pool as descriptions are recieved
@@ -256,6 +268,14 @@ defmodule Mongo.Topology do
     Enum.reduce(removed, state, &remove_address/2)
   end
 
+  defp update_session_pool(%{:session_pool => nil} = state, logical_session_timeout) do
+    {:ok, session_pool} = SessionPool.start_link(self(), logical_session_timeout)
+    %{ state | session_pool: session_pool}
+  end
+  defp update_session_pool(state, _logical_session_timeout) do
+    state
+  end
+
   defp maybe_reinit(%{monitors: monitors} = state) when map_size(monitors) > 0,
     do: state
   defp maybe_reinit(state) do

lib/mongo/topology_description.ex

@@ -48,40 +48,55 @@ defmodule Mongo.TopologyDescription do
     check_server_supported(topology, server_description, num_seeds)
   end
 
-  # steps 3-4
+  def checkout_session(_top, []) do
+    nil
+  end
+  def checkout_session(%{:session_pool => pool}, _servers) do
+    SessionPool.checkout(pool)
+  end
+  def checkout_session(_other, _servers) do
+    nil
+  end
+
+  def select_servers(%{:compatible => false}, _type, _opts) do
+    {:error, :invalid_wire_version}
+  end
   def select_servers(topology, type, opts \\ []) do
-    read_preference = Keyword.get(opts, :read_preference) |> ReadPreference.defaults()
-    if topology[:compatible] == false do
-      {:error, :invalid_wire_version}
-    else
-      {servers, slave_ok, mongos?} = case topology.type do
-        :unknown -> {[], false, false}
-        :single ->
-          server   = topology.servers |> Map.values |> Enum.at(0, %{type: :unknown})
-          slave_ok = type != :write and server.type != :mongos
-          {topology.servers, slave_ok, server.type == :mongos}
-        :sharded ->
-          mongos_servers = topology.servers |> Enum.filter(fn {_, server} -> server.type == :mongos end)
-          {mongos_servers, false, true}
-        _ ->
-          case type do
-            :read -> {select_replica_set_server(topology, read_preference.mode, read_preference), true, false}
-            :write ->
-              if topology.type == :replica_set_with_primary do
-                {select_replica_set_server(topology, :primary, ReadPreference.defaults), false, false}
-              else
-                {[], false, false}
-              end
-          end
-      end
+    read_preference = opts
+                      |> Keyword.get(:read_preference)
+                      |> ReadPreference.defaults()
 
-      servers =
-        for {server, _} <- servers do
-          server
+    {servers, slave_ok, mongos?} = case topology.type do
+      :unknown -> {[], false, false}
+
+      :single ->
+        server   = topology.servers
+                   |> Map.values
+                   |> Enum.at(0, %{type: :unknown})
+        slave_ok = type != :write and server.type != :mongos
+        {topology.servers, slave_ok, server.type == :mongos}
+
+      :sharded -> {mongos_servers(topology), false, true}
+      _ ->
+        case type do
+          :read -> {select_replica_set_server(topology, read_preference.mode, read_preference), true, false}
+          :write ->
+            if topology.type == :replica_set_with_primary do
+              {select_replica_set_server(topology, :primary, ReadPreference.defaults), false, false}
+            else
+              {[], false, false}
+            end
         end
-        # todo: Enum.map(elem(0)) ?
-      {:ok, servers, slave_ok, mongos?}
     end
+
+    case Enum.map(servers, fn {server, _} -> server end) do
+      []      -> :empty
+      servers -> {:ok, servers, slave_ok, mongos?, checkout_session(topology)}
+    end
+  end
+
+  defp mongos_servers(topology) do
+    Enum.filter(topology.servers, fn {_, server} -> server.type == :mongos end)
   end
 
   ## Private Functions