neo4j-ruby-driver 1.7.4 → 4.4.0

data/ruby/neo4j/driver/internal/async/network_session.rb

@@ -0,0 +1,150 @@
+module Neo4j::Driver
+  module Internal
+    module Async
+      class NetworkSession
+        attr_reader :retry_logic
+
+        def initialize(connection_provider, retry_logic, database_name, mode, bookmark_holder, impersonated_user, fetch_size, logger)
+          @connection_provider = connection_provider
+          @mode = mode
+          @retry_logic = retry_logic
+          @log = Logging::PrefixedLogger.new("[#{hash}]", logger)
+          @bookmark_holder = bookmark_holder
+          # @database_name = database_name.database_name
+          @connection_context = NetworkSessionConnectionContext.new(database_name, @bookmark_holder.bookmark, impersonated_user)
+          @fetch_size = fetch_size
+          @open = Concurrent::AtomicBoolean.new(true)
+        end
+
+        def run_async(query, **config)
+          new_result_cursor = build_result_cursor_factory(query, config).async_result
+          @result_cursor = new_result_cursor
+          new_result_cursor.map_successful_run_completion_async
+        end
+
+        def begin_transaction_async(mode = @mode, **config)
+          ensure_session_is_open
+          ensure_no_open_tx_before_starting_tx
+          acquire_connection(mode).then do |connection|
+            ImpersonationUtil.ensure_impersonation_support(connection, connection.impersonated_user)
+            tx = UnmanagedTransaction.new(connection, @bookmark_holder, @fetch_size)
+            tx.begin_async(@bookmark_holder.bookmark, config)
+          end&.tap { |new_transaction| @transaction = new_transaction }
+        end
+
+        def reset_async
+          existing_transaction_or_null
+            .then_accept { |tx| tx&.mark_terminated }
+            .then_flat { @connection_stage }
+            .then_flat do |connection|
+            # there exists an active connection, send a RESET message over it
+            connection&.reset || Util::Futures.completed_with_null
+          end
+        end
+
+        def last_bookmark
+          @bookmark_holder.bookmark
+        end
+
+        def release_connection_async
+          @connection&.release
+        end
+
+        def connection_async
+          @connection
+        end
+
+        def open?
+          @open.true?
+        end
+
+        def close_async
+          return unless @open.make_false
+          # there exists a cursor with potentially unconsumed error, try to extract and propagate it
+          error = @result_cursor&.discard_all_failure_async
+          close_transaction_and_release_connection
+        rescue => tx_close_error
+          error = Util::Futures.combine_errors(error, tx_close_error)
+        ensure
+          raise error if error
+        end
+
+        def current_connection_open?
+          @connection&.open? # some connection has actually been acquired and it's still open
+        end
+
+        private
+
+        def build_result_cursor_factory(query, config)
+          ensure_session_is_open
+          ensure_no_open_tx_before_running_query
+          connection = acquire_connection(@mode)
+          ImpersonationUtil.ensure_impersonation_support(connection, connection.impersonated_user)
+          connection.protocol.run_in_auto_commit_transaction(connection, query, @bookmark_holder, config,
+                                                             @fetch_size)
+        end
+
+        def acquire_connection(mode)
+          # make sure previous result is fully consumed and connection is released back to the pool
+          @result_cursor&.pull_all_failure_async&.result!&.then(&method(:raise))
+          if @connection&.open?
+            # there somehow is an existing open connection, this should not happen, just a precondition
+            raise Neo4j::Driver::Exceptions::IllegalStateException.new('Existing open connection detected')
+          end
+
+          @connection = @connection_provider.acquire_connection(@connection_context.context_with_mode(mode))
+        end
+
+        def close_transaction_and_release_connection
+          existing_transaction_or_null&.close_async
+        ensure
+          release_connection_async
+        end
+
+        def ensure_no_open_tx_before_running_query
+          ensure_no_open_tx('Queries cannot be run directly on a session with an open transaction; either run from within the transaction or use a different session.')
+        end
+
+        def ensure_no_open_tx_before_starting_tx
+          ensure_no_open_tx('You cannot begin a transaction on a session with an open transaction; either run from within the transaction or use a different session.')
+        end
+
+        def ensure_no_open_tx(error_message)
+          existing_transaction_or_null&.then do
+            raise Neo4j::Driver::Exceptions::TransactionNestingException.new(error_message)
+          end
+        end
+
+        def existing_transaction_or_null
+          @transaction if @transaction&.open?
+        end
+
+        def ensure_session_is_open
+          unless @open.true?
+            raise Neo4j::Driver::Exceptions::ClientException.new('No more interaction with this session are allowed as the current session is already closed.')
+          end
+        end
+
+        # The {@link NetworkSessionConnectionContext#mode} can be mutable for a session connection context
+        class NetworkSessionConnectionContext
+          # This bookmark is only used for rediscovery.
+          # It has to be the initial bookmark given at the creation of the session.
+          # As only that bookmark could carry extra system bookmarks
+          attr_accessor :database_name
+          attr :mode, :rediscovery_bookmark, :impersonated_user
+
+          def initialize(database_name, bookmark, impersonated_user)
+            @database_name = database_name
+            @rediscovery_bookmark = bookmark
+            @impersonated_user = impersonated_user
+          end
+
+          def context_with_mode(mode)
+            @mode = mode
+            self
+          end
+        end
+      end
+    end
+  end
+end

