require 'rubygems'
require 'memcache'
# It's your job to catch any signals relating to process termination
# (INT, TERM, etc.) and call the shutdown method on each worker. If
# this isn't done you run the risk of losing messages.
class MemcacheQueue
class ShutdownError < StandardError; end
attr_accessor :add_timeout
def initialize(name, *args)
@name = name or raise ArgumentError, "Must specify a unique name for each worker"
@client = MemCache.new *args
@first_key = state(:get, name)
@add_timeout = 10
end
def create_queue
@client.add('latest_added', '0', 0, true)
@client.add('latest_read', '0', 0, true)
end
# ! lost message potential !
# If it takes more than add_timeout seconds from the time incr
# returns to the time the message is add()ed to memcache the message _could_
# be lost.
def add_msg(msg)
raise ShutdownError if @shutdown
begin
latest_msg ||= @client.incr('latest_added')
@client.add(latest_msg.to_s, msg)
@client.set("added_#{latest_msg}", true)
rescue Exception
warn "Error on add: #{$!.inspect}, retrying"
sleep 1
retry
end
end
# ! lost message potential !
# If we die between incr and state(:set) the message _will_ be
# lost. This would have to be due to hard process kill or other like events.
def get_msg
raise ShutdownError if @shutdown
if key = @first_key
@first_key = nil
else
key = @client.incr('latest_read').to_s
end
start_time = Time.now
loop do
begin
if msg = @client.get(key)
begin
@client.delete(key)
rescue Exception
warn "Error while deleting key #{$!.inspect}"
end
return msg
end
rescue Exception
warn "Error on get, #{$!.inspect}, retrying"
end
state(:set, key) unless @state_set
raise ShutdownError if @shutdown
return nil if (Time.now - start_time) > @add_timeout and failed_add?(key)
sleep 1
start_time = Time.now if msgs_left(true) < 0
end
ensure
state(:delete, key) if @state_set and ! @shutdown
end
def msgs_left(neg_value = false)
latest_added = @client.get('latest_added', true).to_i
latest_read = @client.get('latest_read', true).to_i
# check for 64bit boundary crossing
if latest_added > 2**63 and latest_read < 2**62
latest_read += 2**64
elsif latest_read > 2**63 and latest_added < 2**62
latest_added += 2**64
end
diff = latest_added - latest_read
neg_value ? diff : [0, diff].max
end
def shutdown
@shutdown = true
end
private
def failed_add?(key)
msgs_left(true) >= 0 and ! @client.get("added_#{key}")
end
def state(action, key = nil)
return @client.get("state_#@name") if action == :get
@client.send(action, "state_#@name", key)
@state_set = (action == :set)
rescue Exception
warn "Error on #{action} state, #{$!.inspect}, retrying"
sleep 1
retry
end
end
