message_bus 2.0.0.beta.2 → 2.0.0.beta.3

data/lib/message_bus/backends/postgres.rb

@@ -0,0 +1,396 @@
+require 'pg'
+
+module MessageBus::Postgres; end
+
+class MessageBus::Postgres::Client
+  INHERITED_CONNECTIONS = []
+
+  class Listener
+    attr_reader :do_sub, :do_unsub, :do_message
+
+    def subscribe(&block)
+      @do_sub = block
+    end
+
+    def unsubscribe(&block)
+      @do_unsub = block
+    end
+
+    def message(&block)
+      @do_message = block
+    end
+  end
+
+  def initialize(config)
+    @config = config
+    @listening_on = {}
+    @available = []
+    @allocated = {}
+    @mutex = Mutex.new
+    @pid = Process.pid
+  end
+
+  def add(channel, value)
+    hold{|conn| exec_prepared(conn, 'insert_message', [channel, value]){|r| r.getvalue(0,0).to_i}}
+  end
+
+  def clear_global_backlog(backlog_id, num_to_keep)
+    if backlog_id > num_to_keep
+      hold{|conn| exec_prepared(conn, 'clear_global_backlog', [backlog_id - num_to_keep])}
+      nil
+    end
+  end
+
+  def clear_channel_backlog(channel, backlog_id, num_to_keep)
+    hold{|conn| exec_prepared(conn, 'clear_channel_backlog', [channel, backlog_id, num_to_keep])}
+    nil
+  end
+
+  def expire(max_backlog_age)
+    hold{|conn| exec_prepared(conn, 'expire', [max_backlog_age])}
+    nil
+  end
+
+  def backlog(channel, backlog_id)
+    hold{|conn| exec_prepared(conn, 'channel_backlog', [channel, backlog_id]){|r| r.values.each{|a| a[0] = a[0].to_i}}}
+  end
+
+  def global_backlog(backlog_id)
+    hold{|conn| exec_prepared(conn, 'global_backlog', [backlog_id]){|r| r.values.each{|a| a[0] = a[0].to_i}}}
+  end
+
+  def get_value(channel, id)
+    hold{|conn| exec_prepared(conn, 'get_message', [channel, id]){|r| r.getvalue(0,0)}}
+  end
+
+  def reconnect
+    sync do
+      @listening_on.clear
+      @available.clear
+    end
+  end
+
+  # Dangerous, drops the message_bus table containing the backlog if it exists.
+  def reset!
+    hold do |conn|
+      conn.exec 'DROP TABLE IF EXISTS message_bus'
+      create_table(conn)
+    end
+  end
+
+  def max_id(channel=nil)
+    block = proc do |r|
+      if r.ntuples > 0
+        r.getvalue(0,0).to_i
+      else
+        0
+      end
+    end
+
+    if channel
+      hold{|conn| exec_prepared(conn, 'max_channel_id', [channel], &block)}
+    else
+      hold{|conn| exec_prepared(conn, 'max_id', &block)}
+    end
+  end
+
+  def publish(channel, data)
+    hold{|conn| exec_prepared(conn, 'publish', [channel, data])}
+  end
+
+  def subscribe(channel)
+    obj = Object.new
+    sync{@listening_on[channel] = obj}
+    listener = Listener.new
+    yield listener
+    pid = Process.pid
+    
+    conn = raw_pg_connection
+    conn.exec "LISTEN #{channel}"
+    listener.do_sub.call
+    while listening_on?(channel, obj)
+      conn.wait_for_notify(10) do |_,_,payload|
+        listener.do_message.call(nil, payload)
+      end
+      break if pid != Process.pid
+    end
+    listener.do_unsub.call
+
+    if pid != Process.pid
+      sync{INHERITED_CONNECTIONS << conn}
+    else
+      conn.exec "UNLISTEN #{channel}"
+    end
+
+    nil
+  end
+
+  def unsubscribe
+    sync{@listening_on.clear}
+  end
+
+  private
+
+  def exec_prepared(conn, *a)
+    r = conn.exec_prepared(*a)
+    yield r if block_given?
+  ensure
+    r.clear if r.respond_to?(:clear)
+  end
+
+  def create_table(conn)
+    conn.exec 'CREATE TABLE message_bus (id bigserial PRIMARY KEY, channel text NOT NULL, value text NOT NULL, added_at timestamp DEFAULT CURRENT_TIMESTAMP NOT NULL)'
+    conn.exec 'CREATE INDEX table_channel_id_index ON message_bus (channel, id)'
+    conn.exec 'CREATE INDEX table_added_at_index ON message_bus (added_at)'
+    nil
+  end
+
+  def hold
+    current_pid = Process.pid
+    if current_pid != @pid
+      @pid = current_pid
+      sync do
+        @listening_on.clear
+        INHERITED_CONNECTIONS.concat(@available)
+        @available.clear
+      end
+    end
+
+    if conn = sync{@allocated[Thread.current]}
+      return yield(conn)
+    end
+      
+    begin
+      conn = sync{@available.shift} || new_pg_connection
+      sync{@allocated[Thread.current] = conn}
+      yield conn
+    rescue PG::ConnectionBad => e
+      # don't add this connection back to the pool
+    ensure
+      sync{@allocated.delete(Thread.current)}
+      if Process.pid != current_pid
+        sync{INHERITED_CONNECTIONS << conn}
+      elsif conn && !e
+        sync{@available << conn}
+      end
+    end
+  end
+
+  def raw_pg_connection
+    PG::Connection.connect(@config[:backend_options] || {})
+  end
+
+  def new_pg_connection
+    conn = raw_pg_connection
+
+    begin
+      conn.exec("SELECT 'message_bus'::regclass")
+    rescue PG::UndefinedTable
+      create_table(conn)
+    end
+
+    conn.exec 'PREPARE insert_message AS INSERT INTO message_bus (channel, value) VALUES ($1, $2) RETURNING id'
+    conn.exec 'PREPARE clear_global_backlog AS DELETE FROM message_bus WHERE (id <= $1)'
+    conn.exec 'PREPARE clear_channel_backlog AS DELETE FROM message_bus WHERE ((channel = $1) AND (id <= (SELECT id FROM message_bus WHERE ((channel = $1) AND (id <= $2)) ORDER BY id DESC OFFSET $3)))'
+    conn.exec 'PREPARE channel_backlog AS SELECT id, value FROM message_bus WHERE ((channel = $1) AND (id > $2)) ORDER BY id'
+    conn.exec 'PREPARE global_backlog AS SELECT id, channel, value FROM message_bus WHERE (id > $1) ORDER BY id'
+    conn.exec "PREPARE expire AS DELETE FROM message_bus WHERE added_at < CURRENT_TIMESTAMP - ($1::text || ' seconds')::interval"
+    conn.exec 'PREPARE get_message AS SELECT value FROM message_bus WHERE ((channel = $1) AND (id = $2))'
+    conn.exec 'PREPARE max_channel_id AS SELECT max(id) FROM message_bus WHERE (channel = $1)'
+    conn.exec 'PREPARE max_id AS SELECT max(id) FROM message_bus'
+    conn.exec 'PREPARE publish AS SELECT pg_notify($1, $2)'
+
+    conn
+  end
+
+  def listening_on?(channel, obj)
+    sync{@listening_on[channel]} == obj
+  end
+
+  def sync
+    @mutex.synchronize{yield}
+  end
+end
+
+class MessageBus::Postgres::ReliablePubSub
+  attr_reader :subscribed
+  attr_accessor :max_publish_retries, :max_publish_wait, :max_backlog_size,
+                :max_global_backlog_size, :max_in_memory_publish_backlog,
+                :max_backlog_age
+
+  UNSUB_MESSAGE = "$$UNSUBSCRIBE"
+
+  def self.reset!(config)
+    MessageBus::Postgres::Client.new(config).reset!
+  end
+
+  # max_backlog_size is per multiplexed channel
+  def initialize(config = {}, max_backlog_size = 1000)
+    @config = config
+    @max_backlog_size = max_backlog_size
+    @max_global_backlog_size = 2000
+    @max_publish_retries = 10
+    @max_publish_wait = 500 #ms
+    @max_in_memory_publish_backlog = 1000
+    @in_memory_backlog = []
+    @lock = Mutex.new
+    @flush_backlog_thread = nil
+    # after 7 days inactive backlogs will be removed
+    @max_backlog_age = 604800
+    @h = {}
+  end
+
+  def new_connection
+    MessageBus::Postgres::Client.new(@config)
+  end
+
+  def backend
+    :postgres
+  end
+
+  def after_fork
+    client.reconnect
+  end
+
+  def postgresql_channel_name
+    db = @config[:db] || 0
+    "_message_bus_#{db}"
+  end
+
+  def client
+    @client ||= new_connection
+  end
+
+  # use with extreme care, will nuke all of the data
+  def reset!
+    client.reset!
+  end
+
+  def publish(channel, data, queue_in_memory=true)
+    client = self.client
+    backlog_id = client.add(channel, data)
+    msg = MessageBus::Message.new backlog_id, backlog_id, channel, data
+    payload = msg.encode
+    client.publish postgresql_channel_name, payload
+    client.clear_global_backlog(backlog_id, @max_global_backlog_size)
+    client.expire(@max_backlog_age)
+    client.clear_channel_backlog(channel, backlog_id, @max_backlog_size)
+
+    backlog_id
+  end
+
+  def last_id(channel)
+    client.max_id(channel)
+  end
+
+  def backlog(channel, last_id = nil)
+    items = client.backlog channel, last_id.to_i
+
+    items.map! do |id, data|
+      MessageBus::Message.new id, id, channel, data
+    end
+  end
+
+  def global_backlog(last_id = nil)
+    last_id = last_id.to_i
+
+    items = client.global_backlog last_id.to_i
+
+    items.map! do |id, channel, data|
+      MessageBus::Message.new id, id, channel, data
+    end
+  end
+
+  def get_message(channel, message_id)
+    if data = client.get_value(channel, message_id)
+      MessageBus::Message.new message_id, message_id, channel, data
+    else
+      nil
+    end
+  end
+
+  def subscribe(channel, last_id = nil)
+    # trivial implementation for now,
+    #   can cut down on connections if we only have one global subscriber
+    raise ArgumentError unless block_given?
+
+    global_subscribe(last_id) do |m|
+      yield m if m.channel == channel
+    end
+  end
+
+  def process_global_backlog(highest_id)
+    if highest_id > client.max_id
+      highest_id = 0
+    end
+
+    global_backlog(highest_id).each do |old|
+      yield old
+      highest_id = old.global_id
+    end
+
+    highest_id
+  end
+
+  def global_unsubscribe
+    client.publish(postgresql_channel_name, UNSUB_MESSAGE)
+    @subscribed = false
+  end
+
+  def global_subscribe(last_id=nil, &blk)
+    raise ArgumentError unless block_given?
+    highest_id = last_id
+
+    begin
+      client.subscribe(postgresql_channel_name) do |on|
+        h = {}
+
+        on.subscribe do
+          if highest_id
+            process_global_backlog(highest_id) do |m|
+              h[m.global_id] = true
+              yield m
+            end
+          end
+          h = nil if h.empty?
+          @subscribed = true
+        end
+
+        on.unsubscribe do
+          @subscribed = false
+        end
+
+        on.message do |c,m|
+          if m == UNSUB_MESSAGE
+            @subscribed = false
+            return
+          end
+          m = MessageBus::Message.decode m
+
+          # we have 3 options
+          #
+          # 1. message came in the correct order GREAT, just deal with it
+          # 2. message came in the incorrect order COMPLICATED, wait a tiny bit and clear backlog
+          # 3. message came in the incorrect order and is lowest than current highest id, reset
+
+          if h
+            # If already yielded during the clear backlog when subscribing,
+            # don't yield a duplicate copy.
+            unless h.delete(m.global_id)
+              h = nil if h.empty?
+              yield m
+            end
+          else
+            yield m
+          end
+        end
+      end
+    rescue => error
+      MessageBus.logger.warn "#{error} subscribe failed, reconnecting in 1 second. Call stack\n#{error.backtrace.join("\n")}"
+      sleep 1
+      retry
+    end
+  end
+
+  MessageBus::BACKENDS[:postgres] = self
+end

