zkruby 3.4.3

data/lib/zkruby/rubyio.rb

@@ -0,0 +1,310 @@
+require 'socket'
+require 'thread'
+require 'monitor'
+
+# Binding over standard ruby sockets
+#
+# Manages 3 threads per zookeeper session
+#
+# Read thread
+#   manages connecting to and reading from the tcp socket. Uses non blocking io to manage timeouts
+#   and initiate the required ping requests.
+#
+# Write thread
+#   each new connection spawns a new thread. Requests coming from the session in response
+#   to multiple threads are written to a blocking queue. While the connection is alive
+#   this thread reads from the queue and writes to the socket, all in blocking fashion
+#   TODO: Is it really ok to do a non-blocking read during a blocking write?
+# 
+# Event thread
+#   All response and watch callbacks are put on another blocking queue to be read and executed
+#   by this thread.
+#
+# All interaction with the session is synchronized
+#
+# Client synchronous code is implemented with a condition variable that waits on the callback/errback 
+module ZooKeeper::RubyIO
+
+    class Connection
+        include ZooKeeper::Protocol
+        include Slf4r::Logger
+        include Socket::Constants
+
+        def initialize(host,port,timeout,session)
+            @session = session
+            @write_queue = Queue.new()
+
+            # JRuby cannot do non-blocking connects, which means there is
+            # no way to properly implement the connection-timeout
+            # See http://jira.codehaus.org/browse/JRUBY-5165 
+            # In any case this should be encapsulated in TCPSocket.open(host,port,timeout)
+            if RUBY_PLATFORM == "java"
+                begin
+                    sock = TCPSocket.new(host,port.to_i)
+                rescue Errno::ECONNREFUSED
+                    logger.warn("TCP Connection refused to #{host}:#{port}")
+                    sock = nil
+                end
+            else
+                addr = Socket.getaddrinfo(host, nil)
+                sock = Socket.new(Socket.const_get(addr[0][0]), Socket::SOCK_STREAM, 0)
+                sock.setsockopt(SOL_SOCKET, SO_LINGER, [0,-1].pack("ii"))
+                sock.setsockopt(SOL_TCP, TCP_NODELAY,[0].pack("i_"))
+                sockaddr = Socket.pack_sockaddr_in(port, addr[0][3])
+                begin
+                    sock.connect_nonblock(sockaddr)
+                rescue Errno::EINPROGRESS
+                    resp = IO.select(nil, [sock], nil, timeout)
+                    begin
+                        sock.connect_nonblock(sockaddr)
+                    rescue Errno::ECONNREFUSED
+                        logger.warn("Connection refused to #{ host }:#{ port }")
+                        sock = nil
+                    rescue Errno::EISCONN
+                    end
+                end
+            end
+            @socket = sock
+            Thread.new(sock) { |sock| write_loop(sock) } if sock
+        end
+
+        # This is called from random client threads, but only within
+        # a @session.synchronized() block
+        def send_data(data)
+            @write_queue.push(data)
+        end
+
+        # Since this runs in its very own thread
+        # we can use boring blocking IO
+        def write_loop(socket)
+            Thread.current[:name] = "ZooKeeper::RubyIO::WriteLoop"
+            begin
+                while socket
+                    data = @write_queue.pop()
+                    if socket.write(data) != data.length()
+                        #TODO - will this really ever happen
+                        logger.warn("Incomplete write!")
+                    end
+                    logger.debug { "Sending: " + data.unpack("H*")[0] }
+                end
+            rescue Exception => ex
+                logger.warn("Exception in write loop",ex)
+                disconnect()
+            end
+
+        end
+
+        def read_loop()
+            socket = @socket
+            ping = 0
+            while socket # effectively forever
+                begin
+                    data = socket.read_nonblock(1024)
+                    logger.debug { "Received (#{data.length})" + data.unpack("H*")[0] }
+                    receive_data(data)
+                    ping = 0
+                rescue IO::WaitReadable
+                    select_result = IO.select([socket],[],[],@session.ping_interval)
+                    unless select_result
+                        ping += 1
+                        # two timeouts in a row mean we need to send a ping
+                        case ping
+                        when 1 ; @session.synchronize { @session.ping() }
+                        when 2
+                            logger.debug{"No response to ping in #{@session.ping_interval}*2"}
+                            break
+                        end
+                    end
+                rescue EOFError
+                    logger.debug { "EOF reading from socket" }
+                    break
+                rescue Exception => ex
+                    logger.warn( "#{ex.class} exception in readloop",ex )
+                    break
+                end
+            end
+            disconnect()
+        end
+
+        def disconnect()
+            socket = @socket
+            @socket = nil
+            socket.close if socket
+        rescue Exception => ex
+            #oh well
+            logger.debug("Exception closing socket",ex)
+        end
+
+        # Protocol requirement
+        def receive_records(packet_io)
+            @session.synchronize { @session.receive_records(packet_io) }
+        end
+
+    end #Class connection
+
+    class Binding
+        include Slf4r::Logger
+        attr_reader :session
+        def self.available?
+            true
+        end
+
+        def self.context(&context_block)
+           yield Thread 
+        end
+
+        def initialize()
+            @event_queue = Queue.new()
+        end
+
+        def pop_event_queue()
+            queued = @event_queue.pop()
+            return false  unless queued
+            logger.debug { "Invoking #{queued[0]}" }
+            callback,*args = queued
+            callback.call(*args)
+            logger.debug { "Completed #{queued[0]}" }
+            return true
+        rescue Exception => ex
+            logger.warn( "Exception in event thread", ex )
+            return true
+        end
+
+        def start(client,session)
+            @session = session
+            @session.extend(MonitorMixin)
+
+            # start the event thread
+            @event_thread = Thread.new() do
+                Thread.current[:name] = "ZooKeeper::RubyIO::EventLoop"
+
+                # In this thread, the current client is always this client!
+                Thread.current[ZooKeeper::CURRENT] = [client]
+                loop do
+                    break unless pop_event_queue() 
+                end
+            end
+
+            # and the read thread
+            Thread.new() do
+                begin
+                    Thread.current[:name] = "ZooKeeper::RubyIO::ReadLoop"
+                    conn = session.synchronize { session.start(); session.conn() } # will invoke connect 
+                    loop do
+                        break unless conn
+                        conn.read_loop()
+                        conn =  session.synchronize { session.disconnected(); session.conn() }
+                    end
+                    #event of death
+                    logger.debug("Pushing nil (event of death) to event queue")
+                    @event_queue.push(nil)
+                rescue Exception => ex
+                    logger.error( "Exception in session thread", ex )
+                end
+            end
+        end
+
+        # session callback, IO thread
+        def connect(host,port,delay,timeout)
+            sleep(delay)
+            conn = Connection.new(host,port,timeout,session)
+            session.synchronize() { session.prime_connection(conn) }
+        end
+
+
+        def close(&callback)
+            op = AsyncOp.new(self,&callback)
+            
+            session.synchronize do
+                    session.close() do |error,response|
+                        op.resume(error,response)
+                    end
+            end
+
+            op
+            
+        end
+
+        def queue_request(*args,&callback)
+
+            op = AsyncOp.new(self,&callback)
+            
+            session.synchronize do 
+                session.queue_request(*args) do |error,response|
+                    op.resume(error,response)                    
+                end
+            end
+
+            op
+        end
+
+        def event_thread?
+            Thread.current.equal?(@event_thread)
+        end
+
+        def invoke(*args)
+            @event_queue.push(args)
+        end
+
+    end #Binding
+
+    class AsyncOp < ::ZooKeeper::AsyncOp
+
+        def initialize(binding,&callback)
+            @mutex = Monitor.new
+            @cv = @mutex.new_cond()
+            @callback = callback
+            @rubyio = binding
+        end
+
+        def resume(error,response)
+            mutex.synchronize do
+                @error = error
+                @result = nil
+                begin
+                    @result = @callback.call(response) unless error
+                rescue StandardError => ex
+                    @error = ex
+                end
+               
+                if @error && @errback
+                    begin
+                        @result = @errback.call(@error) 
+                        @error = nil
+                    rescue StandardError => ex
+                        @error = ex
+                    end
+                end
+
+                cv.signal()
+            end
+        end
+
+        private
+            attr_reader :mutex, :cv, :error, :result
+
+        def set_error_handler(errback)
+            @errback = errback
+        end
+     
+        def wait_value()
+            if @rubyio.event_thread?
+                #Waiting in the event thread (eg made a synchronous call inside a callback)
+                #Keep processing events until we are resumed
+                until defined?(@result)
+                    break unless @rubyio.pop_event_queue()
+                end
+            else
+                mutex.synchronize do
+                    cv.wait()
+                end
+            end
+
+            raise error if error
+            result
+        end
+    end
+
+end #ZooKeeper::RubyIO
+# Add our binding
+ZooKeeper.add_binding(ZooKeeper::RubyIO::Binding)

