ninjudd-memcache 0.9.0

data/lib/memcache/local_server.rb

@@ -0,0 +1,91 @@
+class Memcache
+  class LocalServer
+    def initialize
+      @data = {}
+      @expiry = {}
+    end
+
+    def name
+      "local:#{hash}"
+    end
+
+    def flush_all(delay = nil)
+      raise 'flush_all not supported with delay' if delay
+      @data.clear
+      @expiry.clear
+    end
+
+    def get(keys)
+      if keys.kind_of?(Array)
+        hash = {}
+        keys.each do |key|
+          key = key.to_s
+          val = get(key)
+          hash[key] = val if val
+        end
+        hash
+      else
+        key = keys.to_s
+        if @expiry[key] and Time.now > @expiry[key]
+          @data[key]   = nil
+          @expiry[key] = nil
+        end
+        @data[key]
+      end
+    end
+
+    def incr(key, amount = 1)
+      key = key.to_s
+      value = get(key)
+      return unless value
+      return unless value =~ /^\d+$/
+
+      value = value.to_i + amount
+      value = 0 if value < 0
+      @data[key] = value.to_s
+      value
+    end
+
+    def decr(key, amount = 1)
+      incr(key, -amount)
+    end
+
+    def delete(key, expiry = nil)
+      if expiry
+        old_expiry = @expiry[key.to_s] || Time.now + expiry
+        @expiry[key.to_s] = [old_expiry, expiry].min
+      else
+        @data.delete(key.to_s)
+      end
+    end
+
+    def set(key, value, expiry = nil)
+      key = key.to_s
+      @data[key]   = value
+      @expiry[key] = Time.now + expiry if expiry and expiry != 0
+      value
+    end
+
+    def add(key, value, expiry = nil)
+      return nil if get(key)
+      set(key, value, expiry)
+    end
+
+    def replace(key, value, expiry = nil)
+      return nil if get(key).nil?
+      set(key, value, expiry)
+    end
+
+    def append(key, value)
+      existing = get(key)
+      return nil if existing.nil?
+      set(key, existing + value)
+    end
+
+    def prepend(key, value)
+      existing = get(key)
+      return nil if existing.nil?
+      set(key, value + existing)
+    end
+  end
+end

data/lib/memcache/migration.rb

@@ -0,0 +1,23 @@
+class Memcache
+  class Migration < ActiveRecord::Migration
+    class << self
+      attr_accessor :table
+    end
+    
+    def self.up
+      create_table table, :id => false do |t|
+        t.string    :key
+        t.text      :value
+        t.timestamp :expires_at
+        t.timestamp :updated_at
+      end
+      
+      add_index table, [:key], :unique => true
+      add_index table, [:expires_at]
+    end
+    
+    def self.down
+      drop_table table
+    end
+  end
+end

data/lib/memcache/null_server.rb

@@ -0,0 +1,30 @@
+class Memcache
+  class NullServer
+    def name
+      "null"
+    end
+
+    def flush_all(delay = nil)
+    end
+
+    def get(keys)
+      keys.kind_of?(Array) ? {} : nil
+    end
+
+    def incr(key, amount = nil)
+      nil
+    end
+
+    def delete(key, expiry = nil)
+      nil
+    end
+
+    def set(key, value, expiry = nil)
+      nil
+    end
+
+    def add(key, value, expiry = nil)
+      nil
+    end
+  end
+end

data/lib/memcache/pg_server.rb

@@ -0,0 +1,159 @@
+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)
+      expiry == 0 ? 'NULL' : "NOW() + interval '#{expiry} seconds'"
+    end
+  end
+end

data/lib/memcache/segmented_server.rb

@@ -0,0 +1,89 @@
+require 'digest/sha1'
+
+class Memcache
+  class SegmentedServer < Server
+    MAX_SIZE = 1000000 # bytes
+    PARTIAL_VALUE = 0x40000000
+    
+    def get(keys, opts = {})
+      return get([keys], opts)[keys.to_s] unless keys.kind_of?(Array)
+      return {} if keys.empty?
+
+      results = super
+      keys = {}
+      keys_to_fetch = []
+      results.each do |key, value|
+        next unless segmented?(value)
+        hash, num = value.split(':')
+        keys[key] = []
+        num.to_i.times do |i|
+          hash_key = "#{hash}:#{i}"
+          keys_to_fetch << hash_key
+          keys[key]     << hash_key 
+        end
+      end
+      
+      parts = super(keys_to_fetch)
+      keys.each do |key, hashes|
+        value = ''
+        hashes.each do |hash_key|
+          value << parts[hash_key]
+        end
+        value.memcache_cas   = results[key].memcache_cas
+        value.memcache_flags = results[key].memcache_flags ^ PARTIAL_VALUE
+        results[key] = value
+      end
+      results
+    end
+
+    def set(key, value, expiry = 0, flags = 0)
+      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)
+        end
+        [master_key, flags | PARTIAL_VALUE]
+      else
+        [value, flags]
+      end
+    end
+  end
+end

data/lib/memcache/server.rb

@@ -0,0 +1,249 @@
+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
+    attr_writer :strict_reads
+
+    class MemcacheError < StandardError; end
+    class ServerError   < MemcacheError; end
+    class ClientError   < MemcacheError; end
+
+    def initialize(opts)
+      @host         = opts[:host]
+      @port         = opts[:port] || DEFAULT_PORT
+      @readonly     = opts[:readonly]
+      @strict_reads = opts[:strict_reads]
+      @status       = 'NOT CONNECTED'
+    end
+
+    def inspect
+      "<Memcache::Server: %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 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 = false)
+      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)
+        value.memcache_flags = flags.to_i
+        value.memcache_cas   = cas
+        results[key] = value
+      end
+      results
+    end
+
+    def incr(key, amount = 1)
+      check_writable!
+      raise MemcacheError, "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 MemcacheError, "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"
+    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
+
+  private
+
+    def check_writable!
+      raise MemcacheError, "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") if command.kind_of?(Array)
+#puts command
+#puts '==========================='
+      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
+    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
+          # Close the socket and retry once.
+          retried = true
+          close
+          retry
+        end
+        close(e) # Mark dead.
+        raise(e)
+      end
+    end
+
+    def read_command(command, &block)
+      raise MemcacheError, "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)
+          match_response!(socket.read(2), "\r\n")
+          response = socket.gets
+        end
+        unexpected_eof!
+      end
+    rescue Exception => e
+      puts "Memcache read error: #{e.class}: #{e.to_s}"
+      close(e) # Mark dead.
+      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 MemcacheError, 'unexpected end of file' 
+    end
+  end
+end