fix(sessions): use static session pool

This should prevent nested transactions after checking out the session
from the pool.

Author: hauleth

lib/mongo/protocol.ex

@@ -172,42 +172,16 @@ defmodule Mongo.Protocol do
   end
 
   @impl DBConnection
-  def handle_begin(opts, %{topology_pid: pid} = state) do
-    with {:ok, session} <- Mongo.start_session(pid, opts),
-         :ok <- Mongo.Session.start_transaction(session, opts) do
-      {:ok, :ok, %{state | session: session}}
-    else
-      _ ->
-        {:idle, state}
-    end
-  end
+  def handle_begin(_opts, state), do: {:ok, nil, state}
 
   @impl DBConnection
-  def handle_commit(_opts, state) do
-    with false <- is_nil(state.session),
-         :ok <- Mongo.Session.commit_transaction(state.session),
-         :ok <- Mongo.Session.end_session(state.session) do
-      {:ok, :ok, %{state | session: nil}}
-    else
-      _ -> {:idle, state}
-    end
-  end
+  def handle_commit(_opts, state), do: {:ok, nil, state}
 
   @impl DBConnection
-  def handle_rollback(_opts, state) do
-    with false <- is_nil(state.session),
-         :ok <- Mongo.Session.abort_transaction(state.session),
-         :ok <- Mongo.Session.end_session(state.session) do
-      {:ok, :ok, %{state | session: nil}}
-    else
-      _ -> {:idle, state}
-    end
-  end
+  def handle_rollback(_opts, state), do: {:ok, nil, state}
 
   @impl DBConnection
-  def handle_close(_query, _opts, state) do
-    {:ok, nil, state}
-  end
+  def handle_close(_query, _opts, state), do: {:ok, nil, state}
 
   @impl DBConnection
   def handle_deallocate(_query, _cursor, _opts, state) do

lib/mongo/session.ex

@@ -1,5 +1,5 @@
 defmodule Mongo.Session do
-  @enforce_keys [:id, :pid]
+  @enforce_keys [:session, :pid]
   defstruct @enforce_keys ++
               [
                 :ref,
@@ -9,11 +9,7 @@ defmodule Mongo.Session do
                 operation_time: nil,
                 causal_consistency: true,
                 retry_writes: false,
-                txn: %{
-                  seq: 0,
-                  read_concern: nil,
-                  write_concern: nil
-                }
+                active_txn: nil
               ]
 
   @opaque session :: pid()
@@ -24,8 +20,8 @@ defmodule Mongo.Session do
   defmodule Supervisor do
     @moduledoc false
 
-    def start_child(conn, id, opts) do
-      DynamicSupervisor.start_child(__MODULE__, {Mongo.Session, {conn, id, opts, self()}})
+    def start_child(conn, session, opts, parent) do
+      DynamicSupervisor.start_child(__MODULE__, {Mongo.Session, {conn, session, opts, parent}})
     end
 
     def child_spec(_) do
@@ -66,9 +62,9 @@ defmodule Mongo.Session do
   """
   @spec end_session(session()) :: :ok
   def end_session(pid) do
-    Mongo.SessionPool.checkin(pid)
-
-    :ok
+    with {:ok, id, txn} <- :gen_statem.call(pid, :end_session) do
+      Mongo.SessionPool.checkin(id, txn)
+    end
   end
 
   @doc """
@@ -105,25 +101,6 @@ defmodule Mongo.Session do
     :gen_statem.call(pid, {:advance_operation_time, timestamp})
   end
 
-  @doc false
-  def set_options(pid, opts) do
-    causal_consistency = Keyword.get(opts, :causal_consistency, true)
-    read_concern = Keyword.get(opts, :read_concern, %{})
-    read_preference = Keyword.get(opts, :read_preference)
-    retry_writes = Keyword.get(opts, :retry_writes, true)
-    write_concern = Keyword.get(opts, :write_concern)
-
-    state = %{
-      causal_consistency: causal_consistency,
-      read_concern: read_concern,
-      read_preference: read_preference,
-      retry_writes: retry_writes,
-      write_concern: write_concern
-    }
-
-    :gen_statem.call(pid, {:set_options, state})
-  end
-
   @doc false
   def update_session(doc, nil), do: doc
 
@@ -151,15 +128,15 @@ defmodule Mongo.Session do
   @outside_txn @states -- @in_txn
 
   @doc false