data/ruby/neo4j/driver/internal/async/outbound/chunk_aware_byte_buf_output.rb

@@ -0,0 +1,110 @@
+module Neo4j::Driver
+  module Internal
+    module Async
+      module Outbound
+        class ChunkAwareByteBufOutput
+          include Packstream::PackStream::Packer
+          include Messaging::Common::CommonValuePacker
+
+          class ChunkBuffer < ::Async::IO::Buffer
+            include Packstream::PackOutput
+            alias write <<
+          end
+
+          def initialize(output, max_chunk_size: Connection::BoltProtocolUtil::DEFAULT_MAX_OUTBOUND_CHUNK_SIZE_BYTES)
+            @output = output
+            @max_chunk_size = verify_max_chunk_size(max_chunk_size)
+            @chunk = ChunkBuffer.new
+          end
+
+          def start
+            assert_not_started
+            @chunk.clear
+          end
+
+          def write_byte(value)
+            ensure_can_fit_in_current_chunk(1)
+            @chunk.write_byte(value)
+            self
+          end
+
+          def write(data)
+            offset = 0
+            length = data.bytesize
+
+            while offset < length
+              # Ensure there is an open chunk, and that it has at least one byte of space left
+              ensure_can_fit_in_current_chunk(1)
+
+              # Write as much as we can into the current chunk
+              amount_to_write = [available_bytes_in_current_chunk, length - offset].min
+
+              @chunk.write(data.byteslice(offset, amount_to_write))
+              offset += amount_to_write
+            end
+
+            self
+          end
+
+          def write_short(value)
+            ensure_can_fit_in_current_chunk(2)
+            @chunk.write_short(value)
+            self
+          end
+
+          def write_int(value)
+            ensure_can_fit_in_current_chunk(4)
+            @chunk.write_int(value)
+            self
+          end
+
+          def write_long(value)
+            ensure_can_fit_in_current_chunk(8)
+            @chunk.write_long(value)
+            self
+          end
+
+          def write_double(value)
+            ensure_can_fit_in_current_chunk(8)
+            @chunk.write_double(value)
+            self
+          end
+
+          def write_message_boundary
+            @output.write_short(0)
+          end
+
+          def write_chunk
+            @output.write_short(@chunk.bytesize)
+            @output.write(@chunk)
+            @chunk.clear
+          end
+
+          alias stop write_chunk
+
+          private
+
+          def ensure_can_fit_in_current_chunk(number_of_bytes)
+            write_chunk if @chunk.bytesize + number_of_bytes > @max_chunk_size
+          end
+
+          def available_bytes_in_current_chunk
+            @max_chunk_size - @chunk.bytesize
+          end
+
+          def assert_not_started
+            raise Neo4j::Driver::Exceptions::IllegalStateException.new('Already started') unless @chunk.empty?
+          end
+
+          def verify_max_chunk_size(max_chunk_size)
+            if max_chunk_size <= 0
+              raise ArgumentError.new("Max chunk size should be > 0, given: #{max_chunk_size}")
+            end
+
+            max_chunk_size
+          end
+        end
+      end
+    end
+  end
+end

data/ruby/neo4j/driver/internal/async/outbound/outbound_message_handler.rb

