message_bus 3.3.4 → 3.3.8

data/lib/message_bus/backends/memory.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require "message_bus/backends/base"
-
 module MessageBus
   module Backends
     # The memory backend stores published messages in a simple array per
@@ -214,6 +212,12 @@ module MessageBus
         client.reset!
       end
 
+      # No-op; this backend doesn't maintain any storage connections.
+      # (see Base#destroy)
+      def destroy
+        nil
+      end
+
       # (see Base#expire_all_backlogs!)
       def expire_all_backlogs!
         client.expire_all_backlogs!

data/lib/message_bus/backends/postgres.rb

@@ -2,8 +2,6 @@
 
 require 'pg'
 
-require "message_bus/backends/base"
-
 module MessageBus
   module Backends
     # The Postgres backend stores published messages in a single Postgres table
@@ -46,6 +44,7 @@ module MessageBus
           @listening_on = {}
           @available = []
           @allocated = {}
+          @subscribe_connection = nil
           @mutex = Mutex.new
           @pid = Process.pid
         end
@@ -87,10 +86,12 @@ module MessageBus
           hold { |conn| exec_prepared(conn, 'get_message', [channel, id]) { |r| r.getvalue(0, 0) } }
         end
 
-        def reconnect
+        def after_fork
           sync do
-            @listening_on.clear
+            @pid = Process.pid
+            INHERITED_CONNECTIONS.concat(@available)
             @available.clear
+            @listening_on.clear
           end
         end
 
@@ -102,6 +103,13 @@ module MessageBus
           end
         end
 
+        def destroy
+          sync do
+            @available.each(&:close)
+            @available.clear
+          end
+        end
+
         # use with extreme care, will nuke all of the data
         def expire_all_backlogs!
           reset!
@@ -133,7 +141,7 @@ module MessageBus
           listener = Listener.new
           yield listener
 
-          conn = raw_pg_connection
+          conn = @subscribe_connection = raw_pg_connection
           conn.exec "LISTEN #{channel}"
           listener.do_sub.call
           while listening_on?(channel, obj)
@@ -147,6 +155,9 @@ module MessageBus
 
           conn.exec "UNLISTEN #{channel}"
           nil
+        ensure
+          @subscribe_connection&.close
+          @subscribe_connection = nil
         end
 
         def unsubscribe
@@ -172,11 +183,7 @@ module MessageBus
         def hold
           current_pid = Process.pid
           if current_pid != @pid
-            @pid = current_pid
-            sync do
-              INHERITED_CONNECTIONS.concat(@available)
-              @available.clear
-            end
+            after_fork
           end
 
           if conn = sync { @allocated[Thread.current] }
@@ -251,12 +258,13 @@ module MessageBus
         # after 7 days inactive backlogs will be removed
         @max_backlog_age = 604800
         @clear_every = config[:clear_every] || 1
+        @mutex = Mutex.new
       end
 
-      # Reconnects to Postgres; used after a process fork, typically triggerd by a forking webserver
+      # Reconnects to Postgres; used after a process fork, typically triggered by a forking webserver
       # @see Base#after_fork
       def after_fork
-        client.reconnect
+        client.after_fork
       end
 
       # (see Base#reset!)
@@ -264,6 +272,11 @@ module MessageBus
         client.reset!
       end
 
+      # (see Base#destroy)
+      def destroy
+        client.destroy
+      end
+
       # (see Base#expire_all_backlogs!)
       def expire_all_backlogs!
         client.expire_all_backlogs!
@@ -279,7 +292,7 @@ module MessageBus
         msg = MessageBus::Message.new backlog_id, backlog_id, channel, data
         payload = msg.encode
         c.publish postgresql_channel_name, payload
-        if backlog_id % clear_every == 0
+        if backlog_id && backlog_id % clear_every == 0
           max_backlog_size = (opts && opts[:max_backlog_size]) || self.max_backlog_size
           max_backlog_age = (opts && opts[:max_backlog_age]) || self.max_backlog_age
           c.clear_global_backlog(backlog_id, @max_global_backlog_size)
@@ -399,11 +412,7 @@ module MessageBus
       private
 
       def client
-        @client ||= new_connection
-      end
-
-      def new_connection
-        Client.new(@config)
+        @client || @mutex.synchronize { @client ||= Client.new(@config) }
       end
 
       def postgresql_channel_name

data/lib/message_bus/backends/redis.rb

@@ -3,8 +3,6 @@
 require 'redis'
 require 'digest'
 
-require "message_bus/backends/base"
-
 module MessageBus
   module Backends
     # The Redis backend stores published messages in Redis sorted sets (using