-  def child_spec({topology_pid, id, opts, parent}) do
+  def child_spec({topology_pid, session, opts, parent}) do
     causal_consistency = Keyword.get(opts, :causal_consistency, true)
     read_concern = Keyword.get(opts, :read_concern, %{})
     read_preference = Keyword.get(opts, :read_preference)
     retry_writes = Keyword.get(opts, :retry_writes, true)
     write_concern = Keyword.get(opts, :write_concern)
 
     state = %__MODULE__{
-      id: id,
+      session: session,
       pid: topology_pid,
       causal_consistency: causal_consistency,
       read_concern: read_concern,
@@ -200,25 +177,27 @@ defmodule Mongo.Session do
   end
 
   # Start new transaction if there isn't one already.
-  def handle_event({:call, from}, {:start_transaction, opts}, state, %{txn: %{seq: seq}} = data)
+  def handle_event({:call, from}, {:start_transaction, opts}, state, %{session: session} = data)
       when state in @outside_txn do
     write_concern = Keyword.get(opts, :write_concern, data.write_concern)
     read_concern = Keyword.get(opts, :read_concern, data.read_concern)
 
     txn = %{
-      seq: seq + 1,
       write_concern: write_concern,
       read_concern: read_concern
     }
 
-    {:next_state, :transaction_started, struct(data, txn: txn), {:reply, from, :ok}}
+    session = Map.update!(session, :txn, & &1 + 1)
+
+    {:next_state, :transaction_started, struct(data, session: session, active_txn: txn),
+     {:reply, from, :ok}}
   end
 
   # Add session information to the query metadata.
   def handle_event({:call, from}, {:add_session, query}, :transaction_started, data) do
     %{
-      txn: %{
-        seq: seq,
+      session: %{txn: seq, id: id} = session,
+      active_txn: %{
         read_concern: read_concern,
         write_concern: write_concern
       }
@@ -227,33 +206,39 @@ defmodule Mongo.Session do
     new_query =
       query
       |> Keyword.new()
-      |> add_option(:lsid, data.id)
+      |> add_option(:lsid, %{id: id})
       |> add_option(:txnNumber, {:long, seq})
       |> add_option(:startTransaction, true)
       |> add_option(:autocommit, false)
       |> add_option(:writeConcern, write_concern)
       |> add_option(:readConcern, read_concern)
       |> set_read_concern(data.operation_time, data.causal_consistency)
 
-    {:next_state, :in_transaction, data, {:reply, from, {:ok, new_query}}}
+    session = Map.put(session, :last_use, :erlang.monotonic_time())
+
+    {:next_state, :in_transaction, struct(data, session: session),
+     {:reply, from, {:ok, new_query}}}
   end
 
   def handle_event({:call, from}, {:add_session, query}, :in_transaction, data) do
     new_query =
       query
       |> Keyword.new()
       |> Keyword.merge(
-        lsid: data.id,
-        txnNumber: {:long, data.txn},
+        lsid: %{id: data.session.id},
+        txnNumber: {:long, data.session.txn},
         autocommit: false
       )
 
+    session = Map.put(data.session, :last_use, :erlang.monotonic_time())
+    data = struct(data, session: session)
+
     case Keyword.get(new_query, :read_preference, %{mode: :primary}) do
       %{mode: :primary} ->
-        {:keep_state_and_data, {:reply, from, {:ok, new_query}}}
+        {:keep_state, data, {:reply, from, {:ok, new_query}}}
 
       %{mode: mode} ->
-        {:keep_state_and_data,
+        {:keep_state, data,
          {:reply, from,
           {:error,
            Mongo.Error.exception(message: "Read preference must be primary, not: #{mode}")}}}
@@ -267,7 +252,7 @@ defmodule Mongo.Session do
       new_query =
         query
         |> Keyword.new()
-        |> add_option(:lsid, data.id)
+        |> add_option(:lsid, data.session.id)
         |> set_read_concern(data.operation_time, data.causal_consistency)
 
       {:next_state, :no_transaction, data, {:reply, from, {:ok, new_query}}}
@@ -303,8 +288,13 @@ defmodule Mongo.Session do
 
   # Finish session by ending process (for further "closing" see `terminate/3`
   # handler.
-  def handle_event({:call, from}, :end_session, _state, _data) do
-    {:stop_and_reply, :normal, {:reply, from, :ok}}
+  def handle_event({:call, from}, :end_session, state, %{session: session} = data) do
+    _ =
+      if state == :in_transaction do
+        try_run_txn_command(data, :abortTransaction)
+      end
+
+    {:stop_and_reply, :normal, [{:reply, from, {:ok, session}}]}
   end
 
   def handle_event({:call, from}, {:advance_operation_time, timestamp}, _state, data) do
@@ -318,10 +308,6 @@ defmodule Mongo.Session do
     end
   end
 
-  def handle_event({:call, from}, {:set_options, opts}, _state, data) do
-    {:keep_state, struct(data, opts), {:reply, from, :ok}}
-  end
-
   # If parent process died before session then stop process and handle aborting
   # sessions in `terminate/3` handler.
   def handle_event(:info, {:DOWN, ref, :process, _pid, _reason}, _state, %{ref: ref}) do
@@ -339,12 +325,13 @@ defmodule Mongo.Session do
   @impl :gen_statem
   # Abort all pending transactions if there any and end session itself.
   def terminate(_reason, state, %{pid: pid} = data) do
-    if state == :in_transaction do
-      _ = try_run_txn_command(data, :abortTransaction)
-    end
+    _ =
+      if state == :in_transaction do
+        try_run_txn_command(data, :abortTransaction)
+      end
 
     query = %{
-      endSessions: [data.id]
+      endSessions: [data.session.id]
     }
 
     with {:ok, conn, _, _} <- Mongo.select_server(pid, :write, []),
@@ -380,9 +367,9 @@ defmodule Mongo.Session do
     query =
       [
         {command, 1},
-        lsid: state.id,
+        lsid: %{id: state.session.id},
         autocommit: false,
-        txnNumber: {:long, state.txn.seq}
+        txnNumber: {:long, state.session.txn}
       ]
       |> add_option(:writeConcern, state.write_concern)