parallel588_polipus 0.4.0

data/lib/polipus/signal_handler.rb

@@ -0,0 +1,42 @@
+# encoding: UTF-8
+require 'singleton'
+module Polipus
+  class SignalHandler
+    include Singleton
+    attr_accessor :terminated
+    attr_accessor :enabled
+
+    def initialize
+      self.terminated = false
+      self.enabled = false
+    end
+
+    def self.enable
+      trap(:INT)  do
+        exit unless self.enabled?
+        terminate
+      end
+      trap(:TERM) do
+        exit unless self.enabled?
+        terminate
+      end
+      instance.enabled = true
+    end
+
+    def self.disable
+      instance.enabled = false
+    end
+
+    def self.terminate
+      instance.terminated = true
+    end
+
+    def self.terminated?
+      instance.terminated
+    end
+
+    def self.enabled?
+      instance.enabled
+    end
+  end
+end

data/lib/polipus/storage.rb

@@ -0,0 +1,31 @@
+require 'polipus/storage/base'
+
+module Polipus
+  module Storage
+    COLLECTION = 'pages'
+
+    def self.mongo_store(mongo = nil, collection = COLLECTION, except = [])
+      require 'polipus/storage/mongo_store'
+      mongo ||= Mongo::Connection.new('localhost', 27_017, pool_size: 15, pool_timeout: 5).db('polipus')
+      fail 'First argument must be an instance of Mongo::DB' unless mongo.is_a?(Mongo::DB)
+      self::MongoStore.new(mongo: mongo, collection: collection, except: except)
+    end
+
+    def self.rethink_store(conn = nil, table = COLLECTION, except = [])
+      require 'polipus/storage/rethink_store'
+      conn ||= RethinkDB::RQL.new.connect(host: 'localhost', port: 28_015, db: 'polipus' )
+      fail "First argument must be a RethinkDB::Connection, got `#{conn.class}`" unless conn.is_a?(RethinkDB::Connection)
+      self::RethinkStore.new(conn: conn, table: table,  except: except)
+    end
+
+    def self.dev_null
+      require 'polipus/storage/dev_null'
+      self::DevNull.new
+    end
+
+    def self.memory_store
+      require 'polipus/storage/memory_store'
+      self::MemoryStore.new
+    end
+  end
+end

data/lib/polipus/storage/base.rb

@@ -0,0 +1,20 @@
+# encoding: UTF-8
+require 'uri'
+
+module Polipus
+  module Storage
+    class Base
+      attr_accessor :include_query_string_in_uuid
+
+      protected
+
+      def uuid(page)
+        if @include_query_string_in_uuid.nil?
+          @include_query_string_in_uuid = true
+        end
+        url_to_hash = @include_query_string_in_uuid ? page.url.to_s : page.url.to_s.gsub(/\?.*$/, '')
+        Digest::MD5.hexdigest(url_to_hash)
+      end
+    end
+  end
+end

data/lib/polipus/storage/dev_null.rb

@@ -0,0 +1,35 @@
+# encoding: UTF-8
+module Polipus
+  module Storage
+    class DevNull < Base
+      def initialize(_options = {})
+      end
+
+      def add(_page)
+      end
+
+      def exists?(_page)
+        false
+      end
+
+      def get(_page)
+        nil
+      end
+
+      def remove(_page)
+        false
+      end
+
+      def count
+        0
+      end
+
+      def each
+        yield nil
+      end
+
+      def clear
+      end
+    end
+  end
+end

data/lib/polipus/storage/memory_store.rb