@@ -64,10 +62,10 @@ module MessageBus
         @max_backlog_age = 604800
       end
 
-      # Reconnects to Redis; used after a process fork, typically triggerd by a forking webserver
+      # Reconnects to Redis; used after a process fork, typically triggered by a forking webserver
       # @see Base#after_fork
       def after_fork
-        pub_redis.disconnect!
+        @pub_redis&.disconnect!
       end
 
       # (see Base#reset!)
@@ -77,6 +75,11 @@ module MessageBus
         end
       end
 
+      # (see Base#destroy)
+      def destroy
+        @pub_redis&.disconnect!
+      end
+
       # Deletes all backlogs and their data. Does not delete ID pointers, so new publications will get IDs that continue from the last publication before the expiry. Use with extreme caution.
       # @see Base#expire_all_backlogs!
       def expire_all_backlogs!
@@ -104,8 +107,8 @@ module MessageBus
 
       local global_id = redis.call("INCR", global_id_key)
       local backlog_id = redis.call("INCR", backlog_id_key)
-      local payload = string.format("%i|%i|%s", global_id, backlog_id, start_payload)
-      local global_backlog_message = string.format("%i|%s", backlog_id, channel)
+      local payload = table.concat({ global_id, backlog_id, start_payload }, "|")
+      local global_backlog_message = table.concat({ backlog_id, channel }, "|")
 
       redis.call("ZADD", backlog_key, backlog_id, payload)
       redis.call("EXPIRE", backlog_key, max_backlog_age)

data/lib/message_bus/client.rb

@@ -46,6 +46,9 @@ class MessageBus::Client
     @bus = opts[:message_bus] || MessageBus
     @subscriptions = {}
     @chunks_sent = 0
+    @async_response = nil
+    @io = nil
+    @wrote_headers = false
   end
 
   # @yield executed with a lock on the Client instance
@@ -220,7 +223,7 @@ class MessageBus::Client
 
   private
 
-  # heavily optimised to avoid all uneeded allocations
+  # heavily optimised to avoid all unneeded allocations
   NEWLINE = "\r\n".freeze
   COLON_SPACE = ": ".freeze
   HTTP_11 = "HTTP/1.1 200 OK\r\n".freeze
@@ -261,7 +264,7 @@ class MessageBus::Client
       @wrote_headers = true
     end
 
-    # chunked encoding may be "re-chunked" by proxies, so add a seperator
+    # chunked encoding may be "re-chunked" by proxies, so add a separator
     postfix = NEWLINE + "|" + NEWLINE
     data = data.gsub(postfix, NEWLINE + "||" + NEWLINE)
     chunk_length = data.bytesize + postfix.bytesize
@@ -275,11 +278,7 @@ class MessageBus::Client
       @async_response << postfix
       @async_response << NEWLINE
     else
-      @io.write(chunk_length.to_s(16))
-      @io.write(NEWLINE)
-      @io.write(data)
-      @io.write(postfix)
-      @io.write(NEWLINE)
+      @io.write(chunk_length.to_s(16) << NEWLINE << data << postfix << NEWLINE)
     end
   end
 

data/lib/message_bus/codec/base.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module MessageBus
+  module Codec
+    class Base
+      def encode(hash)
+        raise ConcreteClassMustImplementError
+      end
+
+      def decode(payload)
+        raise ConcreteClassMustImplementError
+      end
+    end
+
+    autoload :Json, File.expand_path("json", __dir__)
+    autoload :Oj, File.expand_path("oj", __dir__)
+  end
+end

data/lib/message_bus/codec/json.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module MessageBus
+  module Codec
+    class Json < Base
+      def encode(hash)
+        JSON.dump(hash)
+      end
+
+      def decode(payload)
+        JSON.parse(payload)
+      end
+    end
+  end
+end

data/lib/message_bus/codec/oj.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+require 'oj' unless defined? ::Oj
+
+module MessageBus
+  module Codec
+    class Oj < Base
+      def initialize(options = { mode: :compat })
+        @options = options
+      end
+
+      def encode(hash)
+        ::Oj.dump(hash, @options)
+      end
+
+      def decode(payload)
+        ::Oj.load(payload, @options)
+      end
+    end
+  end
+end

data/lib/message_bus/connection_manager.rb

@@ -37,7 +37,7 @@ class MessageBus::ConnectionManager
             rescue
               # pipe may be broken, move on
             end
-            # turns out you can delete from a set while itereating
+            # turns out you can delete from a set while iterating
             remove_client(client) if client.closed?
           end
         end

data/lib/message_bus/distributed_cache.rb

