From c1600a99e29442c76e06dd63c17609f36e105ea7 Mon Sep 17 00:00:00 2001
From: Kevin McConnell <kevin@37signals.com>
Date: Wed, 1 Oct 2025 12:56:17 +0100
Subject: [PATCH] Populate writer and last transaction in responses

In order to control routing dynamically, we need to expose some Beamer
information to the proxy.

The `beamer_writer` value is used by the proxy to keep track of which
node should receive write requests. When it changes due to a failover,
the proxy will update after seeing a new value in this header. We
provide this to the proxy in the `X-Writer` response header.

The `beamer_last_txn` value will be used to control writer pinning.
Setting it in the cookie here is the first stage of this. The second
stage will be to catch situations where a reader gets a request where
this header value is set to a large value than the reader has seen; when
that happens if means there's a risk of reading stale data, so rather
than serve the request we should request the proxy to reproxy it back to
the writer.

We also set `X-Kamal-Target` in the response to match the
`X-Kamal-Target` that was set in the request; this lets us see which
proxy target served each request, which is useful for diagnostics.
---
 app/controllers/application_controller.rb     |  3 +-
 .../concerns/load_balancer_routing.rb         | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 app/controllers/concerns/load_balancer_routing.rb

diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index 1de45b5c9..9e381c604 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -1,5 +1,6 @@
 class ApplicationController < ActionController::Base
-  include Authentication, Authorization, CurrentRequest, CurrentTimezone, SetPlatform, TurboFlash, ViewTransitions, WriterAffinity
+  include Authentication, Authorization, CurrentRequest, CurrentTimezone, LoadBalancerRouting,
+    SetPlatform, TurboFlash, ViewTransitions, WriterAffinity
 
   stale_when_importmap_changes
   allow_browser versions: :modern, block: -> { render "errors/not_acceptable", layout: "error" }
diff --git a/app/controllers/concerns/load_balancer_routing.rb b/app/controllers/concerns/load_balancer_routing.rb
new file mode 100644
index 000000000..200aaebca
--- /dev/null
+++ b/app/controllers/concerns/load_balancer_routing.rb
@@ -0,0 +1,39 @@
+module LoadBalancerRouting
+  extend ActiveSupport::Concern
+
+  ALLOWED_PRAGMAS = %w[ beamer_primary beamer_last_txn ]
+
+  included do
+    before_action :set_target_header, :set_writer_header
+    after_action :set_transaction_cookie
+  end
+
+  private
+    def set_target_header
+        response.headers["X-Kamal-Target"] = request.headers["X-Kamal-Target"]
+    end
+
+    def set_writer_header
+      if ApplicationRecord.current_tenant.present?
+        response.headers["X-Writer"] = pragma("beamer_primary")
+      end
+    end
+
+    def set_transaction_cookie
+      unless safe_request?
+        if ApplicationRecord.current_tenant.present? && Account.sole.present?
+          cookies[:last_transaction] = { value: pragma("beamer_last_txn"), path: Account.sole.slug }
+        end
+      end
+    end
+
+    def pragma(name)
+      if ALLOWED_PRAGMAS.include?(name)
+        ApplicationRecord.connection.execute("pragma #{name}").first&.values&.first
+      end
+    end
+
+    def safe_request?
+      request.get? || request.head?
+    end
+end