@@ -0,0 +1,56 @@
+# encoding: UTF-8
+require 'thread'
+module Polipus
+  module Storage
+    class MemoryStore < Base
+      def initialize(_options = {})
+        @store = {}
+        @semaphore = Mutex.new
+      end
+
+      def add(page)
+        @semaphore.synchronize do
+          u = uuid(page)
+          @store[u] = page
+          u
+        end
+      end
+
+      def exists?(page)
+        @semaphore.synchronize do
+          @store.key?(uuid(page))
+        end
+      end
+
+      def get(page)
+        @semaphore.synchronize do
+          @store[uuid(page)]
+        end
+      end
+
+      def remove(page)
+        @semaphore.synchronize do
+          @store.delete(uuid(page))
+        end
+      end
+
+      def count
+        @semaphore.synchronize do
+          @store.count
+        end
+      end
+
+      def each
+        @store.each do |k, v|
+          yield k, v
+        end
+      end
+
+      def clear
+        @semaphore.synchronize do
+          @store = Hash.new
+        end
+      end
+    end
+  end
+end

data/lib/polipus/storage/mongo_store.rb

@@ -0,0 +1,90 @@
+# encoding: UTF-8
+require 'mongo'
+require 'zlib'
+require 'thread'
+
+module Polipus
+  module Storage
+    class MongoStore < Base
+      BINARY_FIELDS = %w(body headers data)
+      def initialize(options = {})
+        @mongo      = options[:mongo]
+        @collection = options[:collection]
+        @mongo.create_collection(@collection)
+        begin
+          @mongo[@collection].ensure_index(:uuid, unique: true, dropDups: true, background: true)
+        rescue Exception
+        end
+
+        @compress_body = options[:compress_body] ||= true
+        @except = options[:except] ||= []
+        @semaphore = Mutex.new
+      end
+
+      def add(page)
+        @semaphore.synchronize do
+          obj = page.to_hash
+          @except.each { |e| obj.delete e.to_s }
+          obj['uuid'] = uuid(page)
+          obj['body'] = Zlib::Deflate.deflate(obj['body']) if @compress_body && obj['body']
+          BINARY_FIELDS.each do |field|
+            obj[field] = BSON::Binary.new(obj[field]) unless obj[field].nil?
+          end
+          @mongo[@collection].update({ uuid: obj['uuid'] }, obj, upsert: true, w: 1)
+          obj['uuid']
+        end
+      end
+
+      def exists?(page)
+        @semaphore.synchronize do
+          doc = @mongo[@collection].find({ uuid: uuid(page) }, { fields: [:_id] }).limit(1).first
+          !doc.nil?
+        end
+      end
+
+      def get(page)
+        @semaphore.synchronize do
+          data = @mongo[@collection].find(uuid: uuid(page)).limit(1).first
+          return load_page(data) if data
+        end
+      end
+
+      def remove(page)
+        @semaphore.synchronize do
+          @mongo[@collection].remove(uuid: uuid(page))
+        end
+      end
+
+      def count
+        @mongo[@collection].count
+      end
+
+      def each
+        @mongo[@collection].find({}, timeout: false) do |cursor|
+          cursor.each do |doc|
+            page = load_page(doc)
+            yield doc['uuid'], page
+          end
+        end
+      end
+
+      def clear
+        @mongo[@collection].drop
+      end
+
+      private
+
+      def load_page(hash)
+        BINARY_FIELDS.each do |field|
+          hash[field] = hash[field].to_s
+        end
+        hash['body'] = Zlib::Inflate.inflate(hash['body']) if @compress_body && hash['body'] && !hash['body'].empty?
+        page = Page.from_hash(hash)
+        if page.fetched_at.nil?
+          page.fetched_at = hash['_id'].generation_time.to_i
+        end
+        page
+      end
+    end
+  end
+end

data/lib/polipus/storage/rethink_store.rb

