segmented-memcache 1.2.0

data/lib/memcache/pg_server.rb

@@ -0,0 +1,163 @@
+require 'active_record'
+require 'memcache/migration'
+
+class Memcache
+  class PGServer
+    attr_reader :db, :table
+
+    def initialize(opts)
+      @table = opts[:table]
+      @db    = opts[:db] || ActiveRecord::Base.connection.raw_connection
+    end
+
+    def name
+      @name ||= begin
+        db_config = db.instance_variable_get(:@config)
+        "#{db_config[:host]}:#{db_config[:database]}:#{table}"
+      end
+    end
+
+    def flush_all(delay = nil)
+      db.exec("TRUNCATE #{table}")
+    end
+
+    def get(keys)
+      return get([keys])[keys.to_s] unless keys.kind_of?(Array)
+
+      keys = keys.collect {|key| quote(key.to_s)}.join(',')
+      sql = %{
+        SELECT key, value FROM #{table}
+          WHERE key IN (#{keys}) AND #{expiry_clause}
+      }
+      results = {}
+      db.query(sql).each do |key, value|
+        results[key] = value
+      end
+      results
+    end
+
+    def incr(key, amount = 1)
+      transaction do
+        value = get(key)
+        return unless value        
+        return unless value =~ /^\d+$/
+        
+        value = value.to_i + amount
+        value = 0 if value < 0
+        db.exec %{
+          UPDATE #{table} SET value = #{quote(value)}, updated_at = NOW()
+            WHERE key = #{quote(key)}
+        }
+        value
+      end
+    end
+
+    def decr(key, amount = 1)
+      incr(key, -amount)
+    end
+
+    def delete(key)
+      result = db.exec %{
+        DELETE FROM #{table}
+          WHERE key = #{quote(key)}
+      }
+    end
+
+    def set(key, value, expiry = 0)
+      transaction do
+        delete(key)
+        insert(key, value, expiry)
+      end
+      value
+    end
+
+    def add(key, value, expiry = 0)
+      delete_expired(key)
+      insert(key, value, expiry)
+      value
+    rescue PGError => e
+      nil
+    end
+
+    def replace(key, value, expiry = 0)      
+      delete_expired(key)
+      result = update(key, value, expiry)
+      result.cmdtuples == 1 ? value : nil
+    end
+
+    def append(key, value)      
+      delete_expired(key)
+      result = db.exec %{
+        UPDATE #{table}
+          SET value = value || #{quote(value)}, updated_at = NOW()
+          WHERE key = #{quote(key)}
+      }
+      result.cmdtuples == 1
+    end
+
+    def prepend(key, value)      
+      delete_expired(key)
+      result = db.exec %{
+        UPDATE #{table}
+          SET value = #{quote(value)} || value, updated_at = NOW()
+          WHERE key = #{quote(key)}
+      }
+      result.cmdtuples == 1
+    end
+
+  private
+ 
+    def insert(key, value, expiry = 0)
+      db.exec %{
+        INSERT INTO #{table} (key, value, updated_at, expires_at)
+          VALUES (#{quote(key)}, #{quote(value)}, NOW(), #{expiry_sql(expiry)})
+      }
+    end
+
+    def update(key, value, expiry = 0)
+      db.exec %{
+        UPDATE #{table}
+          SET value = #{quote(value)}, updated_at = NOW(), expires_at = #{expiry_sql(expiry)}
+          WHERE key = #{quote(key)}
+      }
+    end
+
+    def transaction
+      return yield if @in_transaction
+      
+      begin 
+        @in_transaction = true
+        db.exec('BEGIN')
+        value = yield
+        db.exec('COMMIT')
+        value
+      rescue Exception => e
+        db.exec('ROLLBACK')
+        raise e
+      ensure
+        @in_transaction = false
+      end
+    end
+      
+    def quote(string)
+      string.to_s.gsub(/'/,"\'")
+      "'#{string}'"
+    end
+
+    def delete_expired(key)
+      db.exec "DELETE FROM #{table} WHERE key = #{quote(key)} AND NOT (#{expiry_clause})"
+    end
+    
+    def expiry_clause
+      "expires_at IS NULL OR expires_at > NOW()"
+    end
+
+    def expiry_sql(expiry)
+      if expiry.kind_of?(Time)
+        quote(expiry.to_s(:db))
+      else
+        expiry == 0 ? 'NULL' : "NOW() + interval '#{expiry} seconds'"
+      end
+    end
+  end
+end

data/lib/memcache/segmented_server.rb

@@ -0,0 +1,116 @@
+require 'digest/sha1'
+
+class Memcache
+  class SegmentedServer < Server
+    MAX_SIZE = 1000000 # bytes
+    PARTIAL_VALUE = 0x40000000
+    
+    alias :super_get :get
+    
+    def get(keys, cas = nil)
+      return get([keys], cas)[keys.to_s] unless keys.kind_of?(Array)
+      return {} if keys.empty?
+
+      results = super
+      keys = {}
+      keys_to_fetch = []
+
+      extract_keys(results, keys_to_fetch, keys)
+      
+      parts = super(keys_to_fetch)
+      keys.each do |key, hashes|
+        value = ''
+        hashes.each do |hash_key|          
+          if part = parts[hash_key]
+            value << part
+          else
+            value = nil
+            break
+          end
+        end
+
+        if value
+          value.memcache_cas   = results[key].memcache_cas
+          value.memcache_flags = results[key].memcache_flags ^ PARTIAL_VALUE
+        end
+        results[key] = value
+      end
+      results
+    end
+    
+    def delete(keys, cas = nil)
+      return delete([keys], cas) unless keys.kind_of?(Array)
+      return {} if keys.empty?
+
+      results = super_get(keys, cas)
+      keys_to_fetch = keys     
+      extract_keys(results, keys_to_fetch)
+      keys_to_fetch.each{|k| super k}
+    end
+
+    def set(key, value, expiry = 0, flags = 0)
+      delete key if !(super_get(key)).nil?
+      value, flags = store_segments(key, value, expiry, flags)
+      super(key, value, expiry, flags) && value
+    end
+
+    def cas(key, value, cas, expiry = 0, flags = 0)
+      value, flags = store_segments(key, value, expiry, flags)
+      super(key, value, cas, expiry, flags)
+    end
+
+    def add(key, value, expiry = 0, flags = 0)
+      value, flags = store_segments(key, value, expiry, flags)
+      super(key, value, expiry, flags)
+    end
+
+    def replace(key, value, expiry = 0, flags = 0)
+      value, flags = store_segments(key, value, expiry, flags)
+      super(key, value, expiry, flags)
+    end
+
+  private
+
+    def segmented?(value)
+      value.memcache_flags & PARTIAL_VALUE == PARTIAL_VALUE
+    end
+
+    def segment(key, value)
+      hash  = Digest::SHA1.hexdigest("#{key}:#{Time.now}:#{rand}")
+      parts = {}
+      i = 0; offset = 0
+      while offset < value.size
+        parts["#{hash}:#{i}"] = value[offset, MAX_SIZE]
+        offset += MAX_SIZE; i += 1
+      end
+      master_key = "#{hash}:#{parts.size}"
+      [master_key, parts]
+    end
+    
+    def store_segments(key, value, expiry = 0, flags = 0)
+      if value and value.size > MAX_SIZE
+        master_key, parts = segment(key, value)
+        parts.each do |hash, data|
+          set(hash, data, expiry==0 ? 0 : expiry + 1) # We want the segments to expire slightly after the master key.
+        end
+        [master_key, flags | PARTIAL_VALUE]
+      else
+        [value, flags]
+      end
+    end
+    
+    def extract_keys(results, keys_to_fetch, keys=nil)
+      results.each do |key, value|
+        next unless segmented?(value)
+        hash, num = value.split(':')
+        keys[key] = [] unless keys.nil?
+        num.to_i.times do |i|
+          hash_key = "#{hash}:#{i}"
+          keys_to_fetch << hash_key
+          keys[key]     << hash_key unless keys.nil?
+        end
+      end
+    end
+    
+  end
+end

data/lib/memcache/server.rb

@@ -0,0 +1,265 @@
+require 'socket'
+require 'thread'
+require 'timeout'
+
+class Memcache
+  class Server
+    CONNECT_TIMEOUT  = 1.0
+    READ_RETRY_DELAY = 5.0
+    DEFAULT_PORT     = 11211
+
+    attr_reader :host, :port, :status, :retry_at
+
+    def initialize(opts)
+      @host         = opts[:host]
+      @port         = opts[:port] || DEFAULT_PORT
+      @readonly     = opts[:readonly]
+      @strict_reads = opts[:strict_reads]
+      @status       = 'NOT CONNECTED'
+    end
+
+    def clone
+      self.class.new(:host => host, :port => port, :readonly => readonly?, :strict_reads => strict_reads?)
+    end
+
+    def inspect
+      "<#{self.class.name}: %s:%d (%s)>" % [@host, @port, @status]
+    end
+
+    def name
+      "#{host}:#{port}"
+    end
+
+    def alive?
+      @retry_at.nil? or @retry_at < Time.now
+    end
+
+    def readonly?
+      @readonly
+    end
+
+    def strict_reads?
+      @strict_reads
+    end
+
+    def close(error = nil)
+      # Close the socket. If there is an error, mark the server dead.
+      @socket.close if @socket and not @socket.closed?
+      @socket = nil
+      
+      if error
+        @retry_at = Time.now + READ_RETRY_DELAY
+        @status   = "DEAD: %s: %s, will retry at %s" % [error.class, error.message, @retry_at]
+      else
+        @retry_at = nil
+        @status   = "NOT CONNECTED"
+      end
+    end
+
+    def stats
+      stats = {}
+      read_command('stats') do |response|
+        key, value = match_response!(response, /^STAT ([\w]+) (-?[\w\.\:]+)/)
+
+        if ['rusage_user', 'rusage_system'].include?(key)
+          seconds, microseconds = value.split(/:/, 2)
+          microseconds ||= 0
+          stats[key] = Float(seconds) + (Float(microseconds) / 1_000_000)
+        else
+          stats[key] = (value =~ /^-?\d+$/ ? value.to_i : value)
+        end
+      end
+      stats
+    end
+    
+    def count
+      stats['curr_items']
+    end
+
+    def flush_all(delay = nil)
+      check_writable!
+      write_command("flush_all #{delay}")
+    end
+
+    alias clear flush_all
+
+    def gets(keys)
+      get(keys, true)
+    end
+
+    def get(keys, cas = nil)
+      return get([keys], cas)[keys.to_s] unless keys.kind_of?(Array)
+      return {} if keys.empty?
+
+      method = cas ? 'gets' : 'get'
+
+      results = {}
+      read_command("#{method} #{keys.join(' ')}") do |response|
+        if cas
+          key, flags, length, cas = match_response!(response, /^VALUE ([^\s]+) ([^\s]+) ([^\s]+) ([^\s]+)/)
+        else
+          key, flags, length = match_response!(response, /^VALUE ([^\s]+) ([^\s]+) ([^\s]+)/)
+        end
+          
+        value = socket.read(length.to_i)
+        match_response!(socket.read(2), "\r\n")
+
+        value.memcache_flags = flags.to_i
+        value.memcache_cas   = cas
+        results[key] = value
+      end
+      results
+    end
+
+    def incr(key, amount = 1)
+      check_writable!
+      raise Error, "incr requires unsigned value" if amount < 0
+      response = write_command("incr #{key} #{amount}")
+      response == "NOT_FOUND\r\n" ? nil : response.to_i
+    end
+
+    def decr(key, amount = 1)
+      check_writable!
+      raise Error, "decr requires unsigned value" if amount < 0
+      response = write_command("decr #{key} #{amount}")
+      response == "NOT_FOUND\r\n" ? nil : response.to_i
+    end
+
+    def delete(key)
+      check_writable!
+      write_command("delete #{key}") == "DELETED\r\n" ? true : nil
+    end
+
+    def set(key, value, expiry = 0, flags = 0)
+      return delete(key) if value.nil?
+      check_writable!
+      write_command("set #{key} #{flags.to_i} #{expiry.to_i} #{value.to_s.size}", value)
+      value
+    end
+
+    def cas(key, value, cas, expiry = 0, flags = 0)
+      check_writable!
+      response = write_command("cas #{key} #{flags.to_i} #{expiry.to_i} #{value.to_s.size} #{cas.to_i}", value)
+      response == "STORED\r\n" ? value : nil
+    end
+
+    def add(key, value, expiry = 0, flags = 0)
+      check_writable!
+      response = write_command("add #{key} #{flags.to_i} #{expiry.to_i} #{value.to_s.size}", value)
+      response == "STORED\r\n" ? value : nil
+    end
+
+    def replace(key, value, expiry = 0, flags = 0)
+      check_writable!
+      response = write_command("replace #{key} #{flags.to_i} #{expiry.to_i} #{value.to_s.size}", value)
+      response == "STORED\r\n" ? value : nil
+    end
+
+    def append(key, value)
+      check_writable!
+      response = write_command("append #{key} 0 0 #{value.to_s.size}", value)
+      response == "STORED\r\n"
+    end
+
+    def prepend(key, value)
+      check_writable!
+      response = write_command("prepend #{key} 0 0 #{value.to_s.size}", value)
+      response == "STORED\r\n"
+    end
+
+    class Error < StandardError; end
+    class ConnectionError < Error
+      def initialize(e)
+        if e.kind_of?(String)
+          super
+        else
+          super("(#{e.class}) #{e.message}")
+          set_backtrace(e.backtrace)
+        end
+      end
+    end
+    class ServerError < Error; end
+    class ClientError < Error; end
+
+  private
+
+    def check_writable!
+      raise Error, "Update of readonly cache" if readonly?
+    end
+
+    def match_response!(response, regexp)
+      # Make sure that the response matches the protocol.
+      unexpected_eof! if response.nil?
+      match = response.match(regexp)
+      raise ServerError, "unexpected response: #{response.inspect}" unless match
+
+      match.to_a[1, match.size]
+    end
+
+    def send_command(*command)
+      command = command.join("\r\n")
+      socket.write("#{command}\r\n")
+      response = socket.gets
+      
+      unexpected_eof! if response.nil?
+      if response =~ /^(ERROR|CLIENT_ERROR|SERVER_ERROR) (.*)\r\n/
+        raise ($1 == 'SERVER_ERROR' ? ServerError : ClientError), $2
+      end
+      
+      block_given? ? yield(response) : response
+    rescue Exception => e
+      close(e) # Mark dead.
+      raise e if e.kind_of?(Error)
+      raise ConnectionError.new(e)
+    end
+
+    def write_command(*command, &block)
+      retried = false
+      begin
+        send_command(*command, &block)
+      rescue Exception => e
+        puts "Memcache write error: #{e.class} #{e.to_s}"
+        unless retried
+          retried = true
+          retry
+        end
+        raise(e)
+      end
+    end
+
+    def read_command(command, &block)
+      raise ConnectionError, "Server dead, will retry at #{retry_at}" unless alive?
+      send_command(command) do |response|
+        while response do
+          return if response == "END\r\n"
+          yield(response)
+          response = socket.gets
+        end
+        unexpected_eof!
+      end
+    rescue Exception => e
+      puts "Memcache read error: #{e.class} #{e.to_s}"
+      raise(e) if strict_reads?
+    end
+
+    def socket
+      if @socket.nil? or @socket.closed?
+        # Attempt to connect.
+        @socket = timeout(CONNECT_TIMEOUT) do
+          TCPSocket.new(host, port)
+        end
+        if Socket.constants.include? 'TCP_NODELAY'
+          @socket.setsockopt Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1
+        end
+        
+        @retry_at = nil
+        @status   = 'CONNECTED'
+      end
+      @socket
+    end
+
+    def unexpected_eof!
+      raise Error, 'unexpected end of file'
+    end
+  end
+end