memcache 1.0.0

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
+      "<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 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 = 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)
+        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

data/lib/memcache.rb

@@ -0,0 +1,382 @@
+require 'zlib'
+
+$:.unshift(File.dirname(__FILE__))
+require 'memcache/server'
+require 'memcache/local_server'
+require 'memcache/segmented_server'
+
+class Memcache
+  DEFAULT_EXPIRY  = 0
+  LOCK_TIMEOUT    = 5
+  WRITE_LOCK_WAIT = 1
+
+  attr_reader :default_expiry, :default_namespace, :servers
+
+  def initialize(opts)
+    @default_expiry    = opts[:default_expiry] || DEFAULT_EXPIRY
+    @default_namespace = opts[:namespace]
+    default_server = opts[:segment_large_values] ? SegmentedServer : Server
+
+    @servers = (opts[:servers] || [ opts[:server] ]).collect do |server|
+      case server
+      when Hash
+        server = default_server.new(opts.merge(server))
+      when String
+        host, port = server.split(':')
+        server = default_server.new(opts.merge(:host => host, :port => port))
+      when Class
+        server = server.new
+      when :local
+        server = Memcache::LocalServer.new
+      end
+      server
+    end
+  end
+
+  def clone
+    self.class.new(
+      :default_expiry    => default_expiry,
+      :default_namespace => default_namespace,
+      :servers           => servers.collect {|s| s.clone}
+    )
+  end
+
+  def inspect
+    "<Memcache: %d servers, ns: %p>" % [@servers.length, namespace]
+  end
+
+  def namespace
+    @namespace || default_namespace
+  end
+
+  def namespace=(namespace)
+    if default_namespace == namespace
+      @namespace = nil
+    else
+      @namespace = namespace
+    end
+  end
+
+  def in_namespace(namespace)
+    # Temporarily change the namespace for convenience.
+    begin
+      old_namespace = self.namespace
+      self.namespace = "#{old_namespace}#{namespace}"
+      yield
+    ensure
+      self.namespace = old_namespace
+    end
+  end
+
+  def get(keys, opts = {})
+    raise 'opts must be hash' unless opts.kind_of?(Hash)
+
+    if keys.kind_of?(Array)
+      multi_get(keys, opts)
+    else
+      key = cache_key(keys)
+
+      if opts[:expiry]
+        value = server(key).gets(key)
+        server(key).cas(key, value, value.memcache_cas, opts[:expiry]) if value
+      else
+        value = server(key).get(key, opts[:cas])
+      end
+      unmarshal(value, opts)
+    end
+  end
+
+  def read(keys, opts = {})
+    get(keys, opts.merge(:raw => true))
+  end
+
+  def set(key, value, opts = {})
+    raise 'opts must be hash' unless opts.kind_of?(Hash)
+
+    expiry = opts[:expiry] || default_expiry
+    flags  = opts[:flags]  || 0
+    key    = cache_key(key)
+    data   = marshal(value, opts)
+    server(key).set(key, data, expiry, flags)
+    value
+  end
+
+  def write(key, value, opts = {})
+    set(key, value, opts.merge(:raw => true))
+  end
+
+  def add(key, value, opts = {})
+    raise 'opts must be hash' unless opts.kind_of?(Hash)
+
+    expiry = opts[:expiry] || default_expiry
+    flags  = opts[:flags]  || 0
+    key    = cache_key(key)
+    data   = marshal(value, opts)
+    server(key).add(key, data, expiry, flags) && value
+  end
+
+  def replace(key, value, opts = {})
+    raise 'opts must be hash' unless opts.kind_of?(Hash)
+
+    expiry = opts[:expiry] || default_expiry
+    flags  = opts[:flags]  || 0
+    key    = cache_key(key)
+    data   = marshal(value, opts)
+    server(key).replace(key, data, expiry, flags) && value
+  end
+
+  def cas(key, value, opts = {})
+    raise 'opts must be hash' unless opts.kind_of?(Hash)
+
+    expiry = opts[:expiry] || default_expiry
+    flags  = opts[:flags]  || 0
+    key    = cache_key(key)
+    data   = marshal(value, opts)
+    server(key).cas(key, data, opts[:cas], expiry, flags) && value
+  end
+
+  def append(key, value)
+    key = cache_key(key)
+    server(key).append(key, value)
+  end
+
+  def prepend(key, value)
+    key = cache_key(key)
+    server(key).prepend(key, value)
+  end
+
+  def count(key)
+    key = cache_key(key)
+    server(key).get(key).to_i
+  end
+
+  def incr(key, amount = 1)
+    key = cache_key(key)
+    server(key).incr(key, amount)
+  end
+
+  def decr(key, amount = 1)
+    key = cache_key(key)
+    server(key).decr(key, amount)
+  end
+
+  def update(key, opts = {})
+    value = get(key, :cas => true)
+    if value
+      cas(key, yield(value), opts.merge!(:cas => value.memcache_cas))
+    else
+      add(key, yield(value), opts)
+    end
+  end
+
+  def get_or_add(key, *args)
+    # Pseudo-atomic get and update.
+    if block_given?
+      opts = args[0] || {}
+      get(key) || add(key, yield, opts) || get(key)
+    else
+      opts = args[1] || {}
+      get(key) || add(key, args[0], opts) || get(key)
+    end
+  end
+
+  def get_or_set(key, *args)
+    if block_given?
+      opts = args[0] || {}
+      get(key) || set(key, yield, opts)
+    else
+      opts = args[1] || {}
+      get(key) || set(key, args[0], opts)
+    end
+  end
+
+  def get_some(keys, opts = {})
+    keys = keys.collect {|key| key.to_s}
+
+    records = opts[:disable] ? {} : self.get(keys, opts)
+    if opts[:validation]
+      records.delete_if do |key, value|
+        not opts[:validation].call(key, value)
+      end
+    end
+
+    keys_to_fetch = keys - records.keys
+    method = opts[:overwrite] ? :set : :add
+    if keys_to_fetch.any?
+      yield(keys_to_fetch).each do |key, value|
+        self.send(method, key, value, opts) unless opts[:disable] or opts[:disable_write]
+        records[key] = value
+      end
+    end
+    records
+  end
+
+  def lock(key, opts = {})
+    # Returns false if the lock already exists.
+    expiry = opts[:expiry] || LOCK_TIMEOUT
+    add(lock_key(key), Socket.gethostname, :expiry => expiry, :raw => true)
+  end
+
+  def unlock(key)
+    delete(lock_key(key))
+  end
+
+  def with_lock(key, opts = {})
+    until lock(key) do
+      return if opts[:ignore]
+      sleep(WRITE_LOCK_WAIT) # just wait
+    end
+    yield
+    unlock(key) unless opts[:keep]
+  end
+
+  def lock_key(key)
+    "lock:#{key}"
+  end
+
+  def locked?(key)
+    get(lock_key(key), :raw => true)
+  end
+
+  def delete(key)
+    key = cache_key(key)
+    server(key).delete(key)
+  end
+
+  def flush_all(opts = {})
+    delay    = opts[:delay].to_i
+    interval = opts[:interval].to_i
+
+    servers.each do |server|
+      server.flush_all(delay)
+      delay += interval
+    end
+  end
+
+  def reset
+    servers.each {|server| server.close if server.respond_to?(:close)}
+  end
+
+  def stats(field = nil)
+    if field
+      servers.collect do |server|
+        server.stats[field]
+      end
+    else
+      stats = {}
+      servers.each do |server|
+        stats[server.name] = server.stats
+      end
+      stats
+    end
+  end
+ 
+  alias clear flush_all
+
+  def [](key)
+    get(key)
+  end
+
+  def []=(key, value)
+    set(key, value)
+  end
+
+protected
+
+  def multi_get(keys, opts = {})
+    return {} if keys.empty?
+    
+    key_to_input_key = {}
+    keys_by_server  = Hash.new { |h,k| h[k] = [] }
+    
+    # Store keys by servers. Also store a mapping from cache key to input key.
+    keys.each do |input_key|
+      key    = cache_key(input_key)
+      server = server(key)
+      key_to_input_key[key] = input_key.to_s
+      keys_by_server[server] << key
+    end
+    
+    # Fetch and combine the results. Also, map the cache keys back to the input keys.
+    results = {}
+    keys_by_server.each do |server, keys|
+      server.get(keys, opts[:cas]).each do |key, value|
+        input_key = key_to_input_key[key]
+        results[input_key] = unmarshal(value, opts)
+      end
+    end
+    results
+  end
+
+  def cache_key(key)
+    safe_key = key ? key.to_s.gsub(/%/, '%%').gsub(/ /, '%s') : key
+    if namespace.nil? then
+      safe_key
+    else
+      "#{namespace}:#{safe_key}"
+    end
+  end
+  
+  def marshal(value, opts = {})
+    opts[:raw] ? value : Marshal.dump(value)
+  end
+
+  def unmarshal(value, opts = {})
+    return value if value.nil? or opts[:raw]
+    
+    object = Marshal.load(value)
+    object.memcache_flags = value.memcache_flags
+    object.memcache_cas   = value.memcache_cas
+    object
+  rescue Exception => e
+    puts "Memcache read error: #{e.class} #{e.to_s} while unmarshalling value: #{value}"
+    nil
+  end
+
+  def server(key)
+    raise ArgumentError, "key too long #{key.inspect}" if key.length > 250
+    return servers.first if servers.length == 1
+
+    hash = (Zlib.crc32(key) >> 16) & 0x7fff
+    servers[hash % servers.length]
+  end
+
+  class Pool
+    attr_reader :fallback
+
+    def initialize
+      @cache_by_scope = {}
+      @cache_by_scope[:default] = Memcache.new(:server => Memcache::LocalServer)
+      @fallback = :default
+    end
+    
+    def include?(scope)
+      @cache_by_scope.include?(scope.to_sym)
+    end
+
+    def fallback=(scope)
+      @fallback = scope.to_sym
+    end
+
+    def [](scope)
+      @cache_by_scope[scope.to_sym] || @cache_by_scope[fallback]
+    end
+    
+    def []=(scope, cache)
+      @cache_by_scope[scope.to_sym] = cache
+    end
+
+    def reset
+      @cache_by_scope.values.each {|c| c.reset}
+    end
+  end
+  
+  def self.pool
+    @@cache_pool ||= Pool.new
+  end
+end
+
+# Add flags and cas
+class Object
+  attr_accessor :memcache_flags, :memcache_cas
+end

