updated bulk write with session support, added better logging

Author: zookzook

config/config.exs

@@ -8,6 +8,12 @@ use Mix.Config
 # if you want to provide default values for your application for third-
 # party users, it should be done in your mix.exs file.
 
+config :logger, :console,
+  level: :debug,
+  truncate: 1024,
+  format: "$time [$level] $message ($metadata)\n\n",
+  metadata: [:module, :function, :line]
+
 # Sample configuration:
 #
 #     config :logger, :console,

lib/mongo.ex

@@ -46,11 +46,12 @@ defmodule Mongo do
       interval, the operation returns an error
   """
 
+  require Logger
+
   use Bitwise
   use Mongo.Messages
   alias Mongo.Query
   alias Mongo.ReadPreference
-  alias Mongo.TopologyDescription
   alias Mongo.Topology
   alias Mongo.UrlParser
   alias Mongo.Session.ServerSession
@@ -304,27 +305,28 @@ defmodule Mongo do
   session pool for reuse.
   """
   def issue_command(topology_pid, cmd, :write, opts) do
+
+    Logger.debug("issue_command: write #{inspect cmd}")
+
     with {:ok, conn, _, _, session_id} <- Topology.select_server(topology_pid, :write, opts),
          {:ok, doc} <- exec_command(conn, update_session_id(cmd, session_id), opts),
-         ok <- checkin_session_id(topology_pid, cmd, session_id) do
+         :ok <- checkin_session_id(topology_pid, cmd, session_id) do
       {:ok, conn, doc}
     else
       {:new_connection, _server} ->
-        IO.puts "too fast, call it again"
         issue_command(topology_pid, cmd, :write, opts)
     end
   end
   def issue_command(topology_pid, cmd, :read, opts) do
 
-    IO.puts "issue_command: read #{inspect cmd}"
+    Logger.debug("issue_command: read #{inspect cmd}")
     with {:ok, conn, slave_ok, _, session_id} <- Topology.select_server(topology_pid, :read, opts),
          opts = Keyword.put(opts, :slave_ok, slave_ok),
          {:ok, doc} <- exec_command(conn, update_session_id(cmd, session_id), opts),
-          ok <- checkin_session_id(topology_pid, cmd, session_id) do
+          :ok <- checkin_session_id(topology_pid, cmd, session_id) do
       {:ok, conn, doc}
       else
       {:new_connection, _server} ->
-        IO.puts "too fast, call it again"
         issue_command(topology_pid, cmd, :read, opts)
     end
   end
@@ -333,10 +335,10 @@ defmodule Mongo do
     cmd
   end
   defp update_session_id(cmd, %ServerSession{:session_id => session_id}) do
-    IO.puts "update_session_id #{inspect session_id}"
+    Logger.debug("update_session_id for cmd #{inspect session_id}")
     Keyword.merge(cmd, [lsid: %{id: session_id}])
   end
-  defp checkin_session_id(topology_pid, cmd, nil), do: :ok
+  defp checkin_session_id(_topology_pid, _cmd, nil), do: :ok
   defp checkin_session_id(topology_pid, cmd, session_id) do
     case Keyword.get(cmd, :lsid, :implicit) do
       :implicit -> Topology.checkin_session_id(topology_pid, session_id)
@@ -618,12 +620,11 @@ defmodule Mongo do
   end
 
   @doc false
-  ## refactor: exec_command
   @spec exec_command(pid, BSON.document, Keyword.t) :: {:ok, BSON.document | nil} | {:error, Mongo.Error.t}
   def exec_command(conn, cmd, opts) do
     action = %Query{action: :command}
 
-    IO.puts "Executing cmd #{inspect cmd}"
+    Logger.debug("Executing cmd: #{inspect cmd}")
 
     with {:ok, _cmd, doc} <- DBConnection.execute(conn, action, [cmd], defaults(opts)),
          {:ok, doc} <- check_for_error(doc) do
@@ -972,14 +973,14 @@ defmodule Mongo do
     end
   end
 
-  def start_session(top, opts \\ []) do
+  def start_session(top) do
     ## todo error code handling
     with {:ok, %{"id" => uuid, "ok" => ok}} when ok == 1 <- command(top, [startSession: 1], database: "admin") do
       uuid
     end
   end
 
-  def end_sessions(top, sessions, opts \\ []) do
+  def end_sessions(top, sessions) do
     ## todo error code handling
     with {:ok, %{"ok" => ok}} when ok == 1 <- command(top, [endSessions: sessions], database: "admin") do
       :ok

lib/mongo/bulk_write.ex

@@ -158,6 +158,8 @@ defmodule Mongo.BulkWrite do
   alias Mongo.UnorderedBulk
   alias Mongo.OrderedBulk
   alias Mongo.BulkWriteResult
