Multi-tenancy vs. Cables: Introducing Action Cable command callbacks

Solving multi-tenancy with command callbacks for Connection classes

Introducing multi-tenancy to a web application usually comes at a price: we need to (re–)design a database schema, ensure all kinds of "requests" are bound to the right tenants, and so on. Luckily, for Rails applications, we have battle-tested tools to make developers' lives easier. However, all of them focus on classic Rails components, controllers, and background jobs. Who will take care of the channels?

Execution context, or how tenant scoping is usually implemented

Multi-tenancy could be implemented in many different ways. Still, most of them include the following phases: 1) retrieving a tenant (e.g., from request properties), and 2) storing the current tenant within the current execution context.

What is execution context? We might say it's a unit of work in a web application with a clearly defined beginning and end. Web requests and background jobs are examples of execution contexts.

In Ruby, an execution context is usually connected to a single Thread or Fiber. Thus, most multi-tenancy libraries use Fiber local variables to store the current tenant information. For example, acts_as_tenant relies on the good ole request_store gem, which provides a wrapper for Thread.current and takes care of clearing the state when a request completes. All you need is to set a tenant in your controller (usually, in a before_action hook):

class ApplicationController < ActionController::Base
  before_action do
    current_account = find_account_from_request_or_whatever
    set_current_tenant(current_account)
  end
end

Easily done. We have lifecycle APIs in our controllers (action callbacks), which make injecting some logic before (or after) any unit of work pretty straightforward. We can also go one step above and rely on Rack middlewares (like Apartment does).

What about Action Cable?

First, let's think—what is the execution context for sockets? Cable connections are persistent and long-lived; they have a beginning (connect) and end (disconnect), but these are not our execution context boundaries. So, what are they then?

The way Action Cable works under the hood could give us a hint. How many concurrent clients could be handled by a Ruby server? (We're not talking about AnyCable right now). Maybe we could spawn a Thread per connection? That would quickly blow up due to high resource usage. Instead, Action Cable relies on an event loop and a Thread pool executor (i.e., a fixed number of worker threads). Whenever we need to process an incoming "message" from a client, we fetch a worker Thread from the pool and use it to process the message. And this is our unit of work (and a random Thread from the pool is our execution context). I put the "message" in quotes because we also use the pool to process connection initialization (Connection#connect) and closure (Connection#disconnect) events, which are not messages.

Now, let's take a look at the naive approach to configuring a tenant for Action Cable connections:

module ApplicationCable
  class Connection < ActionCable::Connection::Base
    identified_by :current_user, :tenant

    def connect
      self.tenant = request.subdomain
      Apartment::Tenant.switch!(tenant)
    end
  end
end

Looks similar to what we do in controllers, right? The problem here is that when the next message (say, channel subscription) is processed by this connection, we may have incorrect tenant information because the execution context has likely changed (a different Thread is processing the message). This could mess things up.

NOTE: AnyCable also uses a thread pool under the hood (part of a gRPC server).

We can probably fix this by adding before_subscribe to our ApplicationCable::Channel and calling switch!(tenant) there, too. And we should probably add after_subscribe to reset the state (otherwise, our tenant could leak into the Connection#connect and Connection#disconnect methods).

Alternatively, we can hack around the Connection class and make sure the correct tenant is set up before we enter channels:

module ApplicationCable
  class Connection < ActionCable::Connection::Base
    # ...

    # Make all channel commands tenant-aware
    def dispatch_websocket_message(*)
      using_current_tenant { super }
    end
  end
end

This is my preferred way of dealing with multi-tenancy. I believe that the connection is the right place for dealing with scoping, and that channels should not deal with it. The only problem with this approach is that it relies on Action Cable internals. And it's also incompatible with AnyCable (which doesn't use #dispatch_websocket_message), so we had to patch two methods to work with AnyCable:

module ApplicationCable
  class Connection < ActionCable::Connection::Base
    # ...

    # Make all channel commands tenant-aware
    def dispatch_websocket_message(*)
      using_current_tenant { super }
    end

    # The same override for AnyCable, which uses a different method
    def handle_channel_command(*)
      using_current_tenant { super }
    end
  end
end

The patching and duplication didn't look good to me, so I decided to fix it once and for all—let me share a bit of Rails 7.1 with you.

Action Cable around_command to the rescue

The search for a better API didn't take long: Rails is built on top of conventions, and there is no better way to extend the framework than to follow these conventions. In this particular case, I decided to go with callbacks. Every Rails developer is familiar with callbacks, right?

I'm glad to introduce command callbacks for Connection classes: before_command, after_command, and around_command. They do literally what they say: allow you to execute the code before, after, or around channel commands.

And this is how our multi-tenancy problem could be solved via the around_command callback:

module ApplicationCable
  class Connection < ActionCable::Connection::Base
    around_command :set_current_tenant

    attr_reader :tenant

    def connect
      @tenant = request.subdomain
    end

    private

    def set_current_tenant
      with_tenant(tenant) { yield }
    end
  end
end

Awesome! The only downside is that it's only available since Rails 7.1.

We made AnyCable compatible with this feature, but there's more: our Rails integration includes a backport for command callbacks for older Rails versions. Just drop anycable-rails into your Gemfile and use future Rails APIs!

Finally, let's talk about one important thing left—tests. How do we make sure our command callbacks actually work? Below, you can find the annotated snippet for RSpec:

describe ApplicationCable::Connection do
  let(:tenant) { create :tenant }
  let(:user) { create(:user, tenant:) }
  let(:chat_room) { create(:chat_room, tenant:) }

  describe "#set_current_tenant callback" do
    # We use a custom channel class just for these tests
    # to avoid depending on real channels from the app
    before do
      stub_const("TestChatChannel", Class.new(ApplicationCable::Channel) do
        cattr_accessor :found_user
        cattr_accessor :subscribed

        def subscribed
          # Use this flag to make sure we reached the #subscribed callback
          self.class.subscribed = true
          # Use this value to verify that tenant scoping has been preserved
          self.class.found_room = ChatRoom.find_by(id: params["id"])
        end
      end)
    end

    # Assume that we use cookies for authentication
    before { cookies.signed["user_id"] = user.id }

    # This is the client's command we want to process by the connection
    let(:command) do
      {
        "identifier" => {channel: "TestChatChannel", id: room.id}.to_json,
        "command" => "subscribe"
      }
    end

    specify do
      connection.handle_channel_command(command)
      expect(TestChatChannel.subscribed).to be true
      expect(TestChatChannel.found_room).to eq(room)
    end

    context "when user is from another tenant" do
      let(:user) { create(:user) }

      specify do
        connection.handle_channel_command(command)
        expect(TestChatChannel.subscribed).to be true
        expect(TestChatChannel.found_room).to be_nil
      end
    end
  end
end

We only considered a single use case for command callbacks, though there are plenty of others. For example, you could set the current user's time zone or locale or provide some context via Current attributes or dry-effects.

Give this feature a try with anycable-rails today! (Even if you're not using AnyCable... yet 😉)

---

P.S. Using the apartment gem with AnyCable requires a bit more attention in case you store Active Record objects as connection identifiers. See this issue.