@@ -0,0 +1,39 @@
+module Neo4j::Driver
+  module Internal
+    module Async
+      module Outbound
+        class OutboundMessageHandler
+          NAME = self.class.name
+
+          def initialize(output, message_format, logger)
+            @output = output
+            @writer = message_format.new_writer(output)
+            @log = logger
+          end
+
+          def handler_added(ctx)
+            @log = Logging::ChannelActivityLogger.new(ctx.channel, @log, self.class)
+          end
+
+          def handler_removed(ctx)
+            @log = nil
+          end
+
+          def encode(msg)
+            @log.debug("C: #{msg}")
+
+            @output.start
+            begin
+              @writer.write(msg)
+            ensure
+              @output.stop
+            end
+
+            @output.write_message_boundary
+            # @log.debug( "C: #{}") if @log.debug_enabled?
+          end
+        end
+      end
+    end
+  end
+end

data/ruby/neo4j/driver/internal/async/pool/channel.rb

@@ -0,0 +1,62 @@
+module Neo4j::Driver
+  module Internal
+    module Async
+      module Pool
+        class Channel
+          attr :stream
+          attr_accessor :version, :protocol, :message_format, :message_dispatcher
+          attr :attributes # should be attr
+          attr_accessor :auto_read
+
+          def initialize(address, connector, logger)
+            @attributes = Connection::ChannelAttributes.new
+            @stream = Connection::Stream.new(connector.connect(address))
+            @stream.write(Connection::BoltProtocolUtil.handshake_buf)
+            @stream.flush
+            Connection::HandshakeHandler.new(logger).decode(self)
+            stream_reader = Connection::StreamReader.new(@stream)
+            stream_writer = Outbound::ChunkAwareByteBufOutput.new(@stream)
+            @message_dispatcher = Inbound::InboundMessageDispatcher.new(self, logger)
+            @attributes[:message_dispatcher] = @message_dispatcher
+            @outbound_handler = Outbound::OutboundMessageHandler.new(stream_writer, message_format, logger)
+            @common_message_reader = Messaging::Common::CommonMessageReader.new(stream_reader)
+            connector.initialize_channel(self, protocol)
+          end
+
+          def close
+            @stream.close
+          end
+
+          def write(message)
+            @outbound_handler.encode(message)
+          end
+
+          def write_and_flush(message)
+            write(message)
+            @stream.flush
+            ensure_response_handling
+          end
+
+          private
+
+          def bracketless(host)
+            host.delete_prefix('[').delete_suffix(']')
+          end
+
+          def ensure_response_handling
+            # probably should be synchronized
+            return if @handling_active
+            @handling_active = true
+            while @message_dispatcher.queued_handlers_count > 0 do
+              @common_message_reader.read(@message_dispatcher)
+            end
+            @handling_active = false
+          rescue
+            @handling_active = false
+            raise
+          end
+        end
+      end
+    end
+  end
+end

data/ruby/neo4j/driver/internal/async/pool/channel_pool.rb

@@ -0,0 +1,31 @@
+module Neo4j::Driver
+  module Internal
+    module Async
+      module Pool
+        class ChannelPool < ConnectionPool
+          def initialize(limit: nil, acquisition_timeout: nil, &block)
+            super(size: limit, timeout: acquisition_timeout, &block)
+            @available = TimedStack.new(@size, &block)
+          end
+
+          def acquire(options = {})
+            @available.pop(options[:timeout] || @timeout)
+          end
+
+          def release(resource)
+            @available.push(resource)
+            nil
+          end
+
+          def close
+            @available.shutdown(&:close)
+          end
+
+          def busy?
+            @available.any_resource_busy?
+          end
+        end
+      end
+    end
+  end
+end

data/ruby/neo4j/driver/internal/async/pool/channel_tracker.rb