data/lib/message_bus/{redis/reliable_pub_sub.rb → backends/redis.rb}

@@ -339,4 +339,5 @@ class MessageBus::Redis::ReliablePubSub
     end
   end
 
+  MessageBus::BACKENDS[:redis] = self
 end

data/lib/message_bus/rack/middleware.rb

@@ -84,13 +84,22 @@ class MessageBus::Rack::Middleware
     client = MessageBus::Client.new(message_bus: @bus, client_id: client_id,
                                     user_id: user_id, site_id: site_id, group_ids: group_ids)
 
-    request = Rack::Request.new(env)
-    request.POST.each do |k,v|
-      if k == "__seq".freeze
-        client.seq = v.to_i
-      else
+    if channels = env['message_bus.channels']
+      if seq = env['message_bus.seq']
+        client.seq = seq.to_i
+      end
+      channels.each do |k, v|
         client.subscribe(k, v)
       end
+    else
+      request = Rack::Request.new(env)
+      request.POST.each do |k,v|
+        if k == "__seq".freeze
+          client.seq = v.to_i
+        else
+          client.subscribe(k, v)
+        end
+      end
     end
 
     headers = {}

data/lib/message_bus/version.rb

@@ -1,3 +1,3 @@
 module MessageBus
-  VERSION = "2.0.0.beta.2"
+  VERSION = "2.0.0.beta.3"
 end

