super_queue 0.2.1 → 0.3.1

data/lib/super_queue.rb

@@ -1,19 +1,32 @@
 require 'aws-sdk'
 require 'base64'
-require 'socket'
 require 'digest/md5'
-require 'zlib'
 
 class SuperQueue
 
+  class S3Pointer < Hash
+    def initialize(key)
+      super
+      self.merge!(:s3_key => key)
+    end
+
+    def s3_key
+      self[:s3_key]
+    end
+
+    def s3_key=(value)
+      self[:s3_key] = value
+    end
+  end
+
   def initialize(opts)
     AWS.eager_autoload! # for thread safety
     check_opts(opts)
-    @should_poll_sqs = opts[:should_poll_sqs]
     @buffer_size = opts[:buffer_size] || 100
+    @use_s3 = opts[:use_s3]
     @queue_name = generate_queue_name(opts)
     @request_count = 0
-    initialize_sqs(opts)
+    initialize_aws(opts)
 
     @waiting = []
     @waiting.taint
@@ -21,16 +34,20 @@ class SuperQueue
     @mutex = Mutex.new
     @in_buffer = []
     @out_buffer = []
+    @deletion_buffer = []
     @deletion_queue = []
 
-    @compressor = Zlib::Deflate.new
-    @decompressor = Zlib::Inflate.new
-
-    @sqs_tracker = Thread.new { poll_sqs } if @should_poll_sqs
-    @gc = Thread.new { collect_garbage }
+    @gc = Thread.new do 
+      begin
+        collect_garbage
+      rescue Exception => @gc_error
+        raise @gc_error
+      end
+    end
   end
 
   def push(p)
