moneta 1.1.0 → 1.4.0

data/lib/moneta/proxy.rb

@@ -22,7 +22,7 @@ module Moneta
       raise NotImplementedError, "each_key is not supported on this proxy" \
         unless supports? :each_key
 
-      return enum_for(:each_key) unless block_given?
+      return enum_for(:each_key) { adapter.each_key.size } unless block_given?
       adapter.each_key(&block)
       self
     end
@@ -128,8 +128,8 @@ module Moneta
       end
 
       # (see Defaults::ClassMethods#not_supports)
-      def not_supports *features
-        @features_mask = (self.features_mask | features).freeze
+      def not_supports(*features)
+        @features_mask = (features_mask | features).freeze
         super
       end
     end

data/lib/moneta/server.rb

@@ -4,6 +4,154 @@ module Moneta
   # Moneta server to be used together with Moneta::Adapters::Client
   # @api public
   class Server
+    TIMEOUT = 1
+    MAXSIZE = 0x100000
+
+    # @api private
+    class Connection
+      def initialize(io, store)
+        @io = io
+        @store = store
+        @fiber = Fiber.new { run }
+      end
+
+      def resume(result = nil)
+        @fiber.resume result
+      end
+
+      private
+
+      # The return value of this function will be sent to the reactor.
+      #
+      # @return [:closed,Exception]
+      def run
+        catch :closed do
+          loop { write_dispatch(read_msg) }
+        end
+        :closed
+      rescue => ex
+        ex
+      ensure
+        @io.close unless @io.closed?
+      end
+
+      def dispatch(method, args)
+        case method
+        when :key?, :load, :delete, :increment, :create, :features
+          @store.public_send(method, *args)
+        when :store, :clear
+          @store.public_send(method, *args)
+          nil
+        when :each_key
+          yield_each(@store.each_key)
+          nil
+        end
+      rescue => ex
+        ex
+      end
+
+      def write_dispatch(msg)
+        method, *args = msg
+        result = dispatch(method, args)
+        write(result)
+      end
+
+      def read_msg
+        size = read(4).unpack('N').first
+        throw :closed, 'Message too big' if size > MAXSIZE
+        Marshal.load(read(size))
+      end
+
+      def read(len)
+        buffer = ''
+        loop do
+          begin
+            case received = @io.recv_nonblock(len)
+            when '', nil
+              throw :closed, 'Closed during read'
+            else
+              buffer << received
+              len -= received.bytesize
+            end
+          rescue IO::WaitReadable, IO::WaitWritable
+            yield_to_reactor(:read)
+          rescue Errno::ECONNRESET
+            throw :closed, 'Closed during read'
+          rescue IOError => ex
+            if ex.message =~ /closed stream/
+              throw :closed, 'Closed during read'
+            else
+              raise
+            end
+          end
+          break if len == 0
+        end
+        buffer
+      end
+
+      def write(obj)
+        buffer = pack(obj)
+        until buffer.empty?
+          begin
+            len = sendmsg(buffer)
+            buffer = buffer.byteslice(len...buffer.length)
+          rescue IO::WaitWritable, Errno::EINTR
+            yield_to_reactor(:write)
+          end
+        end
+        nil
+      end
+
+      # Detect support for socket#sendmsg_nonblock
+      Socket.new(Socket::AF_INET, Socket::SOCK_STREAM).tap do |socket|
+        begin
+          socket.sendmsg_nonblock('probe')
+        rescue Errno::EPIPE, Errno::ENOTCONN
+          def sendmsg(msg)
+            @io.sendmsg_nonblock(msg)
+          end
+        rescue NotImplementedError
+          def sendmsg(msg)
+            @io.write_nonblock(msg)
+          end
+        end
+      end
+
+      def yield_to_reactor(mode = :read)
+        if Fiber.yield(mode) == :close
+          throw :closed, 'Closed by reactor'
+        end
+      end
+
+      def pack(obj)
+        s = Marshal.dump(obj)
+        [s.bytesize].pack('N') << s
+      end
+
+      def yield_each(enumerator)
+        received_break = false
+        loop do
+          case msg = read_msg
+          when %w{NEXT}
+            # This will raise a StopIteration at the end of the enumeration,
+            # which will exit the loop.
+            write(enumerator.next)
+          when %w{BREAK}
+            # This is received when the client wants to stop the enumeration.
+            received_break = true
+            break
+          else
+            # Otherwise, the client is attempting to call another method within
+            # an `each` block.
+            write_dispatch(msg)
+          end
+        end
+      ensure
+        # This tells the client to stop enumerating
+        write(StopIteration.new("Server initiated stop")) unless received_break
+      end
+    end
+
     # @param [Hash] options
     # @option options [Integer] :port (9000) TCP port
     # @option options [String] :socket Alternative Unix socket file name
