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.

1 files changed, 6 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