data/message_bus.gemspec

@@ -15,6 +15,8 @@ Gem::Specification.new do |gem|
   gem.name          = "message_bus"
   gem.require_paths = ["lib"]
   gem.version       = MessageBus::VERSION
+  gem.required_ruby_version = ">= 1.9.3"
   gem.add_runtime_dependency 'rack', '>= 1.1.3'
-  gem.add_runtime_dependency 'redis'
+  gem.add_development_dependency 'redis'
+  gem.add_development_dependency 'pg'
 end

data/spec/lib/message_bus/assets/asset_encoding_spec.rb

@@ -1,10 +1,11 @@
+require_relative '../../../spec_helper'
 asset_directory = File.expand_path('../../../../../assets', __FILE__)
 asset_file_paths = Dir.glob(File.join(asset_directory, 'message-bus.js'))
 asset_file_names = asset_file_paths.map{|e| File.basename(e) }
 
 describe asset_file_names do
   it 'should contain .js files' do
-    expect(asset_file_names).to include('message-bus.js')
+    asset_file_names.must_include('message-bus.js')
   end
 end
 
@@ -12,8 +13,7 @@ asset_file_paths.each do | path |
   describe "Asset file #{File.basename(path).inspect}" do
     it 'should be encodable as UTF8' do
       binary_data = File.open(path, 'rb'){|f| f.read }
-      encode_block = -> { binary_data.encode(Encoding::UTF_8) }
-      expect(encode_block).not_to raise_error
+      binary_data.encode(Encoding::UTF_8)
     end
   end
-end
+end