pakyow-realtime 0.11.3 → 1.0.0.rc1

data/lib/pakyow/realtime/server/adapters/redis.rb

@@ -0,0 +1,277 @@
+# frozen_string_literal: true
+
+require "redis"
+require "concurrent/array"
+require "concurrent/timer_task"
+require "connection_pool"
+
+module Pakyow
+  module Realtime
+    class Server
+      module Adapters
+        # Manages websocket channels in redis.
+        #
+        # Use this in production.
+        #
+        # @api private
+        class Redis
+          KEY_PART_SEPARATOR = "/"
+          KEY_PREFIX = "realtime"
+          INFINITY = "+inf"
+
+          PUBSUB_PREFIX = "pubsub"
+
+          def initialize(server, config)
+            @server, @config = server, config
+            @prefix = [@config[:key_prefix], KEY_PREFIX].join(KEY_PART_SEPARATOR)
+
+            connect
+            cleanup
+          end
+
+          def connect
+            @redis = ConnectionPool.new(**@config[:pool]) {
+              ::Redis.new(@config[:connection])
+            }
+
+            @buffer = Buffer.new(@redis, pubsub_channel)
+            @subscriber = Subscriber.new(::Redis.new(@config[:connection]), pubsub_channel) do |payload|
+              channel, message = Marshal.restore(payload).values_at(:channel, :message)
+              @server.transmit_message_to_connection_ids(message, socket_ids_for_channel(channel), raw: true)
+            end
+          end
+
+          def disconnect
+            @subscriber.disconnect
+          end
+
+          def socket_subscribe(socket_id, *channels)
+            @redis.with do |redis|
+              redis.multi do |transaction|
+                channels.each do |channel|
+                  channel = channel.to_s
+                  transaction.zadd(key_socket_ids_by_channel(channel), INFINITY, socket_id)
+                  transaction.zadd(key_channels_by_socket_id(socket_id), INFINITY, channel)
+                end
+              end
+            end
+          end
+
+          def socket_unsubscribe(*channels)
+            @redis.with do |redis|
+              channels.each do |channel|
+                channel = channel.to_s
+
+                # Channel could contain a wildcard, so this takes some work...
+                redis.scan_each(match: key_socket_ids_by_channel(channel)) do |key|
+                  channel = key.split("channel:", 2)[1]
+
+                  socket_ids = redis.zrangebyscore(
+                    key, Time.now.to_i, INFINITY
+                  )
+
+                  redis.multi do |transaction|
+                    transaction.del(key)
+
+                    socket_ids.each do |socket_id|
+                      transaction.zrem(key_channels_by_socket_id(socket_id), channel)
+                    end
+                  end
+                end
+              end
+            end
+          end
+
+          def subscription_broadcast(channel, message)
+            @buffer << Marshal.dump(channel: channel, message: { payload: message }.to_json)
+          end
+
+          def expire(socket_id, seconds)
+            time_expire = Time.now.to_i + seconds
+            channels = channels_for_socket_id(socket_id)
+
+            @redis.with do |redis|
+              redis.multi do |transaction|
+                channels.each do |channel|
+                  transaction.zadd(key_socket_ids_by_channel(channel), time_expire, socket_id)
+                end
+
+                transaction.expireat(key_channels_by_socket_id(socket_id), time_expire + 1)
+                transaction.expireat(key_socket_instance_id_by_socket_id(socket_id), time_expire + 1)
+              end
+            end
+          end
+
+          def persist(socket_id)
+            channels = channels_for_socket_id(socket_id)
+
+            @redis.with do |redis|
+              redis.multi do |transaction|
+                channels.each do |channel|
+                  transaction.zadd(key_socket_ids_by_channel(channel), INFINITY, socket_id)
+                end
+
+                transaction.persist(key_channels_by_socket_id(socket_id))
+                transaction.persist(key_socket_instance_id_by_socket_id(socket_id))
+              end
+            end
+          end
+
+          def current!(socket_id, socket_instance_id)
+            @redis.with do |redis|
+              redis.set(key_socket_instance_id_by_socket_id(socket_id), socket_instance_id)
+            end
+          end
+
+          def current?(socket_id, socket_instance_id)
+            @redis.with do |redis|
+              redis.get(key_socket_instance_id_by_socket_id(socket_id)) == socket_instance_id.to_s
+            end
+          end
+
+          protected
+
+          def socket_ids_for_channel(channel)
+            @redis.with do |redis|
+              redis.zrangebyscore(
+                key_socket_ids_by_channel(channel), INFINITY, INFINITY
+              )
+            end
+          end
+
+          def channels_for_socket_id(socket_id)
+            @redis.with do |redis|
+              redis.zrangebyscore(
+                key_channels_by_socket_id(socket_id), INFINITY, INFINITY
+              )
+            end
+          end
+
+          def build_key(*parts)
+            [@prefix].concat(parts).join(KEY_PART_SEPARATOR)
+          end
+
+          def key_socket_ids_by_channel(channel)
+            build_key("channel:#{channel}")
+          end
+
+          def key_channels_by_socket_id(socket_id)
+            build_key("socket_id:#{socket_id}")
+          end
+
+          def key_socket_instance_id_by_socket_id(socket_id)
+            build_key("socket_instance_id:#{socket_id}")
+          end
+
+          def pubsub_channel
+            [@prefix, PUBSUB_PREFIX].join(KEY_PART_SEPARATOR)
+          end
+
+          def cleanup
+            Concurrent::TimerTask.new(execution_interval: 300, timeout_interval: 300) {
+              Pakyow.logger.debug "[Pakyow::Realtime::Server::Adapters::Redis] Cleaning up channel keys"
+
+              removed_count = 0
+              @redis.with do |redis|
+                redis.scan_each(match: key_socket_ids_by_channel("*")) do |key|
+                  socket_ids = redis.zrangebyscore(
+                    key, Time.now.to_i, INFINITY
+                  )
+
+                  if socket_ids.empty?
+                    removed_count += 1
+                    redis.del(key)
+                  end
+                end
+              end
+
+              Pakyow.logger.debug "[Pakyow::Realtime::Server::Adapters::Redis] Removed #{removed_count} keys"
+            }.execute
+          end
+
+          class Buffer
+            # The number of publish commands to pipeline to redis.
+            #
+            PUBLISH_BUFFER_SIZE = 1_000
+
+            # How often the publish buffer should be flushed.
+            #
+            PUBLISH_BUFFER_FLUSH_MS = 150
+
+            def initialize(redis, channel)
+              @redis, @channel = redis, channel
+              @buffer = Concurrent::Array.new
+            end
+
+            def <<(payload)
+              @buffer << payload
+              maybe_flush
+            end
+
+            protected
+
+            def maybe_flush
+              if @buffer.count > PUBLISH_BUFFER_SIZE
+                flush
+              end
+
+              unless @task&.pending?
+                @task = Concurrent::ScheduledTask.execute(PUBLISH_BUFFER_FLUSH_MS / 1_000) {
+                  flush
+                }
+              end
+            end
+
+            def flush
+              @redis.with do |redis|
+                redis.pipelined do |pipeline|
+                  until @buffer.empty?
+                    pipeline.publish(@channel, @buffer.shift)
+                  end
+                end
+              end
+            end
+          end
+
+          class Subscriber
+            def initialize(redis, channel, &callback)
+              @redis, @channel, @callback = redis, channel, callback
+
+              @thread = Thread.new do
+                subscribe
+              end
+            end
+
+            def disconnect
+              @thread.exit
+              @redis.disconnect!
+            end
+
+            def subscribe
+              @redis.subscribe(@channel) do |on|
+                on.message do |_, payload|
+                  begin
+                    @callback.call(payload)
+                  rescue => error
+                    Pakyow.logger.error "[Pakyow::Realtime::Server::Adapters::Redis] Subscriber callback failed: #{error}"
+                  end
+                end
+              end
+            rescue ::Redis::CannotConnectError
+              Pakyow.logger.error "[Pakyow::Realtime::Server::Adapters::Redis] Subscriber disconnected"
+              resubscribe
+            rescue => error
+              Pakyow.logger.error "[Pakyow::Realtime::Server::Adapters::Redis] Subscriber crashed: #{error}"
+              resubscribe
+            end
+
+            def resubscribe
+              sleep 0.25
+              subscribe
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pakyow/realtime/server.rb

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+require "concurrent/array"
+require "concurrent/timer_task"
+require "concurrent/executor/thread_pool_executor"
+
+require "pakyow/support/message_verifier"
+
+require "pakyow/realtime/websocket"
+
+module Pakyow
+  module Realtime
+    class Server
+      attr_reader :adapter
+
+      HEARTBEAT_INTERVAL = 3
+
+      def initialize(adapter = :memory, adapter_config, timeout_config)
+        require "pakyow/realtime/server/adapters/#{adapter}"
+        @adapter = Adapters.const_get(adapter.to_s.capitalize).new(self, adapter_config)
+        @sockets = Concurrent::Array.new
+        @timeout_config = timeout_config
+        @executor = Concurrent::ThreadPoolExecutor.new(
+          auto_terminate: false,
+          min_threads: 1,
+          max_threads: 10,
+          max_queue: 0
+        )
+
+        connect
+      rescue LoadError => e
+        Pakyow.logger.error "Failed to load data subscriber store adapter named `#{adapter}'"
+        Pakyow.logger.error e.message
+      end
+
+      def shutdown
+        disconnect
+        @executor.shutdown
+        @executor.wait_for_termination(30)
+      end
+
+      def connect
+        @executor << -> {
+          start_heartbeat; @adapter.connect
+        }
+      end
+
+      def disconnect
+        @executor << -> {
+          stop_heartbeat; @adapter.disconnect
+        }
+      end
+
+      def socket_connect(id_or_socket)
+        @executor << -> {
+          find_socket(id_or_socket) do |socket|
+            @sockets << socket
+            @adapter.persist(socket.id)
+            @adapter.current!(socket.id, socket.object_id)
+          end
+        }
+      end
+
+      def socket_disconnect(id_or_socket)
+        @executor << -> {
+          find_socket(id_or_socket) do |socket|
+            @sockets.delete(socket)
+
+            # If this isn't the current instance for the socket id, it means that a
+            # reconnect probably happened and the new socket connected before we
+            # knew that the old one disconnected. Since there's a newer socket,
+            # don't trigger leave events or expirations for the old one.
+            #
+            if @adapter.current?(socket.id, socket.object_id)
+              socket.leave
+              @adapter.expire(socket.id, @timeout_config.disconnect)
+            end
+          end
+        }
+      end
+
+      def socket_subscribe(id_or_socket, *channels)
+        @executor << -> {
+          find_socket_id(id_or_socket) do |socket_id|
+            @adapter.socket_subscribe(socket_id, *channels)
+            @adapter.expire(socket_id, @timeout_config.initial)
+          end
+        }
+      end
+
+      def socket_unsubscribe(*channels)
+        @executor << -> {
+          @adapter.socket_unsubscribe(*channels)
+        }
+      end
+
+      def subscription_broadcast(channel, message)
+        @executor << -> {
+          @adapter.subscription_broadcast(channel.to_s, channel: channel.name, message: message)
+        }
+      end
+
+      # Called by the adapter, which guarantees that this server has connections for these ids.
+      #
+      def transmit_message_to_connection_ids(message, socket_ids, raw: false)
+        socket_ids.each do |socket_id|
+          find_socket_by_id(socket_id)&.transmit(message, raw: raw)
+        end
+      end
+
+      def find_socket_by_id(socket_id)
+        @sockets.find { |socket| socket.id == socket_id }
+      end
+
+      def find_socket(id_or_socket)
+        socket = if id_or_socket.is_a?(WebSocket)
+          id_or_socket
+        else
+          find_socket_by_id(id_or_socket)
+        end
+
+        yield socket if socket
+      end
+
+      def find_socket_id(id_or_socket)
+        socket_id = if id_or_socket.is_a?(WebSocket)
+          id_or_socket.id
+        else
+          id_or_socket
+        end
+
+        yield socket_id if socket_id
+      end
+
+      private
+
+      def start_heartbeat
+        @heartbeat = Concurrent::TimerTask.new(execution_interval: HEARTBEAT_INTERVAL) do
+          @executor << -> {
+            @sockets.each(&:beat)
+          }
+        end
+
+        @heartbeat.execute
+      end
+
+      def stop_heartbeat
+        @heartbeat.shutdown
+      end
+    end
+  end
+end

