Merge pull request #58 from hummingV/rufus-scheduler

Suyog Rao · web-flow · commit 7dc2429efc0a · 2016-08-14T11:54:51.000-07:00
#34, use rufus-scheduler for more flexible scheduling
diff --git a/lib/logstash/inputs/http_poller.rb b/lib/logstash/inputs/http_poller.rb
@@ -4,6 +4,7 @@
 require "logstash/plugin_mixins/http_client"
 require "socket" # for Socket.gethostname
 require "manticore"
+require "rufus/scheduler"
 
 # This Logstash input plugin allows you to call an HTTP API, decode the output of it into event(s), and
 # send them on their merry way. The idea behind this plugins came from a need to read springboot
@@ -33,7 +34,8 @@
 #       }
 #     }
 #     request_timeout => 60
-#     interval => 60
+#     # Supports "cron", "every", "at" and "in" schedules by rufus scheduler
+#     schedule => { cron => "* * * * * UTC"}
 #     codec => "json"
 #     # A hash of request metadata info (timing, response headers, etc.) will be sent here
 #     metadata_target => "http_poller_metadata"
@@ -86,7 +88,20 @@ class LogStash::Inputs::HTTP_Poller < LogStash::Inputs::Base
   config :urls, :validate => :hash, :required => true
 
   # How often (in seconds) the urls will be called
-  config :interval, :validate => :number, :required => true
+  # DEPRECATED. Use 'schedule' option instead.
+  # If both interval and schedule options are specified, interval
+  # option takes higher precedence
+  config :interval, :validate => :number, :deprecated => true
+
+  # Schedule of when to periodically poll from the urls
+  # Format: A hash with
+  #   + key: "cron" | "every" | "in" | "at"
+  #   + value: string
+  # Examples:
+  #   a) { "every" => "1h" }
+  #   b) { "cron" => "* * * * * UTC" }
+  # See: rufus/scheduler for details about different schedule options and value string format
+  config :schedule, :validate => :hash
 
   # Define the target field for placing the received data. If this setting is omitted, the data will be stored at the root (top level) of the event.
   config :target, :validate => :string
@@ -97,17 +112,19 @@ class LogStash::Inputs::HTTP_Poller < LogStash::Inputs::Base
   config :metadata_target, :validate => :string, :default => '@metadata'
 
   public
+  Schedule_types = %w(cron every at in)
   def register
     @host = Socket.gethostname.force_encoding(Encoding::UTF_8)
 
     @logger.info("Registering http_poller Input", :type => @type,
-                 :urls => @urls, :interval => @interval, :timeout => @timeout)
+                 :urls => @urls, :interval => @interval, :schedule => @schedule, :timeout => @timeout)
 
     setup_requests!
   end
 
   def stop
     Stud.stop!(@interval_thread) if @interval_thread
+    @scheduler.stop if @scheduler
   end
 
   private
@@ -160,13 +177,46 @@ def validate_request!(url_or_spec, request)
 
   public
   def run(queue)
+    #interval or schedule must be provided. Must be exclusively either one. Not neither. Not both.
+    raise LogStash::ConfigurationError, "Invalid config. Neither interval nor schedule was specified." \
+      unless @interval ||  @schedule
+    raise LogStash::ConfigurationError, "Invalid config. Specify only interval or schedule. Not both." \
+      if @interval && @schedule
+
+    if @interval
+      setup_interval(queue)
+    elsif @schedule
+      setup_schedule(queue)
+    else
+      #should not reach here
+      raise LogStash::ConfigurationError, "Invalid config. Neither interval nor schedule was specified."
+    end
+  end
+
+  private
+  def setup_interval(queue)
     @interval_thread = Thread.current
     Stud.interval(@interval) do
       run_once(queue)
     end
   end
 
