Include transaction pinning support

See https://github.com/basecamp/fizzy/pull/1656

saas/lib/fizzy/saas/engine.rb

@@ -1,3 +1,6 @@
+require_relative "metrics"
+require_relative "transaction_pinning"
+
 module Fizzy
   module Saas
     class Engine < ::Rails::Engine
@@ -10,6 +13,10 @@ module Fizzy
         end
       end
 
+      initializer "fizzy_saas.transaction_pinning" do |app|
+        app.config.middleware.insert_after(ActiveRecord::Middleware::DatabaseSelector, TransactionPinning::Middleware)
+      end
+
       config.to_prepare do
         Queenbee::Subscription.short_names = Subscription::SHORT_NAMES
         Queenbee::ApiToken.token = Rails.application.credentials.dig(:queenbee_api_token)

saas/lib/fizzy/saas/metrics.rb

@@ -0,0 +1,13 @@
+Yabeda.configure do
+  SHORT_HISTOGRAM_BUCKETS = [ 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5 ]
+
+  group :fizzy do
+    counter :replica_stale,
+      comment: "Number of requests served from a stale replica"
+
+    histogram :replica_wait,
+      unit: :seconds,
+      comment: "Time spent waiting for replica to catch up with transaction",
+      buckets: SHORT_HISTOGRAM_BUCKETS
+  end
+end

saas/lib/fizzy/saas/transaction_pinning.rb

@@ -0,0 +1,65 @@
+module TransactionPinning
+  class Middleware
+    SESSION_KEY = :last_txn
+    DEFAULT_MAX_WAIT = 0.25
+
+    def initialize(app)
+      @app = app
+      @timeout = Rails.application.config.x.transaction_pinning&.timeout&.to_f || DEFAULT_MAX_WAIT
+    end
+
+    def call(env)
+      request = ActionDispatch::Request.new(env)
+      replica_metrics = {}
+
+      if ApplicationRecord.current_role == :reading
+        wait_for_replica_catchup(request, replica_metrics)
+      end
+
+      status, headers, body = @app.call(env)
+      headers.merge!(replica_metrics.transform_values(&:to_s))
+
+      if ApplicationRecord.current_role == :writing
+        capture_transaction_id(request)
+      end
+
+      [ status, headers, body ]
+    end
+
+    private
+      def wait_for_replica_catchup(request, replica_metrics)
+        if last_txn = request.session[SESSION_KEY].presence
+          has_transaction = tracking_replica_wait_time(replica_metrics) do
+            replica_has_transaction(last_txn)
+          end
+
+          unless has_transaction
+            Yabeda.fizzy.replica_stale.increment
+            replica_metrics["X-Replica-Stale"] = true
+          end
+        end
+      end
+
+      def capture_transaction_id(request)
+        request.session[SESSION_KEY] = ApplicationRecord.connection.show_variable("global.gtid_executed")
+      end
+
+      def replica_has_transaction(txn)
+        sql = ApplicationRecord.sanitize_sql_array([ "SELECT WAIT_FOR_EXECUTED_GTID_SET(?, ?)", txn, @timeout ])
+        ApplicationRecord.connection.select_value(sql) == 0
+      rescue => e
+        Sentry.capture_exception(e, extra: { gtid: txn })
+        true # Treat as if we're up to date, since we don't know
+      end
+
+      def tracking_replica_wait_time(replica_metrics)
+        started_at = Time.current
+
+        Yabeda.fizzy.replica_wait.measure do
+          yield
+        end.tap do
+          replica_metrics["X-Replica-Wait"] = Time.current - started_at
+        end
+      end
+  end
+end