+  alias Mongo.Session.ServerSession
+  alias Mongo.Topology
 
   @doc """
   Executes unordered and ordered bulk writes.
@@ -181,27 +183,39 @@ defmodule Mongo.BulkWrite do
   def write(topology_pid, %UnorderedBulk{} = bulk, opts) do
 
     write_concern = write_concern(opts)
-    with {:ok, conn, _, _} <- Mongo.select_server(topology_pid, :write, opts) do
-      one_bulk_write(conn, bulk, write_concern, opts)
+
+    with {:ok, conn, _, _, session_id} <- Topology.select_server(topology_pid, :write, opts),
+         result                        = one_bulk_write(conn, bulk, write_concern, Keyword.merge(opts, [lsid: session_id])),
+         :ok                           <- checkin_session_id(topology_pid, Keyword.get(opts, :lsid, :implicit), session_id) do
+      result
+      else
+      {:new_connection, _server} ->
+        write(topology_pid, bulk, opts)
     end
 
   end
 
-  def write(topology_pid, %OrderedBulk{coll: coll, ops: ops}, opts) do
+  def write(topology_pid, %OrderedBulk{coll: coll, ops: ops} = bulk, opts) do
 
     write_concern = write_concern(opts)
 
     empty = %BulkWriteResult{acknowledged: acknowledged(write_concern)}
 
-    with {:ok, conn, _, _} <- Mongo.select_server(topology_pid, :write, opts),
+    with {:ok, conn, _, _, session_id} <- Topology.select_server(topology_pid, :write, opts),
          {:ok, limits} <- Mongo.limits(conn),
-         max_batch_size <- limits.max_write_batch_size do
+         max_batch_size <- limits.max_write_batch_size,
+         result = ops
+                  |> get_op_sequence()
+                  |> Enum.map(fn {cmd, docs} -> one_bulk_write_operation(conn, cmd, coll, docs, write_concern, max_batch_size, Keyword.merge(opts, [lsid: session_id])) end)
+                  |> BulkWriteResult.reduce(empty),
+         :ok <- checkin_session_id(topology_pid, Keyword.get(opts, :lsid, :implicit), session_id) do
 
-      ops
-      |> get_op_sequence()
-      |> Enum.map(fn {cmd, docs} -> one_bulk_write_operation(conn, cmd, coll, docs, write_concern, max_batch_size, opts) end)
-      |> BulkWriteResult.reduce(empty)
+      result
+    else
+      {:new_connection, _server} ->
+        write(topology_pid, bulk, opts)
     end
+
   end
 
   ##
@@ -231,13 +245,23 @@ defmodule Mongo.BulkWrite do
          max_batch_size <- limits.max_write_batch_size,
          insert_result <- one_bulk_write_operation(conn, :insert, coll, inserts, write_concern, max_batch_size, opts),
          update_result <- one_bulk_write_operation(conn, :update, coll, updates, write_concern, max_batch_size, opts),
-         delete_result <- one_bulk_write_operation(conn, :delete, coll, deletes, write_concern, max_batch_size, opts) do
+         delete_result <- one_bulk_write_operation(conn, :delete, coll, deletes, write_concern, max_batch_size, opts)  do
 
       [insert_result, update_result, delete_result]
       |> BulkWriteResult.reduce(%BulkWriteResult{acknowledged: acknowledged(write_concern)})
     end
   end
 
+  defp update_session_id(cmd, nil) do
+    cmd
+  end
+  defp update_session_id(cmd, %ServerSession{:session_id => session_id}) do
+    Keyword.merge(cmd, [lsid: %{id: session_id}])
+  end
+
+  defp checkin_session_id(topology_pid, :implicit, session_id), do: Topology.checkin_session_id(topology_pid, session_id)
+  defp checkin_session_id(_, _, _), do: :ok
+
   ###
   # Executes the command `cmd` and collects the result.
   #
@@ -340,13 +364,14 @@ defmodule Mongo.BulkWrite do
     Enum.map(cmds, fn cmd -> Mongo.exec_command(conn, cmd, opts) end)
   end
 
-  defp get_insert_cmds(coll, docs, write_concern, max_batch_size, _opts) do
+  defp get_insert_cmds(coll, docs, write_concern, max_batch_size, opts) do
 
     {ids, docs} = assign_ids(docs)
 
     cmds = docs
            |> Enum.chunk_every(max_batch_size)
            |> Enum.map(fn inserts -> get_insert_cmd(coll, inserts, write_concern) end)
+           |> Enum.map(fn cmd -> update_session_id(cmd, Keyword.get(opts, :lsid)) end)
     {cmds, ids}
 
   end
@@ -364,6 +389,7 @@ defmodule Mongo.BulkWrite do
     docs
     |> Enum.chunk_every(max_batch_size)
     |> Enum.map(fn deletes -> get_delete_cmd(coll, deletes, write_concern, opts) end)
+    |> Enum.map(fn cmd -> update_session_id(cmd, Keyword.get(opts, :lsid)) end)
 
   end
 
@@ -389,6 +415,7 @@ defmodule Mongo.BulkWrite do
     docs
     |> Enum.chunk_every(max_batch_size)
     |> Enum.map(fn updates -> get_update_cmd(coll, updates, write_concern, opts) end)
+    |> Enum.map(fn cmd -> update_session_id(cmd, Keyword.get(opts, :lsid)) end)
 
   end
 

lib/mongo/monitor.ex

@@ -23,7 +23,7 @@ defmodule Mongo.Monitor do
 
   # this is not configurable because the specification says so
   # see https://github.com/mongodb/specifications/blob/master/source/server-discovery-and-monitoring/server-discovery-and-monitoring.rst#minheartbeatfrequencyms
-  @min_heartbeat_frequency_ms 500
+  # not used @min_heartbeat_frequency_ms 500
 
   def start_link(args) do
     GenServer.start_link(__MODULE__, args)

lib/mongo/topology.ex

@@ -1,14 +1,15 @@
 defmodule Mongo.Topology do
   @moduledoc false
 
+  require Logger
+
   use GenServer
   alias Mongo.Events.{ServerDescriptionChangedEvent, ServerOpeningEvent, ServerClosedEvent,
                       TopologyDescriptionChangedEvent, TopologyOpeningEvent, TopologyClosedEvent}
   alias Mongo.TopologyDescription
   alias Mongo.ServerDescription
   alias Mongo.Monitor
   alias Mongo.Session.SessionPool
-  alias Mongo.Session.ServerSession
 
   # 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
@@ -109,14 +110,6 @@ defmodule Mongo.Topology do
     })
   end
 
-  def handle_call(:topology, _from, state) do
-    {:reply, state.topology, state}
-  end
-
-  def handle_call({:connection, address}, _from, state) do
-    {:reply, Map.fetch(state.connection_pools, address), state}
-  end
-
   # see https://github.com/mongodb/specifications/blob/master/source/server-discovery-and-monitoring/server-discovery-and-monitoring.rst#updating-the-topologydescription
   def handle_cast({:server_description, server_description}, state) do
     new_state = handle_server_description(state, server_description)
@@ -226,30 +219,37 @@ defmodule Mongo.Topology do
     state
   end
 
+  def handle_call(:topology, _from, state) do
+    {:reply, state.topology, state}
+  end
+
+  def handle_call({:connection, address}, _from, state) do
+    {:reply, Map.fetch(state.connection_pools, address), state}
+  end
+
   def handle_call({:select_server, type, opts}, from, %{:topology => topology, :waiting_pids => waiting, connection_pools: pools} = state) do
     case TopologyDescription.select_servers(topology, type, opts) do
 
       :empty ->
-        IO.puts "select_servers: Empty "
+        Logger.debug("select_server: empty")
         {:noreply, %{state | waiting_pids: [from | waiting]}} ## no servers available, wait for connection
 
       {:ok, servers, slave_ok, mongos?} ->                ## found, select randomly a server and return its connection_pool
-        IO.puts "select_servers: found #{inspect servers}"
-        IO.puts "select_servers: found #{inspect pools}"
+        Logger.debug("select_server: found #{inspect servers}, pools: #{inspect pools}")
         with {:ok, connection} <- servers
                      |> Enum.take_random(1)
                      |> get_connection(state) do
 
           session_id = checkout_session(state)
-        IO.puts "connection is #{inspect connection}"
-        IO.puts "session_id is #{inspect session_id}"
 
 
+        Logger.debug("select_server: connection is #{inspect connection}, session_id is #{inspect session_id}")
+
         {:reply, {:ok, connection, slave_ok, mongos?, session_id}, state}
         end
       error ->
-        IO.puts "select_servers: error "
-        {:reply, error, state}                                 ## in case of an error, just return the error
+        Logger.debug("select_servers: #{inspect error}")
+        {:reply, error, state} ## in case of an error, just return the error
     end
   end
 
@@ -292,12 +292,11 @@ defmodule Mongo.Topology do
   end
 
   defp update_session_pool(%{:session_pool => nil} = state, logical_session_timeout) do
-    IO.puts "creating session pool"
+    Logger.debug("Creating session pool")
     {: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
-    IO.puts "no session pool created"
     state
   end
 

lib/mongo/topology_description.ex

@@ -8,7 +8,6 @@ defmodule Mongo.TopologyDescription do
 
   alias Mongo.ServerDescription
   alias Mongo.ReadPreference
-  alias Mongo.Session.SessionPool
 
   # see https://github.com/mongodb/specifications/blob/master/source/server-discovery-and-monitoring/server-discovery-and-monitoring.rst#topologydescription
   @type type :: :unknown | :single | :replica_set_no_primary | :replica_set_with_primary | :sharded
@@ -56,10 +55,11 @@ defmodule Mongo.TopologyDescription do
   * slave_ok:
   * mongod?:
   """
+  def select_servers(topology, type, opts \\ [])
   def select_servers(%{:compatible => false}, _type, _opts) do
     {:error, :invalid_wire_version}
   end
-  def select_servers(topology, type, opts \\ []) do
+  def select_servers(topology, type, opts) do
 
     read_preference = opts
                       |> Keyword.get(:read_preference)