@@ -0,0 +1,135 @@
+module Neo4j::Driver
+  module Internal
+    module Async
+      module Pool
+        class ChannelTracker
+          attr_reader :lock, :read, :write, :address_to_in_use_channel_count, :address_to_idle_channel_count, :log,
+                      :metrics_listener, :close_listener, :all_channels
+
+          def initialize(metrics_listener, logger, **options)
+            @metrics_listener = metrics_listener
+            @log = logger
+            @all_channels = options[:channels]
+            @lock = Concurrent::ReentrantReadWriteLock.new
+            @close_listener = method(:channel_closed)
+          end
+
+          def channel_released(channel)
+            @lock.with_write_lock do
+              decrement_in_use(channel)
+              increment_idle(channel)
+              channel.close_future.add_listener(close_listener)
+            end
+
+            log.debug("Channel [0x#{channel.id}] acquired from the pool. Local address: #{channel.local_address}, remote address: #{channel.remote_address}")
+          end
+
+          def channel_created(channel, creating_event = nil)
+            if creating_event.nil?
+              raise Neo4j::Driver::Exceptions::IllegalStateException.new('Untraceable channel created.')
+            else
+              # when it is created, we count it as idle as it has not been acquired out of the pool
+              do_in_write_lock(->() { increment_idle(channel) })
+
+              metrics_listener.after_created(Connection::ChannelAttributes.pool_id(channel), creating_event)
+              all_channels.add(channel)
+              log.debug( "Channel [0x#{channel.id}] created. Local address: #{channel.local_address}, remote address: #{channel.remote_address}")
+            end
+          end
+
+          def channel_creating(pool_id)
+            creating_event = metrics_listener.create_listener_event
+            metrics_listener.before_creating(pool_id, creating_event)
+            creating_event
+          end
+
+          def channel_failed_to_create(pool_id)
+            metrics_listener.after_failed_to_create(pool_id)
+          end
+
+          def channel_closed(channel)
+            do_in_write_lock(-> () { decrement_idle(channel) })
+            metrics_listener.after_closed(Connection::ChannelAttributes.pool_id(channel))
+          end
+
+          def in_use_channel_count(address)
+            retrieve_in_read_lock(-> () { address_to_in_use_channel_count.get_or_default(address, 0) })
+          end
+
+          def idle_channel_count(address)
+            retrieve_in_read_lock(-> () { address_to_idle_channel_count.get_or_default(address, 0) })
+          end
+
+          def prepare_to_close_channels
+            all_channels.each do |channel|
+              protocol = Messaging::BoltProtocol.for_channel(channel)
+              begin
+                protocol.prepare_to_close_channel(channel)
+              rescue Exception => e
+                # only logging it
+                log.debug( "Failed to prepare to close Channel #{channel} due to error #{e.get_message}. It is safe to ignore this error as the channel will be closed despite if it is successfully prepared to close or not.")
+              end
+            end
+          end
+
+
+          private
+
+          def increment_in_use(channel)
+            increment(channel, address_to_in_use_channel_count)
+          end
+
+          def decrement_in_use(channel)
+            address = Connection::ChannelAttributes.server_address(channel)
+
+            unless address_to_in_use_channel_count.contains_key(address)
+              raise Neo4j::Driver::Exceptions::IllegalStateException.new("No count exists for address '#{address}' in the 'in use' count")
+            end
+
+            count = address_to_in_use_channel_count.get(address)
+            address_to_in_use_channel_count.put(address, count - 1)
+          end
+
+          def increment_idle(channel)
+            increment(channel, address_to_idle_channel_count)
+          end
+
+          def decrement_idle(channel)
+            address = Connection::ChannelAttributes.server_address(channel)
+
+            unless address_to_idle_channel_count.contains_key(address)
+              raise Neo4j::Driver::Exceptions::IllegalStateException.new("No count exists for address '#{address}' in the 'idle' count")
+            end
+
+            count = address_to_idle_channel_count.get(address)
+            address_to_idle_channel_count.put(address, count - 1)
+          end
+
+          def increment(channel, count_map)
+            address = Connection::ChannelAttributes.server_address(channel)
+            count = count_map.compute_if_absent(address, -> (k) { 0 })
+            count_map.put(address, count + 1)
+          end
+
+          def do_in_write_lock(work)
+            begin
+              write.lock
+              work.run
+            ensure
+              write.unlock
+            end
+          end
+
+          def retrieve_in_read_lock(work)
+            begin
+              read.lock
+              work.get
+            ensure
+              read.unlock
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/ruby/neo4j/driver/internal/async/pool/connection_pool_impl.rb