@@ -22,6 +22,7 @@ module MessageBus
         @lock = Mutex.new
         @message_bus = message_bus || MessageBus
         @publish_queue_in_memory = publish_queue_in_memory
+        @app_version = nil
       end
 
       def subscribers
@@ -45,7 +46,8 @@ module MessageBus
             hash = current.hash(message.site_id || DEFAULT_SITE_ID)
 
             case payload["op"]
-            when "set" then hash[payload["key"]] = payload["marshalled"] ? Marshal.load(Base64.decode64(payload["value"])) : payload["value"]
+            # TODO: consider custom marshal support with a restricted set
+            when "set" then hash[payload["key"]] = payload["marshalled"] ? Marshal.load(Base64.decode64(payload["value"])) : payload["value"] # rubocop:disable Security/MarshalLoad
             when "delete" then hash.delete(payload["key"])
             when "clear"  then hash.clear
             end

data/lib/message_bus/http_client.rb

@@ -146,8 +146,8 @@ module MessageBus
     #
     # A last_message_id may be provided.
     #  * -1 will subscribe to all new messages
-    #  * -2 will recieve last message + all new messages
-    #  * -3 will recieve last 2 message + all new messages
+    #  * -2 will receive last message + all new messages
+    #  * -3 will receive last 2 message + all new messages
     #
     # @example Subscribing to a channel with `last_message_id`
     #   client.subscribe("/test", last_message_id: -2) do |payload|

data/lib/message_bus/rack/diagnostics.rb

@@ -13,6 +13,15 @@ class MessageBus::Rack::Diagnostics
     @bus = config[:message_bus] || MessageBus
   end
 
+  JS_ASSETS = %w{
+    jquery-1.8.2.js
+    react.js
+    react-dom.js
+    babel.min.js
+    message-bus.js
+    application.jsx
+  }
+
   # Process an HTTP request from a subscriber client
   # @param [Rack::Request::Env] env the request environment
   def call(env)
@@ -39,9 +48,9 @@ class MessageBus::Rack::Diagnostics
     end
 
     asset = route.split('/assets/')[1]
-    if asset && !asset !~ /\//
+
+    if asset && JS_ASSETS.include?(asset)
       content = asset_contents(asset)
-      split = asset.split('.')
       return [200, { 'Content-Type' => 'application/javascript;charset=UTF-8' }, [content]]
     end
 
@@ -75,6 +84,23 @@ class MessageBus::Rack::Diagnostics
     File.expand_path("../../../../assets/#{asset}", __FILE__)
   end
 
+  def script_tags
+    tags = []
+
+    JS_ASSETS.each do |asset|
+      type =
+        if asset.end_with?('.js')
+          'text/javascript'
+        elsif asset.end_with?('.jsx')
+          'text/jsx'
+        end
+
+      tags << js_asset(asset, type)
+    end
+
+    tags.join("\n")
+  end
+
   def index
     html = <<~HTML
       <!DOCTYPE html>
@@ -83,12 +109,8 @@ class MessageBus::Rack::Diagnostics
         </head>
         <body>
           <div id="app"></div>
-          #{js_asset "jquery-1.8.2.js"}
-          #{js_asset "react.js"}
-          #{js_asset "react-dom.js"}
-          #{js_asset "babel.min.js"}
-          #{js_asset "message-bus.js"}
-          #{js_asset "application.jsx", "text/jsx"}
+
+          #{script_tags}
         </body>
       </html>
     HTML

data/lib/message_bus/rack/middleware.rb

@@ -66,6 +66,12 @@ class MessageBus::Rack::Middleware
   private
 
   def handle_request(env)
+    # Prevent simple polling from clobbering the session
+    # See: https://github.com/discourse/message_bus/issues/257
+    if (rack_session_options = env[Rack::RACK_SESSION_OPTIONS])
+      rack_session_options[:skip] = true
+    end
+
     # special debug/test route
     if @bus.allow_broadcast? && env['PATH_INFO'] == @broadcast_route
       parsed = Rack::Request.new(env)
@@ -81,6 +87,22 @@ class MessageBus::Rack::Middleware
     client_id = env['PATH_INFO'][@base_route_length..-1].split("/")[0]
     return [404, {}, ["not found"]] unless client_id
 