data/lib/pakyow/realtime/websocket.rb

@@ -0,0 +1,157 @@
+# frozen_string_literal: true
+
+require "json"
+require "securerandom"
+
+require "pakyow/helpers/app"
+require "pakyow/helpers/connection"
+
+require "async/websocket"
+
+require "protocol/websocket/connection"
+require "protocol/websocket/headers"
+
+module Pakyow
+  module Realtime
+    class WebSocket
+      Frame = ::Protocol::WebSocket::Frame
+
+      class Connection < ::Protocol::WebSocket::Connection
+        include ::Protocol::WebSocket::Headers
+
+        def self.call(framer, protocol = [], **options)
+          return self.new(framer, Array(protocol).first, **options)
+        end
+
+        def initialize(framer, protocol = nil, **options)
+          super(framer, **options)
+          @protocol = protocol
+        end
+
+        attr :protocol
+
+        def call
+          self.close
+        end
+      end
+
+      include Pakyow::Helpers::App
+      include Pakyow::Helpers::Connection
+
+      attr_reader :id
+
+      def initialize(id, connection)
+        @id, @connection, @open = id, connection, false
+        @logger = Logger.new(:sock, id: @id[0..7], output: Pakyow.global_logger, level: Pakyow.config.logger.level)
+        @server = @connection.app.websocket_server
+
+        response = Async::WebSocket::Adapters::Native.open(@connection.request, handler: Connection) do |socket|
+          @socket = socket
+
+          handle_open
+          while message = socket.read
+            handle_message(message)
+          end
+        rescue EOFError, Protocol::WebSocket::ClosedError
+        ensure
+          @socket&.close; shutdown
+        end
+
+        @connection.__getobj__.instance_variable_set(:@response, response)
+      end
+
+      def open?
+        @open == true
+      end
+
+      def transmit(message, raw: false)
+        if open?
+          if raw
+            @socket.write(message)
+          else
+            @socket.write(JSON.dump(payload: message))
+          end
+
+          @socket.flush
+        end
+      end
+
+      def beat
+        transmit("beat")
+      end
+
+      # @api private
+      def leave
+        trigger_presence(:leave)
+      end
+
+      private
+
+      def shutdown
+        if open?
+          @server.socket_disconnect(self)
+          @open = false
+          @logger.info "shutdown"
+        end
+      end
+
+      def handle_open
+        @server.socket_connect(self)
+        @open = true
+        trigger_presence(:join)
+        @logger.info "opened"
+        transmit_system_info
+      end
+
+      def handle_message(message)
+        @logger.internal("< " + message)
+      end
+
+      def trigger_presence(event)
+        @connection.app.hooks(:before, event).each do |hook, _|
+          instance_exec(&hook[:block])
+        end
+      end
+
+      def transmit_system_info
+        transmit(
+          channel: "system",
+          message: {
+            version: @connection.app.config.version
+          }
+        )
+      end
+    end
+  end
+end
+
+module Async
+  module WebSocket
+    module Adapters
+      module Native
+        include ::Protocol::WebSocket::Headers
+
+        def self.websocket?(request)
+          request.headers.include?("upgrade")
+        end
+
+        def self.open(request, headers: [], protocols: [], handler: Connection, **options)
+          if websocket?(request) && Array(request.protocol).include?(PROTOCOL)
+            # Select websocket sub-protocol:
+            if requested_protocol = request.headers[SEC_WEBSOCKET_PROTOCOL]
+              protocol = (requested_protocol & protocols).first
+            end
+
+            Response.for(request, headers, protocol: protocol, **options) do |stream|
+              framer = Protocol::WebSocket::Framer.new(stream)
+
+              yield handler.call(framer, protocol)
+            end
+          else
+            nil
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pakyow/realtime.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require "pakyow/support"
+require "pakyow/routing"
+require "pakyow/presenter"
+
+require "pakyow/realtime/framework"
+
+require "pakyow/environment/realtime/config"
+
+module Pakyow
+  include Environment::Realtime::Config
+end