rubydb 0.1.0

data/lib/rubydb/server/connection.rb

@@ -0,0 +1,366 @@
+# frozen_string_literal: true
+
+require "time"
+require "monitor"
+
+module RubyDB
+  module Server
+    # Connection - Represents a client connection
+    class Connection
+      RequestTooLarge = Class.new(StandardError)
+      RequestTimeout = Class.new(StandardError)
+
+      attr_reader :id, :client, :remote_addr, :remote_port
+      attr_reader :created_at, :last_activity, :state, :session
+      attr_reader :bytes_received, :bytes_sent, :requests_processed
+
+      # Connection states
+      STATE_INIT = :init
+      STATE_HANDSHAKING = :handshaking
+      STATE_AUTHENTICATING = :authenticating
+      STATE_AUTHENTICATED = :authenticated
+      STATE_READY = :ready
+      STATE_BUSY = :busy
+      STATE_CLOSED = :closed
+
+      def initialize(client, id, config = {})
+        @id = id
+        @client = client
+        @remote_addr = client.peeraddr[2] rescue "unknown"
+        @remote_port = client.peeraddr[1] rescue 0
+        @config = config
+        @created_at = Time.now
+        @last_activity = Time.now
+        @state = STATE_INIT
+        @bytes_received = 0
+        @bytes_sent = 0
+        @requests_processed = 0
+        @session = Session.new(self, config) if config[:engine]
+        @session_id = nil
+        @authenticated = false
+        @username = nil
+        @read_timeout = config[:read_timeout] || 30
+        @write_timeout = config[:write_timeout] || 30
+        @idle_timeout = config[:idle_timeout] || 300
+        @buffer = ""
+        @lock = Monitor.new
+        @write_lock = Mutex.new
+        @active_request_lock = Mutex.new
+        @active_request_id = nil
+        @active_request_thread = nil
+        @closed = false
+        @pending_requests = []
+      end
+
+      def handshake(protocol)
+        @lock.synchronize do
+          return false if @closed
+
+          @state = STATE_HANDSHAKING
+
+          begin
+            # Read handshake request
+            data = read_data(protocol)
+
+            # Process handshake
+            response = protocol.handshake.start(data)
+
+            # Send handshake response
+            write_data(protocol, response)
+
+            @state = STATE_AUTHENTICATING
+            true
+
+          rescue => e
+            close
+            false
+          end
+        end
+      end
+
+      def authenticate(credentials)
+        @lock.synchronize do
+          return false if @closed || @state == STATE_CLOSED
+
+          @state = STATE_AUTHENTICATING
+
+          begin
+            # Authenticate
+            @authenticated = @session&.authenticate(credentials) || false
+
+            if @authenticated
+              @state = STATE_AUTHENTICATED
+              @username = credentials[:username]
+              true
+            else
+              @state = STATE_CLOSED
+              close
+              false
+            end
+
+          rescue => e
+            @state = STATE_CLOSED
+            close
+            false
+          end
+        end
+      end
+
+      def ready?
+        @state == STATE_READY
+      end
+
+      def busy?
+        @state == STATE_BUSY
+      end
+
+      def authenticated?
+        @authenticated
+      end
+
+      def closed?
+        @closed
+      end
+
+      def idle?
+        Time.now - @last_activity > @idle_timeout
+      end
+
+      def process_request(request, protocol)
+        @lock.synchronize do
+          return false if @closed
+
+          @state = STATE_BUSY
+          @last_activity = Time.now
+          @requests_processed += 1
+
+          begin
+            # Process request through session
+            response = @session&.process(request) || error_response("No session")
+
+            # Send response
+            write_data(protocol, response)
+
+            @state = STATE_READY
+            true
+
+          rescue => e
+            @state = STATE_READY
+            false
+          end
+        end
+      end
+
+      def start(protocol)
+        @thread = Thread.new { serve(protocol) }
+        true
+      end
+
+      def close
+        @lock.synchronize do
+          return if @closed
+
+          @state = STATE_CLOSED
+          @closed = true
+
+          if @client
+            @client.close rescue nil
+            @client = nil
+          end
+
+          if @session
+            @session.close
+            @session = nil
+          end
+
+          true
+        end
+      end
+
+      def to_hash
+        {
+          id: @id,
+          remote_addr: @remote_addr,
+          remote_port: @remote_port,
+          state: @state,
+          authenticated: @authenticated,
+          username: @username,
+          created_at: @created_at.iso8601,
+          last_activity: @last_activity.iso8601,
+          bytes_received: @bytes_received,
+          bytes_sent: @bytes_sent,
+          requests_processed: @requests_processed,
+          idle_time: (Time.now - @last_activity).round(2)
+        }
+      end
+
+      private
+
+      def serve(protocol)
+        first = protocol.decoder.decode(read_frame, protocol.encoder.format)
+        handshake = protocol.handshake.start(first.payload)
+        write_data(protocol, Protocol::Message.new(Protocol::Message::TYPE_HANDSHAKE_RESPONSE, handshake))
+        return close unless handshake[:success]
+
+        auth = protocol.decoder.decode(read_frame, protocol.encoder.format)
+        authenticated = protocol.handshake.authenticate(auth.payload)
+        @session&.authenticate(auth.payload)
+        write_data(protocol, Protocol::Message.new(Protocol::Message::TYPE_AUTHENTICATION, authenticated.merge(session_id: @session&.id)))
+        return close unless authenticated[:success]
+
+        sync = protocol.decoder.decode(read_frame, protocol.encoder.format)
+        negotiated = protocol.handshake.negotiate(sync.payload[:capabilities] || {})
+        @state = STATE_READY
+        write_data(protocol, Protocol::Message.new(Protocol::Message::TYPE_READY_FOR_QUERY, negotiated))
+
+        request_queue = Queue.new
+        request_worker = Thread.new do
+          while (entry = request_queue.pop)
+            begin
+              message, request = entry
+              response = begin
+                @session.process(request)
+              rescue StandardError => e
+                { success: false, error: e.message, timestamp: Time.now.iso8601 }
+              end
+              response[:request_id] ||= message.id
+              @config[:connection_pool]&.record_request(response[:success] != false)
+              response_type = response[:type] || "#{message.type}_response"
+              # The request is no longer executing once the session has
+              # produced its response. Clear this before writing so a fast
+              # client can issue the next request as soon as it receives the
+              # frame; keeping it set until after the write creates a
+              # connection-level false "request in flight" rejection.
+              @active_request_lock.synchronize do
+                @active_request_id = nil if message && @active_request_id == message.id
+              end
+              begin
+                write_data(protocol, Protocol::Message.new(response_type.to_sym, response))
+              rescue IOError, SystemCallError
+                break
+              end
+            ensure
+              @active_request_lock.synchronize do
+                @active_request_id = nil if message && @active_request_id == message.id
+              end
+            end
+          end
+        end
+        @active_request_lock.synchronize { @active_request_thread = request_worker }
+
+        while !@closed && (line = read_frame)
+          begin
+            message = protocol.decoder.decode(line, protocol.encoder.format)
+          rescue StandardError => e
+            write_data(protocol, Protocol::Message.new(
+              Protocol::Message::TYPE_ERROR,
+              { success: false, error: "Invalid protocol frame: #{e.message}" }
+            ))
+            next
+          end
+          break if message.type.to_s == "terminate"
+          request = message.payload.merge(type: message.type.to_s, request_id: message.id)
+
+          if message.type.to_s == "cancel"
+            active_target = @active_request_lock.synchronize { @active_request_id == request[:target_request_id] }
+            request[:allow_pending_cancellation] = active_target
+            response = begin
+              @session.process(request)
+            rescue StandardError => e
+              { success: false, error: e.message, timestamp: Time.now.iso8601 }
+            end
+            @config[:connection_pool]&.record_request(response[:success] != false)
+            write_data(protocol, Protocol::Message.new(:cancel_response, response))
+            next
+          end
+
+          already_busy = @active_request_lock.synchronize { !@active_request_id.nil? }
+          if already_busy
+            response = { success: false, error: "Connection has a request in flight", code: "busy", request_id: message.id, timestamp: Time.now.iso8601 }
+            write_data(protocol, Protocol::Message.new(:error, response))
+            next
+          end
+
+          @active_request_lock.synchronize { @active_request_id = message.id }
+          request_queue << [message, request]
+        end
+        request_queue << nil
+        request_worker.join(5)
+        @active_request_thread&.join(5)
+      rescue RequestTooLarge => e
+        write_data(protocol, Protocol::Message.new(
+          Protocol::Message::TYPE_ERROR,
+          { success: false, error: e.message }
+        )) rescue nil
+        close unless @closed
+      rescue RequestTimeout => e
+        write_data(protocol, Protocol::Message.new(
+          Protocol::Message::TYPE_ERROR,
+          { success: false, error: e.message }
+        )) rescue nil
+        close unless @closed
+      rescue StandardError => e
+        warn "RubyDB connection #{@id} failed: #{e.class}: #{e.message}" if ENV["RUBYDB_DEBUG"]
+        close unless @closed
+      end
+
+      def read_frame
+        limit = Integer(@config[:max_request_size] || 10 * 1024 * 1024)
+        timeout = @read_timeout
+        if timeout && !IO.select([@client], nil, nil, timeout)
+          raise RequestTimeout, "Request timed out after #{timeout} seconds"
+        end
+        # Passing the separator explicitly keeps the byte limit compatible
+        # with both plain sockets and OpenSSL::SSL::SSLSocket.
+        line = @client.gets("\n", limit + 1)
+        return nil unless line
+
+        if line.bytesize > limit || !line.end_with?("\n")
+          raise RequestTooLarge, "Request frame exceeds maximum size of #{limit} bytes"
+        end
+
+        @bytes_received += line.bytesize
+        line
+      end
+
+      def read_data(protocol)
+        data = ""
+
+        while true
+          chunk = @client.recv(4096)
+          break if chunk.empty?
+          data << chunk
+          @bytes_received += chunk.bytesize
+          break if data.bytesize >= @config[:max_request_size]
+        end
+
+        data
+      end
+
+      def write_data(protocol, data)
+        @write_lock.synchronize do
+          return false if @closed || @client.nil?
+
+          client = @client
+          if data.is_a?(String)
+            client.write(data)
+            @bytes_sent += data.bytesize
+          else
+            encoded = protocol.encoder.encode(data)
+            client.write(encoded)
+            @bytes_sent += encoded.bytesize
+          end
+          client.flush
+        end
+      end
+
+      def error_response(message)
+        {
+          success: false,
+          error: message,
+          timestamp: Time.now.iso8601
+        }
+      end
+    end
+  end
+end