@@ -11,7 +159,9 @@ module Moneta
       @store = store
       @server = start(options)
       @ios = [@server]
-      @clients = {}
+      @reads = @ios.dup
+      @writes = []
+      @connections = {}
       @running = false
     end
 
@@ -26,97 +176,89 @@ module Moneta
     #
     # @note This method blocks!
     def run
-      raise 'Already running' if @running
+      raise 'Already running' if running?
       @stop = false
       @running = true
       begin
-        until @stop
-          mainloop
-        end
+        mainloop until @stop
       ensure
-        File.unlink(@socket) if @socket
-        @ios.each{ |io| io.close rescue nil }
+        @running = false
+        @server.close unless @server.closed?
+        @ios
+          .reject { |io| io == @server }
+          .each { |io| close_connection(io) }
+        File.unlink(@socket) if @socket rescue nil
       end
     end
 
     # Stop the server
     def stop
-      raise 'Not running' unless @running
+      raise 'Not running' unless running?
       @stop = true
       @server.close
-      @server = nil
+      nil
     end
 
     private
 
-    TIMEOUT = 1
-    MAXSIZE = 0x100000
-
     def mainloop
-      if ios = IO.select(@ios, nil, @ios, TIMEOUT)
-        ios[2].each do |io|
-          io.close
-          delete_client(io)
-        end
-        ios[0].each do |io|
-          if io == @server
-            if client = @server.accept
-              @ios << client
-              @clients[client] = ''
-            end
-          elsif io.closed? || io.eof?
-            delete_client(io)
-          else
-            handle(io, @clients[io] << io.readpartial(0xFFFF))
-          end
+      if ready = IO.select(@reads, @writes, @ios, TIMEOUT)
+        reads, writes, errors = ready
+        errors.each { |io| close_connection(io) }
+
+        @reads -= reads
+        reads.each do |io|
+          io == @server ? accept_connection : resume(io)
         end
+
+        @writes -= writes
+        writes.each { |io| resume(io) }
       end
-    rescue SignalException => ex
-      warn "Moneta::Server - #{ex.message}"
-      raise if ex.signo == 15 || ex.signo == 2 # SIGTERM or SIGINT
-    rescue Exception => ex
+    rescue SignalException => signal
+      warn "Moneta::Server - received #{signal}"
+      case signal.signo
+      when Signal.list['INT'], Signal.list['TERM']
+        @stop = true # graceful exit
+      end
+    rescue IOError => ex
+      # We get a lot of these "closed stream" errors, which we ignore
+      raise unless ex.message =~ /closed stream/
+    rescue Errno::EBADF => ex
       warn "Moneta::Server - #{ex.message}"
     end
 
-    def delete_client(io)
+    def accept_connection
+      io = @server.accept
+      @connections[io] = Connection.new(io, @store)
+      @ios << io
+      resume(io)
+    ensure
+      @reads << @server
+    end
+
+    def delete_connection(io)
       @ios.delete(io)
-      @clients.delete(io)
+      @reads.delete(io)
+      @writes.delete(io)
     end
 
-    def pack(o)
-      s = Marshal.dump(o)
-      [s.bytesize].pack('N') << s
+    def close_connection(io)
+      delete_connection(io)
+      @connections.delete(io).resume(:close)
     end
 
-    def handle(io, buffer)
-      buffer = @clients[io]
-      return if buffer.bytesize < 8 # At least 4 bytes for the marshalled array
-      size = buffer[0,4].unpack('N').first
-      if size > MAXSIZE
-        delete_client(io)
-        return
-      end
-      return if buffer.bytesize < 4 + size
-      buffer.slice!(0, 4)
-      method, *args = Marshal.load(buffer.slice!(0, size))
-      case method
-      when :key?, :load, :delete, :increment, :create
-        io.write(pack @store.send(method, *args))
-      when :features
-        # all features except each_key are supported
-        io.write(pack(@store.features - [:each_key]))
-      when :store, :clear
-        @store.send(method, *args)
-        io.write(@nil ||= pack(nil))
-      else
-        raise 'Invalid method call'
+    def resume(io)
+      case result = @connections[io].resume
+      when :closed # graceful exit
+        delete_connection(io)
+      when Exception # messy exit
+        delete_connection(io)
+        raise result
+      when :read
+        @reads << io
+      when :write
+        @writes << io
       end