+    check_for_errors
     @mutex.synchronize {
       @in_buffer.push p
       clear_in_buffer if @in_buffer.size >= @buffer_size
@@ -44,8 +61,9 @@ class SuperQueue
   end
 
   def pop(non_block=false)
+    check_for_errors
     @mutex.synchronize {
-      while true
+      loop do
         if @out_buffer.empty?
           if fill_out_buffer_from_sqs_queue || fill_out_buffer_from_in_buffer
             return pop_out_buffer
@@ -62,19 +80,22 @@ class SuperQueue
   end
 
   def length
+    check_for_errors
     @mutex.synchronize {
-      return sqs_length + @in_buffer.size + @out_buffer.size
+      sqsl = sqs_length
+      return sqsl + @in_buffer.size + @out_buffer.size
     }
   end
 
-  def empty?
-    self.length == 0
-  end
-
   def num_waiting
+    check_for_errors
     @waiting.size
   end
 
+  def empty?
+    self.length == 0
+  end
+
   def clear
     begin
       self.pop(true)
@@ -84,19 +105,21 @@ class SuperQueue
   end
 
   def shutdown
-    @sqs_tracker.terminate if @should_poll_sqs
     @mutex.synchronize { clear_in_buffer }
     @gc.terminate
-    @mutex.synchronize { clear_deletion_queue }
+    @mutex.synchronize { fill_deletion_queue_from_buffer } if @deletion_buffer.any?
+    @mutex.synchronize { clear_deletion_queue } if @deletion_queue.any?
+    @done = true
   end
 
   def destroy
-    @sqs_tracker.terminate if @should_poll_sqs
     @gc.terminate
-    delete_queue
+    delete_aws_resources
+    @done = true
   end
 
   def sqs_requests
+    check_for_errors
     @request_count
   end
 
@@ -119,10 +142,15 @@ class SuperQueue
   private
 
   #
-  # Amazon SQS methods
+  # Amazon AWS methods
   #
-  def initialize_sqs(opts)
-    create_sqs_connection(opts)
+  def initialize_aws(opts)
+    aws_options = {
+      :access_key_id => opts[:aws_access_key_id], 
+      :secret_access_key => opts[:aws_secret_access_key]
+    }
+    @sqs = AWS::SQS.new(aws_options)
+    @s3 = AWS::S3.new(aws_options)
     create_sqs_queue(opts)
     if opts[:replace_existing_queue] && (sqs_length > 0)
       delete_queue
@@ -130,30 +158,16 @@ class SuperQueue
       sleep 62 # You must wait 60s after deleting a q to create one with the same name
       create_sqs_queue(opts)
     end
+    @bucket = open_s3_bucket
   end
 
-  def create_sqs_connection(opts)
-    aws_options = {
-      :access_key_id => opts[:aws_access_key_id], 
-      :secret_access_key => opts[:aws_secret_access_key]
-    }
-    begin
-      @sqs = AWS::SQS.new(aws_options)
-    rescue Exception => e
-      raise e
-    end
+  def create_sqs_queue(opts)
+    @sqs_queue = find_queue_by_name || new_sqs_queue(opts)
+    check_for_queue_creation_success
   end
 
-  def create_sqs_queue(opts)
-    retries = 0
-    begin
-      @sqs_queue = find_queue_by_name || new_sqs_queue(opts)
-      check_for_queue_creation_success
-    rescue RuntimeError => e
-      retries += 1
-      sleep 1
-      (retries >= 20) ? retry : raise(e)
-    end
+  def open_s3_bucket
+    @s3.buckets[queue_name].exists? ? @s3.buckets[queue_name] : @s3.buckets.create(queue_name)
   end
 
   def find_queue_by_name
@@ -189,17 +203,20 @@ class SuperQueue
     number_of_batches.times do
       batch = []
       10.times do
+        next if @in_buffer.empty?
         p = @in_buffer.shift
-        obj = is_a_link?(p) ? p : encode(p)
-        batch << obj
+        unless should_send_to_s3?(p)
+          batch << encode(p)
+        else
+          batch << encode(send_payload_to_s3(p))
+        end
       end
       batches << batch
     end
 
-    #Ugliness! But I'm not sure how else to tackle this at the moment
     batches.each do |b|
       @request_count += 1
-      @sqs_queue.batch_send(b)
+      @sqs_queue.batch_send(b) if b.any?
     end
   end
 
@@ -207,34 +224,72 @@ class SuperQueue
     messages = []
     number_of_batches = number_of_messages_to_receive / 10
     number_of_batches += 1 if number_of_messages_to_receive % 10
-    number_of_batches.times do |c|
-      batch = @sqs_queue.receive_messages(:limit => 10)
-      batch.each do |m|
-        if is_a_link?(m.body)
-          obj = {:message => m, :payload => m.body}
+    number_of_batches.times do
+      batch = @sqs_queue.receive_messages(:limit => 10).compact
+      batch.each do |message|
+        obj = decode(message.body)
+        unless obj.is_a?(SuperQueue::S3Pointer)
+          messages << {
+            :payload    => obj,
+            :sqs_handle => message
+          }
         else
-          obj = { :message => m, :payload => decode(m.body) }
+          p = fetch_payload_from_s3(obj)
+          messages << {
+            :payload    => p,
+            :sqs_handle => message,
+            :s3_key     => obj.s3_key } if p
         end
-        messages << obj
       end
       @request_count += 1
     end
     messages
   end
 
+  def send_payload_to_s3(p)
+    dump = Marshal.dump(p)
+    digest = Digest::MD5.hexdigest(dump)
+    return digest if @bucket.objects[digest].exists?
+    retryable(:tries => 5) { @bucket.objects[digest].write(dump) }
+    S3Pointer.new(digest)
+  end
+
+  def fetch_payload_from_s3(pointer)
+    payload = nil
+    retries = 0
+    begin
+      payload = Marshal.load(@bucket.objects[pointer.s3_key].read)
+    rescue AWS::S3::Errors::NoSuchKey
+      return nil
+    rescue
+      retries +=1
+      retry if retries < 5
+    end
+    payload
+  end
+
+  def should_send_to_s3?(p)
+    @use_s3
+  end
+
   def sqs_length
     n = @sqs_queue.approximate_number_of_messages
     return n.is_a?(Integer) ? n : 0
   end
 
-  def delete_queue
+  def delete_aws_resources
     @request_count += 1
     @sqs_queue.delete
+    @bucket.delete!
   end
 
   def clear_deletion_queue
     while !@deletion_queue.empty?
-      @sqs_queue.batch_delete(@deletion_queue[0..9])
+      sqs_handles = @deletion_queue[0..9].map { |m| m[:sqs_handle] }.compact
+      s3_keys = @deletion_queue[0..9].map { |m| m[:s3_key] }.compact
+      10.times { @deletion_queue.shift }
+      @sqs_queue.batch_delete(sqs_handles) if sqs_handles.any?
+      s3_keys.each { |key| @bucket.objects[key].delete }
       @request_count += 1
     end
   end
@@ -244,10 +299,15 @@ class SuperQueue
   #
   def fill_out_buffer_from_sqs_queue
     return false if sqs_length == 0
-    @gc.wakeup if @gc.stop? # This is the best time to do GC, because there are no pops happening.
-    while (@out_buffer.size < @buffer_size)
+    nil_count = 0
+    while (@out_buffer.size < @buffer_size) && (nil_count < 2)
       messages = get_messages_from_queue(@buffer_size - @out_buffer.size)
-      messages.each { |m| @out_buffer.push m }
+      if messages.empty?
+        nil_count += 1
+      else
+        messages.each { |m| @out_buffer.push(m) }
+        nil_count = 0
+      end
     end
     !@out_buffer.empty?
   end
@@ -260,9 +320,14 @@ class SuperQueue
     !@out_buffer.empty?
   end
 
+  def fill_deletion_queue_from_buffer
+    @deletion_queue += @deletion_buffer
+    @deletion_buffer = []
+  end
+
   def pop_out_buffer
     m = @out_buffer.shift
-    @deletion_queue << m[:message] if m[:message]
+    @deletion_buffer << { :sqs_handle => m[:sqs_handle], :s3_key => m[:s3_key] }
     m[:payload]
   end
 
@@ -272,6 +337,9 @@ class SuperQueue
     end
   end
 
+  #
+  # Misc helper methods
+  #
   def check_opts(opts)
     raise "Options can't be nil!" if opts.nil?
     raise "Minimun :buffer_size is 5." if opts[:buffer_size] && (opts[:buffer_size] < 5)
@@ -279,30 +347,37 @@ class SuperQueue
     raise "Visbility timeout must be an integer (in seconds)!" if opts[:visibility_timeout] && !opts[:visibility_timeout].is_a?(Integer)
   end
 
-  #
-  # Misc helper methods
-  #
+  def check_for_errors
+    raise @gc_error if @gc_error
+    raise @sqs_error if @sqs_error
+    raise "Queue is no longer available!" if @done == true
+  end
+
   def encode(p)
     Base64.urlsafe_encode64(Marshal.dump(p))
   end
 
   def decode(ser_obj)
-    #puts "Decode: Decoding message #{ser_obj} from base64..."
-    obj = Base64.urlsafe_decode64(ser_obj)
-    #puts "Decode: Object decoded as #{obj}. Doing Marshal load..."
-    ret = Marshal.load(obj)
-    #puts "Decode: Marshal loaded as #{ret}"
-    ret
-  end
-
-  def is_a_link?(s)
-    return false unless s.is_a? String
-    (s[0..6] == "http://") || (s[0..7] == "https://")
+    Marshal.load(Base64.urlsafe_decode64(ser_obj))
   end
 
   def generate_queue_name(opts)
     q_name = opts[:name] || random_name
-    return opts[:namespace] ? "#{@namespace}-#{q_name}" : q_name
+    return opts[:namespace] ? "queue-#{opts[:namespace]}-#{q_name}" : "queue-#{q_name}"
+  end
+
+  def retryable(options = {}, &block)
+    opts = { :tries => 1, :on => Exception }.merge(options)
+
+    retry_exception, retries = opts[:on], opts[:tries]
+
+    begin
+      return yield
+    rescue retry_exception
+      retry if (retries -= 1) > 0
+    end
+
+    yield
   end
 
   #
@@ -315,8 +390,9 @@ class SuperQueue
   end
 
   def random_name
-    o =  [('a'..'z'),('A'..'Z')].map{|i| i.to_a}.flatten
-    (0...15).map{ o[rand(o.length)] }.join
+    o =  [('a'..'z'),('1'..'9')].map{|i| i.to_a}.flatten
+    random_element = (0...25).map{ o[rand(o.length)] }.join
+    "temp-name-#{random_element}"
   end
 
   def queue_name
@@ -326,19 +402,13 @@ class SuperQueue
   #
   # Maintence thread-related methods
   #
-  def poll_sqs
-    loop do
-      @mutex.synchronize { fill_out_buffer_from_sqs_queue || fill_out_buffer_from_in_buffer } if @out_buffer.empty?
-      @mutex.synchronize { clear_in_buffer } if !@in_buffer.empty? && (@in_buffer.size > @buffer_size)
-      Thread.pass
-    end
-  end
 
   def collect_garbage
     loop do
       #This also needs a condition to clear the del queue if there are any handles where the invisibility is about to expire
-      @mutex.synchronize { clear_deletion_queue } if !@deletion_queue.empty? && (@deletion_queue.size >= (@buffer_size / 2))
-      sleep
+      @mutex.synchronize { fill_deletion_queue_from_buffer } if @deletion_buffer.any?
+      Thread.pass
+      @mutex.synchronize { clear_deletion_queue } if @deletion_queue.any?
     end
   end
 end

metadata

@@ -1,15 +1,15 @@
 --- !ruby/object:Gem::Specification
 name: super_queue
 version: !ruby/object:Gem::Version
-  version: 0.2.1
   prerelease:
+  version: 0.3.1
 platform: ruby
 authors:
 - Jon Stokes
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2013-02-11 00:00:00.000000000 Z
+date: 2013-02-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: aws-sdk
@@ -27,7 +27,7 @@ dependencies:
     none: false
   prerelease: false
   type: :runtime
-description: A an SQS-backed queue.
+description: A thread-safe SQS- and S3-backed queue.
 email: jon@jonstokes.com
 executables: []
 extensions: []
@@ -56,8 +56,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
   none: false
 requirements: []
 rubyforge_project:
-rubygems_version: 1.8.25
+rubygems_version: 1.8.24
 signing_key:
 specification_version: 3
-summary: An SQS-backed queue structure for ruby that works just like a normal queue, except it's essentially infinite and can use very little memory.
+summary: A thread-safe, SQS- and S3-backed queue structure for ruby that works just like a normal queue, except it's essentially infinite and can use very little memory.
 test_files: []