data/lib/zkruby/session.rb

@@ -0,0 +1,445 @@
+require 'set'
+module ZooKeeper
+
+    # Represents an session that may span connections
+    class Session
+
+        DEFAULT_TIMEOUT = 4
+        DEFAULT_CONNECT_DELAY = 0.2
+        DEFAULT_PORT = 2181
+
+        include Slf4r::Logger
+
+        attr_reader :ping_interval
+        attr_reader :ping_logger
+        attr_reader :timeout
+        attr_reader :conn
+        attr_accessor :watcher
+
+        def initialize(binding,addresses,options=nil)
+
+            @binding = binding
+
+            @addresses = parse_addresses(addresses)
+            parse_options(options)
+
+            # These are the server states
+            # :disconnected, :connected, :auth_failed, :expired
+            @keeper_state = nil
+
+            # Client state is
+            # :ready, :closing, :closed
+            @client_state = :ready
+
+            @xid=0
+            @pending_queue = []
+
+            # Create the watch list
+            # hash by watch type of hashes by path of set of watchers
+            @watches = [ :children, :data, :exists ].inject({}) do |ws,wtype| 
+                ws[wtype] = Hash.new() { |h,k| h[k] = Set.new() }
+                ws
+            end
+
+            @watcher = nil
+
+            @ping_logger = Slf4r::LoggerFacade.new("ZooKeeper::Session::Ping")
+
+        end
+
+        def chroot(path)
+            return @chroot + path
+        end
+
+        def unchroot(path)
+            return path unless path
+            path.slice(@chroot.length..-1)
+        end
+
+        # close won't run your block if the connection is 
+        # already closed, so this is how you can check
+        def closed?
+            @client_state == :closed
+        end
+
+        # Connection API - testing whether to send a ping
+        def connected?()
+            @keeper_state == :connected
+        end
+
+        # Connection API - Injects a new connection that is ready to receive records
+        # @param conn that responds to #send_records(record...) and #disconnect()
+        def prime_connection(conn)
+            @conn = conn
+            send_session_connect()
+            send_auth_data()
+            reset_watches()
+        end
+
+
+        # Connection API - called when data is available, reads and processes one packet/event
+        # @param io <IO> 
+        def receive_records(io)
+            case @keeper_state
+            when :disconnected
+                complete_connection(io)
+            when :connected
+                process_reply(io)
+            else
+                logger.warn { "Receive packet for closed session #{@keeper_state}" }
+            end
+        end
+
+        # Connection API - called when no data has been received for #ping_interval
+        def ping()
+            if @keeper_state == :connected
+                ping_logger.debug { "Ping send" }
+                hdr = Proto::RequestHeader.new(:xid => -2, :_type => 11)
+                conn.send_records(hdr)
+            end
+        end
+
+        # Connection API - called when the connection has dropped from either end
+        def disconnected()
+            @conn = nil
+            logger.info { "Disconnected id=#{@session_id}, keeper=:#{@keeper_state}, client=:#{@client_state}" }
+
+            # We keep trying to reconnect until the session expiration time is reached
+            @disconnect_time = Time.now if @keeper_state == :connected
+            time_since_first_disconnect = (Time.now - @disconnect_time) 
+
+            if @client_state == :closed || time_since_first_disconnect > timeout
+                session_expired()
+            else
+                # if we are connected then everything in the pending queue has been sent so
+                # we must clear
+                # if not, then we'll keep them and hope the next reconnect works
+                if @keeper_state == :connected
+                    clear_pending_queue(:disconnected)
+                    invoke_watch(@watcher,KeeperState::DISCONNECTED,nil,WatchEvent::NONE) if @watcher
+                end
+                @keeper_state = :disconnected
+                reconnect()
+            end
+        end
+
+        # Start the session - called by the ProtocolBinding
+        def start()
+            raise ProtocolError, "Already started!" unless @keeper_state.nil?
+            @keeper_state = :disconnected
+            @disconnect_time = Time.now
+            logger.debug("Starting new zookeeper client session")
+            reconnect()
+        end
+
+        def queue_request(request,op,opcode,response=nil,watch_type=nil,watcher=nil,ptype=Packet,&callback)
+            raise Error.SESSION_EXPIRED, "Session expired due to client state #{@client_state}"  unless @client_state == :ready
+            watch_type, watcher = resolve_watcher(watch_type,watcher)
+
+            xid = next_xid
+
+            packet = ptype.new(xid,op,opcode,request,response,watch_type,watcher, callback)
+
+            queue_packet(packet)
+        end
+
+        def close(&callback)
+            case @client_state
+            when :ready
+                # we keep the requested block in a close packet
+                @close_packet = ClosePacket.new(next_xid(),:close,-11,nil,nil,nil,nil,callback)
+                close_packet = @close_packet 
+                @client_state = :closing
+
+                # If there are other requests in flight, then we wait for them to finish
+                # before sending the close packet since it immediately causes the socket
+                # to close.
+                queue_close_packet_if_necessary()
+                @close_packet
+            when :closed, :closing
+                raise ProtocolError, "Already closed"
+            else
+                raise ProtocolError, "Unexpected state #{@client_state}"
+            end
+        end
+
+        private
+        attr_reader :watches
+        attr_reader :binding
+
+        def parse_addresses(addresses)
+            case addresses
+            when String
+                parse_addresses(addresses.split(","))
+            when Array
+                result = addresses.collect() { |addr| parse_address(addr) }
+                #Randomise the connection order
+                result.shuffle!
+            else
+                raise ArgumentError "Not able to parse addresses from #{addresses}"
+            end
+        end
+
+        def parse_address(address)
+            case address
+            when String
+                host,port = address.split(":")
+                port = DEFAULT_PORT unless port
+                [host,port]
+            when Array
+                address[0..1]
+            end
+        end
+
+        def parse_options(options)
+            @timeout = options.fetch(:timeout,DEFAULT_TIMEOUT)
+            @max_connect_delay = options.fetch(:connect_delay,DEFAULT_CONNECT_DELAY)
+            @connect_timeout = options.fetch(:connect_timeout,@timeout * 1.0 / 7.0)
+            @scheme = options.fetch(:scheme,nil)
+            @auth = options.fetch(:auth,nil)
+            @chroot = options.fetch(:chroot,"").chomp("/")
+        end
+
+        def reconnect()
+
+            #Rotate address
+            host,port = @addresses.shift
+            @addresses.push([host,port])
+
+            delay = rand() * @max_connect_delay
+
+            logger.debug { "Connecting id=#{@session_id} to #{host}:#{port} with delay=#{delay}, timeout=#{@connect_timeout}" } 
+            binding.connect(host,port,delay,@connect_timeout)
+        end
+
+
+        def session_expired(reason=:expired)
+            clear_pending_queue(reason)
+
+            invoke_response(*@close_packet.error(reason)) if @close_packet
+
+            if @client_state == :closed
+                logger.info { "Session closed id=#{@session_id}, keeper=:#{@keeper_state}, client=:#{@client_state}" }
+            else
+                logger.warn { "Session expired id=#{@session_id}, keeper=:#{@keeper_state}, client=:#{@client_state}" }
+            end
+
+            invoke_watch(@watcher,KeeperState::EXPIRED,nil,WatchEvent::NONE) if @watcher
+            @keeper_state = reason
+            @client_state = :closed
+        end
+
+        def complete_connection(response)
+            result = Proto::ConnectResponse.read(response)
+            if (result.time_out <= 0)
+                #We're dead!
+                session_expired()
+            else
+                @timeout = result.time_out.to_f / 1000.0
+                @keeper_state = :connected
+
+                # Why 2 / 7 of the timeout?. If a binding sees no server response in this period it is required to
+                # generate a ping request
+                # if 2 periods go by without activity it is required to disconnect
+                # so we are already more than half way through the session timeout
+                # and we need to give ourselves time to reconnect to another server
+                @ping_interval = @timeout * 2.0 / 7.0
+                @session_id = result.session_id
+                @session_passwd = result.passwd
+                logger.info { "Connected session_id=#{@session_id}, timeout=#{@time_out}, ping=#{@ping_interval}" }
+
+                logger.debug { "Sending #{@pending_queue.length} queued packets" }
+                @pending_queue.each { |p| send_packet(p) }
+
+                queue_close_packet_if_necessary()
+                invoke_watch(@watcher,KeeperState::CONNECTED,nil,WatchEvent::NONE) if @watcher
+            end
+        end
+
+        def send_session_connect() 
+            req = Proto::ConnectRequest.new( :timeout => timeout )
+            req.last_zxid_seen = @last_zxid_seen if @last_zxid_seen
+            req.session_id =  @session_id if @session_id
+            req.passwd = @session_passwd if @session_passwd
+
+            conn.send_records(req)
+        end
+
+        def send_auth_data()
+            if @scheme
+                req = Proto::AuthPacket.new(:scheme => @scheme, :auth => @auth)
+                packet = Packet.new(-4,:auth,100,req,nil,nil,nil,nil)
+                send_packet(packet)
+            end
+        end
+        # Watches are dropped on disconnect, we reset them here
+        # dropping connections can be a good way of cleaning up on the server side
+        # #TODO If watch reset is disabled the watches will be notified of connection loss
+        # otherwise they will be seemlessly re-added
+        # This way a watch is only ever triggered exactly once
+        def reset_watches()
+            unless watches[:children].empty? && watches[:data].empty? && watches[:exists].empty?
+                req = Proto::SetWatches.new()
+                req.relative_zxid = @last_zxid_seen
+                req.data_watches = watches[:data].keys
+                req.child_watches = watches[:children].keys
+                req.exist_watches = watches[:exists].keys
+
+                packet = Packet.new(-8,:set_watches,101,req,nil,nil,nil,nil)
+                send_packet(packet)
+            end
+        end
+
+        def process_reply(packet_io)
+            header = Proto::ReplyHeader.read(packet_io)
+
+            case header.xid.to_i
+            when -2
+                ping_logger.debug { "Ping reply" }
+            when -4
+                logger.debug { "Auth reply" }
+                session_expired(:auth_failed) unless header.err.to_i == 0
+            when -1
+                #Watch notification
+                event = Proto::WatcherEvent.read(packet_io)
+                logger.debug { "Watch notification #{event.inspect} " }
+                process_watch_notification(event.state.to_i,event.path,event._type.to_i)
+            when -8
+                #Reset watch reply
+                logger.debug { "SetWatch reply"}
+                #TODO If error, send :disconnected to all watches
+            else
+                # A normal packet reply. They should come in the order we sent them
+                # so we just match it to the packet at the front of the queue
+                packet = @pending_queue.shift
+                logger.debug { "Packet reply: #{packet.inspect}" }
+
+                if (packet.xid.to_i != header.xid.to_i)
+
+                    logger.error { "Bad XID! expected=#{packet.xid}, received=#{header.xid}" }
+
+                    # Treat this like a dropped connection, and then force the connection
+                    # to be dropped. But wait for the connection to notify us before
+                    # we actually update our keeper_state
+                    invoke_response(*packet.error(:disconnected))
+                    @conn.disconnect() 
+                else
+
+                    @last_zxid_seen = header.zxid
+                    
+                    callback, error, response, watch_type =  packet.result(header.err.to_i)
+                    invoke_response(callback, error, response, packet_io)
+
+                    if (watch_type) 
+                        @watches[watch_type][packet.path] << packet.watcher 
+                        logger.debug { "Registered #{packet.watcher} for type=#{watch_type} at #{packet.path}" }
+                    end
+                    queue_close_packet_if_necessary()
+                end
+            end
+        end
+
+
+        def process_watch_notification(state,path,event)
+
+            watch_event = WatchEvent.fetch(event)
+            watch_types = watch_event.watch_types()
+
+            keeper_state = KeeperState.fetch(state)
+
+            watches = watch_types.inject(Set.new()) do | result, watch_type |
+                more_watches = @watches[watch_type].delete(path)
+            result.merge(more_watches) if more_watches
+            result
+            end
+
+            if watches.empty?
+                logger.warn { "Received notification for unregistered watch #{state} #{path} #{event}" }
+            end
+            watches.each { | watch | invoke_watch(watch,keeper_state,path,watch_event) }      
+
+        end
+
+        def invoke_watch(watch,state,path,event)
+            logger.debug { "Watch #{watch} triggered with #{state}, #{path}. #{event}" }
+            if watch.respond_to?(:process_watch)
+                callback = Proc.new() { |state,path,event| watch.process_watch(state,path,event) } 
+            elsif watch.respond_to?(:call)
+                callback = watch
+            else
+                raise ProtocolError("Bad watcher #{watch}")
+            end
+
+            binding.invoke(callback,state,unchroot(path),event)
+        end
+
+        def clear_pending_queue(reason)
+            @pending_queue.each  { |p| invoke_response(*p.error(reason)) }
+            @pending_queue.clear
+        end
+
+        def queue_close_packet_if_necessary
+            if @pending_queue.empty? && @keeper_state == :connected && @close_packet
+                logger.debug { "Sending close packet!" }
+                @client_state = :closed
+                queue_packet(@close_packet)
+                @close_packet = nil
+            end
+        end
+
+        def invoke_response(callback,error,response,packet_io = nil)
+            if callback
+                
+                result = if error
+                             nil
+                       elsif response.respond_to?(:read) && packet_io
+                           response.read(packet_io)
+                       elsif response
+                           response
+                       else
+                           nil
+                       end
+
+                logger.debug { "Invoking response cb=#{callback} err=#{error} resp=#{result}"  }
+                binding.invoke(callback,error,result)
+            end
+        end
+
+        def resolve_watcher(watch_type,watcher)
+            if watcher == true && @watcher
+                #the actual TrueClass refers to the default watcher
+                watcher = @watcher
+            elsif watcher.respond_to?(:call) || watcher.respond_to?(:process_watch)
+                # ok a proc or quacks like a watcher
+            elsif watcher
+                # something, but not something we can handle
+                raise ArgumentError, "Not a watcher #{watcher.inspect}"
+            else
+                watch_type = nil
+            end
+            [watch_type,watcher]
+        end
+
+
+        def queue_packet(packet)
+            @pending_queue.push(packet)
+            logger.debug { "Queued: #{packet.inspect}" }
+
+            if @keeper_state == :connected
+                send_packet(packet)
+            end
+        end
+
+        def next_xid
+            @xid += 1
+        end
+
+        def send_packet(packet)
+            records = [] << Proto::RequestHeader.new(:xid => packet.xid, :_type => packet.opcode)
+            records << packet.request if packet.request
+            conn.send_records(*records)
+        end
+
+
+    end # Session
+end