net-ldap-map-message-id.rb

# Proposal for mapping responses back to requests for Net::LDAP.
require "fiber"

module Net
  module LDAP
    class Connection
      # Hash of Responses keyed by message_id
      attr_reader :responses

      # Fake socket to read/write from
      attr_reader :socket

      def initialize
        @responses = {}
        @message_id = 0
        @socket = Socket.new
      end

      # Returns a Response object. If a block is given, iterate over result
      # PDU's read from `response`.
      def search(&blk)
        message_id = next_message_id
        request = Request.new(message_id)
        response = write(request)
        response.each(&blk) if blk
        response
      end

      def next_message_id
        @message_id += 1
      end

      # Returns an instance of Response immediately after writing `request`.
      def write(request)
        message_id = request.message_id
        response = Response.new(message_id)
        @responses[message_id] = response

        # Create a Fiber to read PDU's. This fiber is resumed in the Response
        # class when someone attempts to read PDU's. We attempt to read a single
        # PDU from the socket, then yield control back to the caller. This would
        # also allow us to abandon an ongoing request.
        socket_fiber = Fiber.new do
          while [email protected]?  # or request_abandoned or request_completed
            pdu = PDU.new(@socket.read)
            @responses[pdu.message_id].enqueue_pdu(pdu)
            Fiber.yield
          end

          # When stream is closed, mark all responses as finished
          @responses.each { |message_id, response| response.finish! }
        end
        response.socket_fiber = socket_fiber

        socket.write(request)
        response
      end
    end

    class Request
      attr_reader :message_id
      def initialize(message_id)
        @message_id = message_id
      end

      def to_s
        "<Request: message_id: #{message_id}>"
      end
    end

    # A Response is a promise for future incoming results for a given message_id
    class Response
      def initialize(message_id)
        @message_id = message_id
        @queued_pdus = []
        @complete = false
      end

      def socket_fiber=(fiber)
        @socket_fiber = fiber
      end

      # Yields PDU to `blk`. This method blocks the socket until all results are
      # read or the socket is closed.
      def each(&blk)
        while !finished?
          @socket_fiber.resume if @socket_fiber.alive?
          if pdu = @queued_pdus.shift
            blk.call(pdu)
          end
        end
      end

      def enqueue_pdu(pdu)
        puts "Response #{@message_id} queued #{pdu}"
        @queued_pdus << pdu
      end

      def finish!
        @complete = true
      end

      # A Response is finished when it has no queued results and the stream has
      # marked it as finished.
      #
      # @socket_fiber will mark this response as finished when the stream has
      # closed. In a real implementation, this would also stop if we received
      # the last search result.
      def finished?
        @complete && @queued_pdus.empty?
      end

      def to_s
        "<Response message_id: #{@message_id} finished?: #{finished?} queued_pdus: #{@queued_pdus}>"
      end
    end

    class PDU
      attr_reader :message_id, :data
      def initialize(raw_data)
        @message_id = raw_data[:message_id]
        @data = raw_data[:payload]
      end

      def to_s
        "<PDU: message_id:#{message_id} data:#{data}>"
      end
    end

    # Simulates a LDAP server responding to two search requests and interleaves
    # results.
    class Socket
      attr_reader :requests  # requests we've seen

      def initialize
        @data = [
          {:message_id => 1, :payload => "1: one"},
          {:message_id => 1, :payload => "1: two"},
          {:message_id => 1, :payload => "1: three"},
          {:message_id => 1, :payload => "1: four"},
          {:message_id => 2, :payload => "2: one"},
          {:message_id => 2, :payload => "2: two"},
          {:message_id => 1, :payload => "1: five"},
          {:message_id => 1, :payload => "1: six"},
          {:message_id => 2, :payload => "2: three"},
        ]
        @requests = []
      end

      def write(request)
        puts "write: #{request}"
        @requests << request
      end

      def read
        pdu = @data.shift
        puts "read: #{pdu}"
        pdu
      end

      def closed?
        @data.empty?
      end
    end
  end
end

results = []
conn = Net::LDAP::Connection.new

# Searching writes the request to the socket, but doesn't attempt to read from it
response1 = conn.search  # first search. message_id 1
response2 = conn.search  # second search. message_id 2

# Although the server may interleave results from different searches, the
# iterators will see results that match their message_id in the order read from
# the stream.
response2.each {|pdu| results << pdu.to_s}
response1.each {|pdu| results << pdu.to_s}

puts "\nResults:\n"
puts results.join("\n")

puts
puts response1

puts
puts response2