cli.rb in jflow-0.4.3

- old
+ new

@@ -23,11 +23,11 @@
 
     def start_workers
       number_of_workers.times do
         worker_threads << worker_thread
       end
-      worker_threads << stats_thread if enable_stats
+      worker_threads << maintenance_thread if enable_stats || is_ec2_instance?
       worker_threads.each(&:join)
     end
 
     # here we want to handle all cases for clean kill of the workers.
     # there is two state on which the workers can be
@@ -61,22 +61,38 @@
         c.cloudwatch_client = Aws::CloudWatch::Client.new
       end
       JFlow.load_activities
     end
 
-    def stats_thread
+    def maintenance_thread
       JFlow::WorkerThread.new do
         Thread.current.set_state(:polling)
         stats = JFlow::Stats.new(@domain, @tasklist)
+        protector = JFlow::TerminationProtector.new
         loop do
           break if Thread.current.marked_for_shutdown?
-          stats.tick
+          JFlow.configuration.logger.debug "Should protect?: #{should_protect?}"
+          protector.set_protection(should_protect?) if is_ec2_instance?
+          stats.tick if enable_stats
           sleep 30
         end
       end
     end
 
+    # returns true if any thread if working on someting
+    def should_protect?
+      worker_threads.each do |thread|
+        return true if thread.currently_working?
+      end
+      return false
+    end
+
+    # This should exist on all EC2 instances
+    def is_ec2_instance?
+      @ec2_instance ||= File.exist?('/sys/hypervisor/uuid')
+    end
+
     def log(str)
       JFlow.configuration.logger.info str
     end
 
     def worker_thread
@@ -93,6 +109,6 @@
       validator = HashValidator.validate(options, VALIDATION)
       raise "configuration is invalid! #{validator.errors}" unless validator.valid?
     end
 
   end
-end
\ No newline at end of file
+end