@@ -0,0 +1,156 @@
+module Neo4j::Driver
+  module Internal
+    module Async
+      module Pool
+        class ConnectionPoolImpl
+          def initialize(connector, settings, logger)
+            @connector = connector
+            @settings = settings
+            @log = logger
+            @address_to_pool_lock = Concurrent::ReentrantReadWriteLock.new
+            @address_to_pool = {}
+            @closed = Concurrent::AtomicBoolean.new
+          end
+
+          def acquire(address)
+            @log.debug("Acquiring a connection from pool towards #{address}")
+
+            assert_not_closed
+            pool = get_or_create_pool(address)
+
+            begin
+              channel = pool.acquire
+              @log.debug { "Channel #{channel.object_id} acquired" }
+            rescue => error
+              process_acquisition_error(pool, address, error)
+            end
+            assert_not_closed(address, channel, pool)
+            NetworkConnection.new(channel, pool, @log) { remove(pool) }
+          end
+
+          def retain_all(addresses_to_retain)
+            @address_to_pool_lock.with_write_lock do
+              @address_to_pool.each do |address, pool|
+                unless addresses_to_retain.include?(address)
+                  unless pool.busy?
+                    # address is not present in updated routing table and has no active connections
+                    # it's now safe to terminate corresponding connection pool and forget about it
+                    @address_to_pool.delete(address)
+                    if pool
+                      @log.info("Closing connection pool towards #{address}, it has no active connections and is not in the routing table registry.")
+                      close_pool_in_background(address, pool)
+                    end
+                  end
+                end
+              end
+            end
+          end
+
+          def remove(pool)
+            @address_to_pool_lock.with_write_lock do
+              @address_to_pool.each { |address, value| @address_to_pool.delete(address) if value == pool }
+            end
+          end
+
+          def in_use_connections(address)
+            @address_to_pool[address]&.size || 0
+            # @netty_channel_tracker.in_use_channel_count(address)
+          end
+
+          def idle_connections(address)
+            @netty_channel_tracker.idle_channel_count(address)
+          end
+
+          def close
+            if @closed.make_true
+              @address_to_pool_lock.with_write_lock do
+                # We can only shutdown event loop group when all netty pools are fully closed,
+                # otherwise the netty pools might missing threads (from event loop group) to execute clean ups.
+                close_all_pools
+                @address_to_pool.clear
+              end
+            end
+          end
+
+          def open?(address)
+            @address_to_pool_lock.with_read_lock { @address_to_pool.key?(address) }
+          end
+
+          def to_string
+            @address_to_pool_lock.with_read_lock { "ConnectionPoolImpl{ pools=#{@address_to_pool}}" }
+          end
+
+          private
+
+          def process_acquisition_error(pool, server_address, error)
+            if error.is_a?(ConnectionPool::TimeoutError)
+              # NettyChannelPool returns future failed with TimeoutException if acquire operation takes more than
+              # configured time, translate this exception to a prettier one and re-throw
+              raise Neo4j::Driver::Exceptions::ClientException.new("Unable to acquire connection from the pool within configured maximum time of #{@settings.connection_acquisition_timeout.inspect}")
+              # elsif pool.closed?
+              # There is a race condition where a thread tries to acquire a connection while the pool is closed by another concurrent thread.
+              # Treat as failed to obtain connection for a direct driver. For a routing driver, this error should be retried.
+              # raise Neo4j::Driver::Exceptions::ServiceUnavailableException, "Connection pool for server #{server_address} is closed while acquiring a connection."
+            else
+              # some unknown error happened during connection acquisition, propagate it
+              raise error
+            end
+          end
+
+          def assert_not_closed(address = nil, channel = nil, pool = nil)
+            if @closed.true?
+              if address
+                pool.release(channel)
+                close_pool_in_background(address, pool)
+                @address_to_pool_lock.with_write_lock { @address_to_pool.delete(address) }
+                assert_not_closed
+              end
+              raise Exceptions::IllegalStateException, Spi::ConnectionPool::CONNECTION_POOL_CLOSED_ERROR_MESSAGE
+            end
+          end
+
+          # for testing only
+          protected def pool(address)
+            @address_to_pool_lock.with_read_lock { @address_to_pool[address] }
+          end
+
+          def new_pool(address)
+            ChannelPool.new(limit: @settings.max_connection_pool_size, acquisition_timeout: @settings.connection_acquisition_timeout) { Channel.new(address, @connector, @log) }
+          end
+
+          def get_or_create_pool(address)
+            @address_to_pool_lock.with_read_lock { @address_to_pool[address] } ||
+              @address_to_pool_lock.with_write_lock do
+                new_pool(address)&.tap do |pool|
+                  # before the connection pool is added I can add the metrics for the pool.
+                  # @metrics_listener.put_pool_metrics(pool.object_id, address, self)
+                  @address_to_pool[address] = pool
+                end
+              end
+          end
+
+          def close_pool(pool)
+            pool.close
+          end
+
+          def close_pool_in_background(address, pool)
+            Async do
+              # Close in the background
+              close_pool(pool)
+            rescue => error
+              @log.warn("An error occurred while closing connection pool towards #{address}.", error)
+            end
+          end
+
+          def close_all_pools
+            @address_to_pool.map do |address, pool|
+              @log.info("Closing connection pool towards #{address}")
+              # Wait for all pools to be closed.
+              close_pool(pool)
+            end
+          end
+        end
+      end
+    end
+  end
+end