+    headers = {}
+    headers["Cache-Control"] = "must-revalidate, private, max-age=0"
+    headers["Content-Type"] = "application/json; charset=utf-8"
+    headers["Pragma"] = "no-cache"
+    headers["Expires"] = "0"
+
+    if @bus.extra_response_headers_lookup
+      @bus.extra_response_headers_lookup.call(env).each do |k, v|
+        headers[k] = v
+      end
+    end
+
+    if env["REQUEST_METHOD"] == "OPTIONS"
+      return [200, headers, ["OK"]]
+    end
+
     user_id = @bus.user_id_lookup.call(env) if @bus.user_id_lookup
     group_ids = @bus.group_ids_lookup.call(env) if @bus.group_ids_lookup
     site_id = @bus.site_id_lookup.call(env) if @bus.site_id_lookup
@@ -111,22 +133,6 @@ class MessageBus::Rack::Middleware
       end
     end
 
-    headers = {}
-    headers["Cache-Control"] = "must-revalidate, private, max-age=0"
-    headers["Content-Type"] = "application/json; charset=utf-8"
-    headers["Pragma"] = "no-cache"
-    headers["Expires"] = "0"
-
-    if @bus.extra_response_headers_lookup
-      @bus.extra_response_headers_lookup.call(env).each do |k, v|
-        headers[k] = v
-      end
-    end
-
-    if env["REQUEST_METHOD"] == "OPTIONS"
-      return [200, headers, ["OK"]]
-    end
-
     long_polling = @bus.long_polling_enabled? &&
                    env['QUERY_STRING'] !~ /dlp=t/ &&
                    @connection_manager.client_count < @bus.max_active_clients

data/lib/message_bus/rack/thin_ext.rb

@@ -9,6 +9,7 @@ module Thin
 
       def initialize
         @queue = []
+        @body_callback = nil
       end
 
       def call(body)
@@ -38,7 +39,7 @@ module Thin
     end
   end
 
-  # Response whos body is sent asynchronously.
+  # Response which body is sent asynchronously.
   class AsyncResponse
     include Rack::Response::Helpers
 

data/lib/message_bus/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module MessageBus
-  VERSION = "3.3.4"
+  VERSION = "3.3.8"
 end

data/lib/message_bus.rb

@@ -2,18 +2,22 @@
 
 require "monitor"
 require "set"
-require "message_bus/version"
-require "message_bus/message"
-require "message_bus/client"
-require "message_bus/connection_manager"
-require "message_bus/diagnostics"
-require "message_bus/rack/middleware"
-require "message_bus/rack/diagnostics"
-require "message_bus/timer_thread"
+
+require_relative "message_bus/version"
+require_relative "message_bus/message"
+require_relative "message_bus/client"
+require_relative "message_bus/connection_manager"
+require_relative "message_bus/diagnostics"
+require_relative "message_bus/rack/middleware"
+require_relative "message_bus/rack/diagnostics"
+require_relative "message_bus/timer_thread"
+require_relative "message_bus/codec/base"
+require_relative "message_bus/backends"
+require_relative "message_bus/backends/base"
 
 # we still need to take care of the logger
-if defined?(::Rails)
-  require 'message_bus/rails/railtie'
+if defined?(::Rails::Engine)
+  require_relative 'message_bus/rails/railtie'
 end
 
 # @see MessageBus::Implementation
@@ -37,6 +41,10 @@ module MessageBus::Implementation
   def initialize
     @config = {}
     @mutex = Synchronizer.new
+    @off = false
+    @destroyed = false
+    @timer_thread = nil
+    @subscriber_thread = nil
   end
 
   # @param [Boolean] val whether or not to cache static assets for the diagnostics pages
@@ -96,14 +104,14 @@ module MessageBus::Implementation
     configure(long_polling_enabled: val)
   end
 
-  # @param [Integer] val The number of simultanuous clients we can service;
+  # @param [Integer] val The number of simultaneous clients we can service;
   #   will revert to polling if we are out of slots
   # @return [void]
   def max_active_clients=(val)
     configure(max_active_clients: val)
   end
 
-  # @return [Integer] The number of simultanuous clients we can service;
+  # @return [Integer] The number of simultaneous clients we can service;
   #   will revert to polling if we are out of slots. Defaults to 1000 if not
   #   explicitly set.
   def max_active_clients
@@ -271,13 +279,24 @@ module MessageBus::Implementation
   #   set, defaults to false unless we're in Rails test or development mode.
   def allow_broadcast?
     @config[:allow_broadcast] ||=
-      if defined? ::Rails
+      if defined? ::Rails.env
         ::Rails.env.test? || ::Rails.env.development?
       else
         false
       end
   end
 
+  # @param [MessageBus::Codec::Base] codec used to encode and decode Message payloads
+  # @return [void]
+  def transport_codec=(codec)
+    configure(transport_codec: codec)
+  end
+
+  # @return [MessageBus::Codec::Base] codec used to encode and decode Message payloads
+  def transport_codec
+    @config[:transport_codec] ||= MessageBus::Codec::Json.new
+  end
+
   # @param [MessageBus::Backend::Base] pub_sub a configured backend
   # @return [void]
   def reliable_pub_sub=(pub_sub)