data/memcache.gemspec

@@ -0,0 +1,69 @@
+# Generated by jeweler
+# DO NOT EDIT THIS FILE DIRECTLY
+# Instead, edit Jeweler::Tasks in Rakefile, and run the gemspec command
+# -*- encoding: utf-8 -*-
+
+Gem::Specification.new do |s|
+  s.name = %q{memcache}
+  s.version = "1.0.0"
+
+  s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version=
+  s.authors = ["Justin Balthrop"]
+  s.date = %q{2009-11-20}
+  s.description = %q{Ruby client for memcached supporting advanced protocol features and pluggable architecture.}
+  s.email = %q{code@justinbalthrop.com}
+  s.extensions = ["ext/extconf.rb"]
+  s.extra_rdoc_files = [
+    "LICENSE",
+     "README.rdoc"
+  ]
+  s.files = [
+    ".gitignore",
+     "LICENSE",
+     "README.rdoc",
+     "Rakefile",
+     "VERSION.yml",
+     "lib/memcache.rb",
+     "lib/memcache/local_server.rb",
+     "lib/memcache/migration.rb",
+     "lib/memcache/null_server.rb",
+     "lib/memcache/pg_server.rb",
+     "lib/memcache/segmented_server.rb",
+     "lib/memcache/server.rb",
+     "memcache.gemspec",
+     "test/memcache_local_server_test.rb",
+     "test/memcache_null_server_test.rb",
+     "test/memcache_pg_server_test.rb",
+     "test/memcache_segmented_server_test.rb",
+     "test/memcache_server_test.rb",
+     "test/memcache_server_test_helper.rb",
+     "test/memcache_test.rb",
+     "test/test_helper.rb"
+  ]
+  s.homepage = %q{http://github.com/ninjudd/memcache}
+  s.rdoc_options = ["--charset=UTF-8"]
+  s.require_paths = ["lib"]
+  s.rubygems_version = %q{1.3.5}
+  s.summary = %q{Advanced ruby memcache client}
+  s.test_files = [
+    "test/memcache_local_server_test.rb",
+     "test/memcache_null_server_test.rb",
+     "test/memcache_pg_server_test.rb",
+     "test/memcache_segmented_server_test.rb",
+     "test/memcache_server_test.rb",
+     "test/memcache_server_test_helper.rb",
+     "test/memcache_test.rb",
+     "test/test_helper.rb"
+  ]
+
+  if s.respond_to? :specification_version then
+    current_version = Gem::Specification::CURRENT_SPECIFICATION_VERSION
+    s.specification_version = 3
+
+    if Gem::Version.new(Gem::RubyGemsVersion) >= Gem::Version.new('1.2.0') then
+    else
+    end
+  else
+  end
+end
+