Write update_last_write_timestamp after request

We need to update using the timestamp from the end of the request, not
the start. For example, if a request spends 5+ seconds writing, we still
want to wait another 5 seconds for replication lag.

Since we now run the update after we yield, we need to use ensure to
make sure we update the timestamp even if there is an exception.

2 files changed, 43 insertions, 2 deletions

diff --git a/activerecord/lib/active_record/middleware/database_selector/resolver.rb b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
index 1c056e2156..930ddf900e 100644
--- a/activerecord/lib/active_record/middleware/database_selector/resolver.rb
+++ b/activerecord/lib/active_record/middleware/database_selector/resolver.rb
@@ -64,8 +64,12 @@ module ActiveRecord
           def write_to_primary(&blk)
             ActiveRecord::Base.connected_to(role: :writing) do
               instrumenter.instrument("database_selector.active_record.wrote_to_primary") do
-                resolver.update_last_write_timestamp
-                yield
+                begin
+                  ret = yield
+                ensure
+                  resolver.update_last_write_timestamp
+                end
+                ret
               end
             end
           end
diff --git a/activerecord/test/cases/database_selector_test.rb b/activerecord/test/cases/database_selector_test.rb
index f961fa754b..6142e223ce 100644
--- a/activerecord/test/cases/database_selector_test.rb
+++ b/activerecord/test/cases/database_selector_test.rb
@@ -58,6 +58,43 @@ module ActiveRecord
       assert called
     end
 
+    def test_write_to_primary
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session)
+
+      # Session should start empty
+      assert_nil @session_store[:last_write]
+
+      called = false
+      resolver.write do
+        assert ActiveRecord::Base.connected_to?(role: :writing)
+        called = true
+      end
+      assert called
+
+      # and be populated by the last write time
+      assert @session_store[:last_write]
+    end
+
+    def test_write_to_primary_with_exception
+      resolver = ActiveRecord::Middleware::DatabaseSelector::Resolver.new(@session)
+
+      # Session should start empty
+      assert_nil @session_store[:last_write]
+
+      called = false
+      assert_raises(ActiveRecord::RecordNotFound) do
+        resolver.write do
+          assert ActiveRecord::Base.connected_to?(role: :writing)
+          called = true
+          raise ActiveRecord::RecordNotFound
+        end
+      end
+      assert called
+
+      # and be populated by the last write time
+      assert @session_store[:last_write]
+    end
+
     def test_the_middleware_chooses_writing_role_with_POST_request
       middleware = ActiveRecord::Middleware::DatabaseSelector.new(lambda { |env|
         assert ActiveRecord::Base.connected_to?(role: :writing)