-  private
+  def setup_schedule(queue)
+    #schedule hash must contain exactly one of the allowed keys
+    msg_invalid_schedule = "Invalid config. schedule hash must contain " +
+      "exactly one of the following keys - cron, at, every or in"
+    raise Logstash::ConfigurationError, msg_invalid_schedule if @schedule.keys.length !=1
+    schedule_type = @schedule.keys.first
+    schedule_value = @schedule[schedule_type]
+    raise LogStash::ConfigurationError, msg_invalid_schedule unless Schedule_types.include?(schedule_type)
+
+    @scheduler = Rufus::Scheduler.new(:max_work_threads => 1)
+    #as of v3.0.9, :first_in => :now doesn't work. Use the following workaround instead
+    opts = schedule_type == "every" ? { :first_in => 0.01 } : {} 
+    @scheduler.send(schedule_type, schedule_value, opts) { run_once(queue) }
+    @scheduler.join
+  end
+
   def run_once(queue)
     @requests.each do |name, request|
       request_async(queue, name, request)
diff --git a/logstash-input-http_poller.gemspec b/logstash-input-http_poller.gemspec
@@ -1,6 +1,6 @@
 Gem::Specification.new do |s|
   s.name        = 'logstash-input-http_poller'
-  s.version         = '3.0.3'
+  s.version         = '3.1.0'
   s.licenses    = ['Apache License (2.0)']
   s.summary     = "Poll HTTP endpoints with Logstash."
   s.description     = "This gem is a Logstash plugin required to be installed on top of the Logstash core pipeline using $LS_HOME/bin/logstash-plugin install gemname. This gem is not a stand-alone program"
@@ -22,8 +22,10 @@ Gem::Specification.new do |s|
   s.add_runtime_dependency 'logstash-codec-plain'
   s.add_runtime_dependency 'logstash-mixin-http_client', ">= 2.2.4", "< 5.0.0"
   s.add_runtime_dependency 'stud', "~> 0.0.22"
+  s.add_runtime_dependency 'rufus-scheduler', "~>3.0.9"
 
   s.add_development_dependency 'logstash-codec-json'
   s.add_development_dependency 'logstash-devutils'
   s.add_development_dependency 'flores'
+  s.add_development_dependency 'timecop'
 end
diff --git a/spec/inputs/http_poller_spec.rb b/spec/inputs/http_poller_spec.rb
@@ -1,11 +1,14 @@
 require "logstash/devutils/rspec/spec_helper"
 require 'logstash/inputs/http_poller'
 require 'flores/random'
+require "timecop"
 
 describe LogStash::Inputs::HTTP_Poller do
   let(:metadata_target) { "_http_poller_metadata" }
   let(:queue) { Queue.new }
