<commit_before>require 'mongo'

module Karait
  class Queue
    
    include Karait
    
    MESSAGES_READ = 10
    NO_OBJECT_FOUND_ERROR = 'No matching object found'
    
    def initialize(opts={})
      set_instance_variables opts
      create_mongo_connection
    end
    
    def write(message, opts={})
      if message.class == Hash
        message_dict = message
      else
        message_dict = message.to_hash
      end
      
      message_dict[:_meta] = {
        :expire => opts.fetch(:expire, -1.0),
        :timestamp => Time.now().to_f,
        :expired => false,
        :visible_after => -1.0
      }
      
      message_dict[:_meta][:routing_key] = opts.fetch(:routing_key) if opts[:routing_key]
      
      @queue_collection.insert(message_dict, :safe => true)
    end
    
    def read(opts={})
      opts = {
        :messages_read => 10,
        :visibility_timeout => -1.0,
        :routing_key => nil
      }.update(opts)
      
      current_time = Time.new.to_f
      messages = []
      
      query = {
          '_meta.expired' => false,
          '_meta.visible_after' => {
            '$lt' => current_time
          }
      }
      if opts[:routing_key] != nil
        query['_meta.routing_key'] = opts[:routing_key]
      else
        query['_meta.routing_key'] = {
          '$exists' => false
        }
      end
      
      update = false
      if opts[:visibility_timeout] > -1.0
        update = {
            '$set' => {
              '_meta.visible_after' => current_time + opts[:visibility_timeout]
            }
        }
      end
      
      raw_messages = []
      
      if update
        (0..opts[:messages_read]).each do
          begin
            
            raw_message = @queue_collection.find_and_modify(:query => query, :update => update)
            
            if raw_message
              raw_messages << raw_message
            else
              break
            end
          
          rescue Mongo::OperationFailure => operation_failure
            if not  operation_failure.to_s.match(Queue::NO_OBJECT_FOUND_ERROR)
              raise operation_failure
            end
          end
          
        end
      else<commit_msg>Added blocking functionality for ruby, same argument names and defaults as pythong<commit_after>require 'mongo'

module Karait
  class Queue
    
    include Karait
    
    MESSAGES_READ = 10
    NO_OBJECT_FOUND_ERROR = 'No matching object found'
    
    def initialize(opts={})
      set_instance_variables opts
      create_mongo_connection
    end
    
    def write(message, opts={})
      if message.class == Hash
        message_dict = message
      else
        message_dict = message.to_hash
      end
      
      message_dict[:_meta] = {
        :expire => opts.fetch(:expire, -1.0),
        :timestamp => Time.now().to_f,
        :expired => false,
        :visible_after => -1.0
      }
      
      message_dict[:_meta][:routing_key] = opts.fetch(:routing_key) if opts[:routing_key]
      
      @queue_collection.insert(message_dict, :safe => true)
    end
    
    def read(opts={})
      opts = {
        :messages_read => 10,
        :visibility_timeout => -1.0,
        :routing_key => nil,
        :block => false,
        :polling_interval_seconds =>10,
        :timeout_in_seconds=> nil
      }.update(opts)
      
      current_time = Time.new.to_f
      messages = []
      
      query = {
          '_meta.expired' => false,
          '_meta.visible_after' => {
            '$lt' => current_time
          }
      }
      if opts[:routing_key] != nil
        query['_meta.routing_key'] = opts[:routing_key]
      else
        query['_meta.routing_key'] = {
          '$exists' => false
        }
      end
      
      update = false
      if opts[:visibility_timeout] > -1.0
        update = {
            '$set' => {
              '_meta.visible_after' => current_time + opts[:visibility_timeout]
            }
        }
      end
      
      raw_messages = []
      # if we want to block, loop and sleep until messages are available (or we time out)
      if opts[:block]
          # as long as there are no message, block (unless we time out)
          while @queue_collection.find(query).count() ==0
               # if we have timed out, exit loop
              if opts[:timeout_in_seconds] != nil and current_time + opts[:timeout_in_seconds] <= time.new.to_f
                  break
              end
              #sleep before polling again
              sleep opts[:polling_interval_seconds]
          end
       end


      if update
        (0..opts[:messages_read]).each do
          begin
            
            raw_message = @queue_collection.find_and_modify(:query => query, :update => update)
            
            if raw_message
              raw_messages << raw_message
            else
              break
            end
          
          rescue Mongo::OperationFailure => operation_failure
            if not  operation_failure.to_s.match(Queue::NO_OBJECT_FOUND_ERROR)
              raise operation_failure
            end
          end
          
        end
      else
{'fraction_non_alphanumeric': 0.09633385335413416, 'fraction_numerical': 0.005655226209048362, 'mean_word_length': 1.5786827551533433, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 42, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1154681', 'n_tokens_mistral': 1474, 'n_tokens_neox': 1461, 'n_words': 315}