-    rescue IOError => ex
-      warn "Moneta::Server - #{ex.message}" unless ex.message =~ /closed/
-      delete_client(io)
-    rescue Exception => ex
-      warn "Moneta::Server - #{ex.message}"
-      io.write(pack Exception.new(ex.message))
     end
 
     def start(options)
@@ -136,5 +278,14 @@ module Moneta
         TCPServer.open(options[:host] || '127.0.0.1', options[:port] || 9000)
       end
     end
+
+    def stats
+      {
+        connections: @connections.length,
+        reading: @reads.length,
+        writing: @writes.length,
+        total: @ios.length
+      }
+    end
   end
 end

data/lib/moneta/shared.rb

@@ -11,8 +11,6 @@ module Moneta
   #
   # @api public
   class Shared < Wrapper
-    not_supports :each_key
-
     # @param [Hash] options
     # @option options [Integer] :port (9000) TCP port
     # @option options [String] :host Server hostname
@@ -25,7 +23,7 @@ module Moneta
 
     # (see Proxy#close)
     def close
-      if @server
+      if server?
         @server.stop
         @thread.join
         @server = @thread = nil
@@ -36,13 +34,20 @@ module Moneta
       end
     end
 
+    # Returns true if this wrapper is running as the server
+    #
+    # @return [Boolean] wrapper is a server
+    def server?
+      @server != nil
+    end
+
     protected
 
     def wrap(*args)
       connect
       yield
     rescue Errno::ECONNRESET, Errno::EPIPE, IOError, SystemCallError
-      @connect_lock.synchronize{ close unless @server }
+      @connect_lock.synchronize { close unless server? }
       tries ||= 0
       (tries += 1) < 3 ? retry : raise
     end
@@ -52,21 +57,22 @@ module Moneta
       @connect_lock.synchronize do
         @adapter ||= Adapters::Client.new(@options)
       end
-    rescue Errno::ECONNREFUSED, Errno::ENOENT => ex
+    rescue Errno::ECONNREFUSED, Errno::ENOENT, IOError => ex
       start_server
       tries ||= 0
       warn "Moneta::Shared - Failed to connect: #{ex.message}" if tries > 0
-      (tries += 1) < 3 ? retry : raise
+      (tries += 1) < 10 ? retry : raise
     end
 
     # TODO: Implement this using forking (MRI) and threading (JRuby)
     # to get maximal performance
     def start_server
       @connect_lock.synchronize do
+        return if server?
         begin
           raise "Adapter already set" if @adapter
           @adapter = Lock.new(@builder.build.last)
-          raise "Server already set" if @server
+          raise "Server already set" if server?
           @server = Server.new(@adapter, @options)
           @thread = Thread.new { @server.run }
           sleep 0.1 until @server.running?

data/lib/moneta/stack.rb

@@ -42,7 +42,7 @@ module Moneta
 
     # (see Proxy#key?)
     def key?(key, options = {})
-      @stack.any? {|s| s.key?(key, options) }
+      @stack.any? { |s| s.key?(key, options) }
     end
 
     # (see Proxy#load)
@@ -56,21 +56,21 @@ module Moneta
 
     # (see Proxy#store)
     def store(key, value, options = {})
-      @stack.each {|s| s.store(key, value, options) }
+      @stack.each { |s| s.store(key, value, options) }
       value
     end
 
     # (see Proxy#increment)
     def increment(key, amount = 1, options = {})
       last = nil
-      @stack.each {|s| last = s.increment(key, amount, options) }
+      @stack.each { |s| last = s.increment(key, amount, options) }
       last
     end
 
     # (see Proxy#create)
     def create(key, value, options = {})
       last = false
-      @stack.each {|s| last = s.create(key, value, options) }
+      @stack.each { |s| last = s.create(key, value, options) }
       last
     end
 
@@ -84,13 +84,13 @@ module Moneta
 
     # (see Proxy#clear)
     def clear(options = {})
-      @stack.each {|s| s.clear(options) }
+      @stack.each { |s| s.clear(options) }
       self
     end
 
     # (see Proxy#close)
     def close
-      @stack.each {|s| s.close }
+      @stack.each { |s| s.close }
       nil
     end