@@ -0,0 +1,90 @@
+# encoding: UTF-8
+require 'rethinkdb'
+require 'thread'
+require 'zlib'
+
+module Polipus
+  module Storage
+    class RethinkStore < Base
+      BINARY_FIELDS = %w(body headers data)
+      def initialize(options = {})
+        @r       = RethinkDB::RQL.new
+        @rethink = options[:conn]
+        @table   = options[:table]
+
+        unless @r.table_list.run(@rethink).include?(@table)
+          @r.table_create(@table).run(@rethink)
+          @r.table(@table).index_create('created_at')
+        end
+
+        @compress_body = options[:compress_body] ||= true
+        @except = options[:except] ||= []
+        @semaphore = Mutex.new
+      end
+
+      def add(page)
+        @semaphore.synchronize do
+          obj = page.to_hash
+          @except.each { |e| obj.delete e.to_s }
+          obj[:id] = uuid(page)
+          obj['body'] = Zlib::Deflate.deflate(obj['body']) if @compress_body && obj['body']
+          obj['created_at'] ||= Time.now.to_i
+          BINARY_FIELDS.each do |field|
+            # Use some marshalling?
+            obj[field] = @r.binary(obj[field]) unless obj[field].nil?
+          end
+
+          @r.table(@table).insert(obj).run(@rethink, durability: 'soft')
+          obj[:id]
+        end
+      end
+
+      def exists?(page)
+        @semaphore.synchronize do
+          doc = @r.table(@table).get(uuid(page)).run(@rethink)
+          !doc.nil?
+        end
+      end
+
+      def get(page)
+        @semaphore.synchronize do
+          data = @r.table(@table).get(uuid(page)).run(@rethink)
+          return load_page(data) if data
+        end
+      end
+
+      def remove(page)
+        @semaphore.synchronize do
+          @r.table(@table).get(uuid(page)).delete.run(@rethink)
+        end
+      end
+
+      def count
+        @r.table(@table).count.run(@rethink)
+      end
+
+      def each
+        @r.table(@table).run(@rethink).each do |doc|
+          page = load_page(doc)
+          yield doc[:id], page
+        end
+      end
+
+      def clear
+        @r.table(@table).delete.run(@rethink)
+      end
+
+      private
+
+      def load_page(hash)
+        BINARY_FIELDS.each do |field|
+          hash[field] = hash[field].to_s
+        end
+        hash['body'] = Zlib::Inflate.inflate(hash['body']) if @compress_body && hash['body'] && !hash['body'].empty?
+        page = Page.from_hash(hash)
+        page.fetched_at ||= hash['created_at']
+        page
+      end
+    end
+  end
+end

data/lib/polipus/url_tracker.rb

@@ -0,0 +1,21 @@
+# encoding: UTF-8
+module Polipus
+  module UrlTracker
+    def self.bloomfilter(options = {})
+      require 'polipus/url_tracker/bloomfilter'
+      options[:size]       ||= 1_000_000
+      options[:error_rate] ||= 0.01
+      options[:key_name]   ||= 'polipus-bloomfilter'
+      options[:redis]      ||= Redis.current
+      options[:driver]     ||= 'lua'
+      self::Bloomfilter.new options
+    end
+
+    def self.redis_set(options = {})
+      require 'polipus/url_tracker/redis_set'
+      options[:redis]      ||= Redis.current
+      options[:key_name]   ||= 'polipus-set'
+      self::RedisSet.new options
+    end
+  end
+end

data/lib/polipus/url_tracker/bloomfilter.rb

@@ -0,0 +1,27 @@
+# encoding: UTF-8
+require 'redis-bloomfilter'
+module Polipus
+  module UrlTracker
+    class Bloomfilter
+      def initialize(options = {})
+        @bf = Redis::Bloomfilter.new options
+      end
+
+      def visited?(url)
+        @bf.include?(url)
+      end
+
+      def visit(url)
+        @bf.insert url
+      end
+
+      def remove(url)
+        @bf.remove url
+      end
+
+      def clear
+        @bf.clear
+      end
+    end
+  end
+end

data/lib/polipus/url_tracker/redis_set.rb

@@ -0,0 +1,27 @@
+# encoding: UTF-8
+module Polipus
+  module UrlTracker
+    class RedisSet
+      def initialize(options = {})
+        @redis    = options[:redis] || Redis.current
+        @set_name = options[:key_name]
+      end
+
+      def visited?(url)
+        @redis.sismember(@set_name, url)
+      end
+
+      def visit(url)
+        @redis.sadd(@set_name, url)
+      end
+
+      def remove(url)
+        @redis.srem(@set_name, url, 0)
+      end
+
+      def clear
+        @redis.del @set_name
+      end
+    end
+  end
+end