@@ -323,7 +342,7 @@ module MessageBus::Implementation
   # @option opts [Array<String,Integer>] :group_ids (`nil`) the group IDs to which the message should be available. If nil, available to all.
   # @option opts [String] :site_id (`nil`) the site ID to scope the message to; used for hosting multiple
   #   applications or instances of an application against a single message_bus
-  # @option opts [nil,Integer] :max_backlog_age the longest amount of time a message may live in a backlog before beging removed, in seconds
+  # @option opts [nil,Integer] :max_backlog_age the longest amount of time a message may live in a backlog before being removed, in seconds
   # @option opts [nil,Integer] :max_backlog_size the largest permitted size (number of messages) for the channel backlog; beyond this capacity, old messages will be dropped
   #
   # @return [Integer] the channel-specific ID the message was given
@@ -358,12 +377,12 @@ module MessageBus::Implementation
       raise ::MessageBus::InvalidMessageTarget
     end
 
-    encoded_data = JSON.dump(
-      data: data,
-      user_ids: user_ids,
-      group_ids: group_ids,
-      client_ids: client_ids
-    )
+    encoded_data = transport_codec.encode({
+      "data" => data,
+      "user_ids" => user_ids,
+      "group_ids" => group_ids,
+      "client_ids" => client_ids
+    })
 
     channel_opts = {}
 
@@ -512,12 +531,13 @@ module MessageBus::Implementation
   end
 
   # Stops listening for publications and stops executing scheduled tasks.
-  # Mostly used in tests to detroy entire bus.
+  # Mostly used in tests to destroy entire bus.
   # @return [void]
   def destroy
     return if @destroyed
 
     reliable_pub_sub.global_unsubscribe
+    reliable_pub_sub.destroy
 
     @mutex.synchronize do
       return if @destroyed
@@ -626,7 +646,7 @@ module MessageBus::Implementation
     channel, site_id = decode_channel_name(msg.channel)
     msg.channel = channel
     msg.site_id = site_id
-    parsed = JSON.parse(msg.data)
+    parsed = transport_codec.decode(msg.data)
     msg.data = parsed["data"]
     msg.user_ids = parsed["user_ids"]
     msg.group_ids = parsed["group_ids"]

data/message_bus.gemspec

@@ -1,5 +1,4 @@
 # frozen_string_literal: true
-# -*- encoding: utf-8 -*-
 
 require File.expand_path('../lib/message_bus/version', __FILE__)
 
@@ -8,7 +7,7 @@ Gem::Specification.new do |gem|
   gem.email         = ["sam.saffron@gmail.com"]
   gem.description   = %q{A message bus for rack}
   gem.summary       = %q{}
-  gem.homepage      = "https://github.com/SamSaffron/message_bus"
+  gem.homepage      = "https://github.com/discourse/message_bus"
   gem.license       = "MIT"
   gem.files         = `git ls-files`.split($\) +
                       ["vendor/assets/javascripts/message-bus.js", "vendor/assets/javascripts/message-bus-ajax.js"]
@@ -17,8 +16,27 @@ Gem::Specification.new do |gem|
   gem.name          = "message_bus"
   gem.require_paths = ["lib"]
   gem.version       = MessageBus::VERSION
-  gem.required_ruby_version = ">= 2.3.0"
+  gem.required_ruby_version = ">= 2.6.0"
+
   gem.add_runtime_dependency 'rack', '>= 1.1.3'
+
+  # Optional runtime dependencies
   gem.add_development_dependency 'redis'
   gem.add_development_dependency 'pg'
+  gem.add_development_dependency 'concurrent-ruby' # for distributed-cache
+
+  gem.add_development_dependency 'minitest'
+  gem.add_development_dependency 'minitest-hooks'
+  gem.add_development_dependency 'minitest-global_expectations'
+  gem.add_development_dependency 'rake'
+  gem.add_development_dependency 'http_parser.rb'
+  gem.add_development_dependency 'thin'
+  gem.add_development_dependency 'rack-test'
+  gem.add_development_dependency 'puma'
+  gem.add_development_dependency 'm'
+  gem.add_development_dependency 'byebug'
+  gem.add_development_dependency 'oj'
+  gem.add_development_dependency 'yard'
+  gem.add_development_dependency 'rubocop-discourse'
+  gem.add_development_dependency 'rubocop-rspec'
 end