data/lib/rubydb/server/connection_pool.rb

@@ -0,0 +1,193 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+module RubyDB
+  module Server
+    # ConnectionPool - Manages database connections
+    class ConnectionPool
+      attr_reader :stats
+
+      def initialize(config, protocol, engine = nil)
+        @config = config
+        @protocol = protocol
+        @engine = engine
+        @max_connections = config[:max_connections] || 100
+        @connections = {}
+        @connection_counter = 0
+        @stats = {
+          total_connections: 0,
+          active_connections: 0,
+          idle_connections: 0,
+          rejected_connections: 0,
+          connection_errors: 0,
+          requests_processed: 0,
+          requests_failed: 0,
+          avg_connection_time_ms: 0,
+          total_connection_time_ms: 0
+        }
+        @lock = Monitor.new
+        @running = false
+        @cleanup_thread = nil
+        @cleanup_interval = config[:cleanup_interval] || 60
+        @idle_timeout = config[:idle_timeout] || 300
+      end
+
+      def start
+        @lock.synchronize do
+          return if @running
+
+          @running = true
+          @cleanup_thread = Thread.new { cleanup_loop }
+          true
+        end
+      end
+
+      def stop
+        @lock.synchronize do
+          return unless @running
+
+          @running = false
+
+          if @cleanup_thread
+            @cleanup_thread.kill
+            @cleanup_thread = nil
+          end
+
+          close_all
+          true
+        end
+      end
+
+      def add_connection(client)
+        @lock.synchronize do
+          prune_closed_connections
+
+          if @connections.size >= @max_connections
+            client.close
+            @stats[:rejected_connections] += 1
+            return false
+          end
+
+          start_time = Time.now
+          conn_id = next_connection_id
+          connection = Connection.new(client, conn_id, @config.merge(engine: @engine, connection_pool: self))
+
+          @connections[conn_id] = connection
+          @stats[:total_connections] += 1
+          connection.start(@protocol)
+
+          elapsed_ms = (Time.now - start_time) * 1000
+          @stats[:total_connection_time_ms] += elapsed_ms
+          @stats[:avg_connection_time_ms] = @stats[:total_connection_time_ms] / @stats[:total_connections] if @stats[:total_connections] > 0
+
+          true
+        end
+      end
+
+      def get_connection(id)
+        @lock.synchronize do
+          @connections[id]
+        end
+      end
+
+      def remove_connection(id)
+        @lock.synchronize do
+          conn = @connections.delete(id)
+          conn&.close
+          conn
+        end
+      end
+
+      def close_all
+        @lock.synchronize do
+          @connections.each do |id, conn|
+            conn.close
+          end
+          @connections.clear
+        end
+      end
+
+      def can_accept?
+        @lock.synchronize do
+          prune_closed_connections
+          @connections.size < @max_connections
+        end
+      end
+
+      def active_count
+        @lock.synchronize do
+          @connections.values.count { |c| !c.closed? && !c.idle? }
+        end
+      end
+
+      def idle_count
+        @lock.synchronize do
+          @connections.values.count { |c| !c.closed? && c.idle? }
+        end
+      end
+
+      def total_connections
+        @connections.size
+      end
+
+      def running?
+        @running
+      end
+
+      # Report whether every tracked connection is still usable. This is
+      # intentionally based on connection state rather than pool cardinality;
+      # a non-empty pool can still contain only closed connections.
+      def healthy?
+        @lock.synchronize do
+          @running && @connections.values.all? { |connection| !connection.closed? }
+        end
+      end
+
+      def stats
+        @lock.synchronize do
+          @stats.merge({
+            current_connections: @connections.size,
+            active_connections: active_count,
+            idle_connections: idle_count,
+            max_connections: @max_connections,
+            running: @running
+          })
+        end
+      end
+
+      def record_request(success)
+        @lock.synchronize do
+          key = success ? :requests_processed : :requests_failed
+          @stats[key] += 1
+        end
+      end
+
+      private
+
+      def next_connection_id
+        @connection_counter += 1
+      end
+
+      def cleanup_loop
+        while @running
+          sleep(@cleanup_interval)
+
+          @lock.synchronize do
+            @connections.each do |id, conn|
+              # Remove idle connections
+              if conn.idle?
+                conn.close
+                @connections.delete(id)
+              end
+            end
+          end
+        end
+      end
+
+      def prune_closed_connections
+        @connections.delete_if { |_id, connection| connection.closed? }
+      end
+    end
+  end
+end