-  let(:default_interval) { 5 }
+  let(:default_schedule) {
+    { "cron" => "* * * * * UTC" }
+  }
   let(:default_name) { "url1 " }
   let(:default_url) { "http://localhost:1827" }
   let(:default_urls) {
@@ -15,7 +18,7 @@
   }
   let(:default_opts) {
     {
-      "interval" => default_interval,
+      "schedule" => default_schedule,
       "urls" => default_urls,
       "codec" => "json",
       "metadata_target" => metadata_target
@@ -31,9 +34,13 @@
     end
 
     describe "#run" do
-      it "should run at the specified interval" do
-        expect(Stud).to receive(:interval).with(default_interval).once
-        subject.run(double("queue"))
+      it "should setup a scheduler" do
+        runner = Thread.new do
+          subject.run(double("queue"))
+          expect(subject.instance_variable_get("@scheduler")).to be_a_kind_of(Rufus::Scheduler)
+        end
+        runner.kill
+        runner.join
       end
     end
 
@@ -151,6 +158,163 @@
     end
   end
 
+  describe "scheduler configuration" do
+    context "given an interval" do
+      let(:opts) {
+        {
+          "interval" => 2,
+          "urls" => default_urls,
+          "codec" => "json",
+          "metadata_target" => metadata_target
+        }
+      }
+      it "should run once in each interval" do
+        instance = klass.new(opts)
+        instance.register
+        queue = Queue.new
+        runner = Thread.new do
+          instance.run(queue)
+        end
+        #T       0123456
+        #events  x x x x
+        #expects 3 events at T=5
+        sleep 5
+        instance.stop
+        runner.kill
+        runner.join
+        expect(queue.size).to eq(3)
+      end
+    end
+
+    context "given both interval and schedule options" do
+      let(:opts) {
+        {
+          "interval" => 1,
+          "schedule" => { "every" => "5s" },
+          "urls" => default_urls,
+          "codec" => "json",
+          "metadata_target" => metadata_target
+        }
+      }
+      it "should raise ConfigurationError" do
+        instance = klass.new(opts)
+        instance.register
+        queue = Queue.new
+        runner = Thread.new do
+          expect{instance.run(queue)}.to raise_error(LogStash::ConfigurationError)
+        end
+        instance.stop
+        runner.kill
+        runner.join
+      end
+    end
+
+    context "given 'cron' expression" do
+      let(:opts) {
+        {
+          "schedule" => { "cron" => "* * * * * UTC" },
+          "urls" => default_urls,
+          "codec" => "json",
+          "metadata_target" => metadata_target
+        }
+      }
+      it "should run at the schedule" do
+        instance = klass.new(opts)
+        instance.register
+        Timecop.travel(Time.new(2000,1,1,0,0,0,'+00:00'))
+        Timecop.scale(60)
+        queue = Queue.new
+        runner = Thread.new do
+          instance.run(queue)
+        end
+        sleep 3
+        instance.stop
+        runner.kill
+        runner.join
+        expect(queue.size).to eq(2)
+        Timecop.return
+      end
+    end
+
+    context "given 'at' expression" do
+      let(:opts) {
+        {
+          "schedule" => { "at" => "2000-01-01 00:05:00 +0000"},
+          "urls" => default_urls,
+          "codec" => "json",
+          "metadata_target" => metadata_target
+        }
+      }
+      it "should run at the schedule" do
+        instance = klass.new(opts)
+        instance.register
+        Timecop.travel(Time.new(2000,1,1,0,0,0,'+00:00'))
+        Timecop.scale(60 * 5)
+        queue = Queue.new
+        runner = Thread.new do
+          instance.run(queue)
+        end
+        sleep 2
+        instance.stop
+        runner.kill
+        runner.join
+        expect(queue.size).to eq(1)
+        Timecop.return
+      end
+    end
+
+    context "given 'every' expression" do
+      let(:opts) {
+        {
+          "schedule" => { "every" => "2s"},
+          "urls" => default_urls,
+          "codec" => "json",
+          "metadata_target" => metadata_target
+        }
+      }
+      it "should run at the schedule" do
+        instance = klass.new(opts)
+        instance.register
+        queue = Queue.new
+        runner = Thread.new do
+          instance.run(queue)
+        end
+        #T       0123456
+        #events  x x x x
+        #expects 3 events at T=5
+        sleep 5
+        instance.stop
+        runner.kill
+        runner.join
+        expect(queue.size).to eq(3)
+      end
+    end
+
+    context "given 'in' expression" do
+      let(:opts) {
+        {
+          "schedule" => { "in" => "2s"},
+          "urls" => default_urls,
+          "codec" => "json",
+          "metadata_target" => metadata_target
+        }
+      }
+      it "should run at the schedule" do
+        instance = klass.new(opts)
+        instance.register
+        queue = Queue.new
+        runner = Thread.new do
+          instance.run(queue)
+        end
+        sleep 3
+        instance.stop
+        runner.kill
+        runner.join
+        expect(queue.size).to eq(1)
+      end
+    end
+  end
+
   describe "events" do
     shared_examples("matching metadata") {
       let(:metadata) { event.get(metadata_target) }
@@ -292,7 +456,9 @@
         }
         let(:opts) {
           {
-            "interval" => default_interval,
+            "schedule" => {
+              "cron" => "* * * * * UTC"
+              },
             "urls" => {
               default_name => url
             },
