jp 0.0.2

data/lib/rb/jp/gen-rb/jp_constants.rb

@@ -0,0 +1,10 @@
+#
+# Autogenerated by Thrift
+#
+# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
+#
+
+require 'jp_types'
+
+  module Jp
+end

data/lib/rb/jp/gen-rb/jp_types.rb

@@ -0,0 +1,61 @@
+#
+# Autogenerated by Thrift
+#
+# DO NOT EDIT UNLESS YOU ARE SURE THAT YOU KNOW WHAT YOU ARE DOING
+#
+
+require 'fb303_types'
+
+
+module Jp
+    class Job
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      MESSAGE = 1
+      ID = 2
+
+      FIELDS = {
+        MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message', :binary => true},
+        ID => {:type => ::Thrift::Types::STRING, :name => 'id', :binary => true}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field message is unset!') unless @message
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field id is unset!') unless @id
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class EmptyPool < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+
+      FIELDS = {
+
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class NoSuchPool < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+
+      FIELDS = {
+
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+  end

data/lib/rb/jp/producer.rb

@@ -0,0 +1,45 @@
+#!/usr/bin/env ruby
+
+require 'jp/client'
+
+autoload :JSON, 'json'
+
+module Jp
+	class AbstractProducer < AbstractClient
+		def initialize queue, options = {}
+			super queue, options
+		end
+		def add message
+			@client.add @queue, translate(message)
+		end
+		private
+		def translate message
+			raise NotImplementedError.new
+		end
+	end
+
+	class TextProducer < AbstractProducer
+		private
+		def translate message
+			message
+		end
+	end
+
+	class JsonProducer < AbstractProducer
+		private
+		def translate message
+			JSON::dump message
+		end
+	end
+
+	class ThriftProducer < AbstractProducer
+		def initialize queue, options = {}
+			super queue, options
+			@serializer = Thrift::Serializer.new
+		end
+		private
+		def translate message
+			@serializer.serialize message
+		end
+	end
+end

data/lib/rb/jp/server/configuration_loader.rb

@@ -0,0 +1,71 @@
+#!/usr/bin/env ruby
+
+autoload :YAML, 'yaml'
+
+# Find a config file and load it
+module Jp
+  module Server
+    module ConfigurationLoader
+    	def self.server_config
+    		name = config_file('jp-config')
+        if name =~ /.ya?ml$/
+          load_yaml_config_file(name)
+        else
+          load_ruby_config_file(name)
+        end
+    	end
+
+    	private
+      def self.load_yaml_config_file file
+        config = YAML.load(File.read(file))
+        symbolize_hash(config)
+      end
+
+      def self.symbolize_hash hash
+        hash.dup.each do |k,v|
+          if v.is_a? Hash
+            symbolize_hash v
+          end
+          if k.is_a? String
+            hash[k.to_sym] = v
+          end
+        end
+        hash
+      end
+
+      def self.load_ruby_config_file file
+        load file
+    		c = JpConfig.new
+    		def c.options
+    			m = self.public_methods - Object.public_methods - [:options]
+    			h = Hash.new
+    			m.each do |method|
+    				h[method] = self.send method
+    			end
+    			h
+    		end
+    		c.options
+      end
+
+    	def self.config_file name
+    		if ARGV[0] && File.exists?(ARGV[0])
+    			return File.expand_path ARGV[0]
+    		end
+    
+    		# Prefixes
+    		[
+    			'./',
+    			'~/',
+    			'~/.',
+    			'/etc/',
+    		].each do |prefix|
+    			full = File.expand_path "%s%s.rb" % [prefix, name]
+    			if File.exists? full
+    				return full
+    			end
+    		end
+    		raise 'No config file found.'
+    	end
+    end
+  end
+end

data/lib/rb/jp/server/handler.rb

@@ -0,0 +1,94 @@
+require 'jp/thrift'
+require 'jp/server/mongo_connection'
+require 'jp/server/pools'
+
+module Jp
+  module Server
+    # Implementation class.
+    #
+    # Does not include Thrift code, etc.
+    class Handler
+      include MongoConnection
+      include Pools
+      attr_reader :retry_attempts, :retry_delay
+      def initialize options = {}
+        # Option defaults
+        defaults = {
+          :mongo_retry_attempts => 10,
+          :mongo_retry_delay    => 1,
+        }
+        options = defaults.merge(options)
+        # Copy with/deal with options
+        @retry_attempts = options[:mongo_retry_attempts]
+        @retry_delay    = options[:mongo_retry_delay]
+
+        load_pools(options)
+        connect_to_mongo(options)
+      end
+
+      def add pool, message
+        raise NoSuchPool.new unless pools.member? pool
+
+        doc = {
+          'message'     => message,
+          'locked'       => false,
+        }
+
+        rescue_connection_failure do
+          database[pool].insert doc
+        end
+      end
+
+      def acquire pool
+        raise NoSuchPool.new unless pools.member? pool
+        now = Time.new.to_i
+        doc = {}
+        begin
+          rescue_connection_failure do
+            doc = database[pool].find_and_modify(
+              query: {
+                'locked' => false
+              },
+              update: {
+                '$set' => {
+                  'locked'       => now,
+                  'locked_until' => now + pools[pool][:timeout],
+                },
+              }
+            )
+          end
+        rescue Mongo::OperationFailure => e
+          raise EmptyPool
+        end
+        job = Job.new
+        job.message = doc['message']
+        job.id = doc['_id'].to_s
+        job
+      end
+
+      def purge pool, id
+        raise NoSuchPool.new unless pools.member? pool
+        rescue_connection_failure do
+          database[pool].remove _id: BSON::ObjectId(id)
+        end
+      end
+
+      # Ensure retry upon failure
+      # Based on code from http://www.mongodb.org/display/DOCS/Replica+Pairs+in+Ruby
+      def rescue_connection_failure
+        success = false
+        retries = 0
+        while !success
+          begin
+            yield
+            success = true
+          rescue Mongo::ConnectionFailure => ex
+            retries += 1
+            raise ex if retries >= @retry_attempts
+            sleep(@retry_delay)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rb/jp/server/instrumented_server.rb

@@ -0,0 +1,87 @@
+require 'jp/server/server'
+
+module Jp
+  module Server
+    class InstrumentedServer < Jp::Server::Server
+      def initialize options = {}
+        options[:jp_server] ||= Jp::Server::Server.new options.merge(thrift_processor: JobPoolInstrumented::Processor.new(self))
+        @server = options[:jp_server]
+        @pools = options[:pools].keys
+        @add_count = Hash.new 0
+        @acquire_count = Hash.new 0
+        @purge_count = Hash.new 0
+        @empty_count = Hash.new 0
+      end
+
+      # Readers
+
+      def start_time
+        @server.aliveSince
+      end
+
+      def pools
+        @pools
+      end
+
+      def add_count pool
+        raise Jp::NoSuchPool.new unless @pools.include? pool
+        @add_count[pool]
+      end
+
+      def acquire_count pool
+        raise Jp::NoSuchPool.new unless @pools.include? pool
+        @acquire_count[pool]
+      end
+
+      def empty_count pool
+        raise Jp::NoSuchPool.new unless @pools.include? pool
+        @empty_count[pool]
+      end
+
+      def purge_count pool
+        raise Jp::NoSuchPool.new unless @pools.include? pool
+        @purge_count[pool]
+      end
+
+      # Data collectors
+
+      def serve
+        @server.serve
+      end
+
+      def add pool, message
+        result = @server.add pool, message
+        @add_count[pool] += 1
+        result
+      end
+
+      def acquire pool
+        begin
+          result = @server.acquire pool
+          @acquire_count[pool] += 1
+        rescue EmptyPool => e
+          @empty_count[pool] += 1
+          raise e
+        end
+        result
+      end
+
+      def purge pool, id
+        result = @server.purge pool, id
+        @purge_count[pool] += 1
+        result
+      end
+
+      def getCounters
+        counters = Hash.new
+        pools.each do |pool|
+          counters["#{pool}.added"]    = add_count(pool)
+          counters["#{pool}.acquired"] = acquire_count(pool)
+          counters["#{pool}.empty"]    = empty_count(pool)
+          counters["#{pool}.purged"]   = purge_count(pool)
+        end
+        counters
+      end
+    end
+  end
+end

data/lib/rb/jp/server/mongo_connection.rb

@@ -0,0 +1,29 @@
+module Jp
+  module Server
+    module MongoConnection
+      attr_reader :database
+      def connect_to_mongo options
+        defaults = {
+          :mongo_uri            => 'mongodb://localhost',
+          :mongo_pool_size      => 10,
+          :mongo_pool_timeout   => 60,
+        }
+        options = defaults.merge(options)
+        unless options[:mongo_db]
+          raise ArgumentError.new "mongo_db option must be specified"
+        end
+        # Connect to mongodb
+        if options.member? :injected_mongo_database then
+          @database = options[:injected_mongo_database]
+        else
+          connection = Mongo::Connection.from_uri(
+            options[:mongo_uri],
+            :pool_size => options[:mongo_pool_size],
+            :timeout   => options[:mongo_pool_timeout],
+          )
+          @database = connection.db(options[:mongo_db])
+        end
+      end
+    end
+  end
+end

data/lib/rb/jp/server/pools.rb

@@ -0,0 +1,28 @@
+module Jp
+  module Server
+    module Pools
+      attr_reader :pools
+      protected
+      def load_pools options
+        defaults = {
+          :default_timeout => 3600, # 1 hour
+        }
+        options = defaults.merge(options)
+
+        unless options[:pools]
+          raise ArgumentError.new "pools option must be specified"
+        end
+        if options[:pools].empty?
+          raise ArgumentError.new "pools option must not be empty"
+        end
+
+        @pools = Hash.new
+        options[:pools].each do |name, data|
+          data[:timeout]          ||= options[:default_timeout]
+          data[:cleanup_interval] ||= data[:timeout]
+          @pools[name] = data
+        end
+      end
+    end
+  end
+end

data/lib/rb/jp/server/server.rb

@@ -0,0 +1,74 @@
+require 'jp/thrift'
+require 'jp/server/handler'
+require 'jp/server/unlocker'
+
+require 'mongo'
+include Jp
+
+module Jp
+  module Server
+    class Server < Jp::Server::Handler
+      def initialize options = {}
+        super options
+        options[:port_number] ||= 9090
+    
+        # Setup Thrift server (allowing dependency injection)
+        if options.member? :injected_thrift_server then
+          @server = options[:injected_thrift_server]
+        else
+          processor = options[:thrift_processor] # For testing, and allow instrumented server to override
+          processor ||= JobPool::Processor.new self
+          socket = Thrift::ServerSocket.new options[:port_number]
+          transportFactory = Thrift::BufferedTransportFactory.new
+    
+          @server = Thrift::ThreadedServer.new processor, socket, transportFactory
+        end
+    
+        @unlocker = nil
+        unless options[:skip_embedded_unlocker]
+          if options.member? :injected_unlocker
+            @unlocker = options[:injected_unlocker]
+          else
+            @unlocker = Jp::Server::Unlocker.new options
+          end
+        end
+    
+        @start_time = Time.new.to_i
+      end
+    
+      def serve
+        # Look for expired entries
+        @unlocker ||= nil
+        unlocker_thread = nil
+        if @unlocker
+          unlocker_thread = Thread.new do
+            @unlocker.run
+          end
+        end
+        @server.serve
+        unlocker_thread.join if unlocker_thread
+      end
+    
+    
+      # fb303:
+      def getName; 'jp'; end
+      def getVersion; '0.0.1'; end
+      def getStatus; Fb_status::ALIVE; end
+      def getStatusDetails; 'nothing to see here; move along'; end
+      def aliveSince; @start_time; end
+      def shutdown
+        STDERR.write "Shutdown requested via fb303\n"
+        exit
+      end
+      # fb303 stubs:
+      def setOption(key, value); end
+      def getOption(key); end
+      def getOptions; Hash.new; end
+      def getCpuProfile(seconds); String.new; end
+      def reinitialize; end
+      # fb303 stubs properly implemented in JpInstrumentedServer:
+      def getCounters; Hash.new; end
+      def getCounter(name); 0; end
+    end
+  end
+end