oneapm_rpm 1.1.0

data/lib/one_apm/agent/datastores/mongo/obfuscator.rb

@@ -0,0 +1,37 @@
+# encoding: utf-8
+
+module OneApm
+  module Agent
+    module Datastores
+      module Mongo
+        module Obfuscator
+
+          WHITELIST = [:operation].freeze
+
+          def self.obfuscate_statement(source, whitelist=WHITELIST)
+            obfuscated = {}
+            source.each do |key, value|
+              if whitelist.include?(key)
+                obfuscated[key] = value
+              else
+                obfuscated[key] = obfuscate_value(value, whitelist)
+              end
+            end
+
+            obfuscated
+          end
+
+          def self.obfuscate_value(value, whitelist)
+            if value.is_a?(Hash)
+              obfuscate_statement(value, whitelist)
+            elsif value.is_a?(Array)
+              value.map {|v| obfuscate_value(v, whitelist)}
+            else
+              '?'
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/one_apm/agent/datastores/mongo/statement_formatter.rb

@@ -0,0 +1,51 @@
+# encoding: utf-8
+
+require 'one_apm/agent/datastores/mongo/obfuscator'
+
+module OneApm
+  module Agent
+    module Datastores
+      module Mongo
+        module StatementFormatter
+
+          PLAINTEXT_KEYS = [
+            :database,
+            :collection,
+            :operation,
+            :fields,
+            :skip,
+            :limit,
+            :order
+          ]
+
+          OBFUSCATE_KEYS = [
+            :selector
+          ]
+
+          def self.format(statement, operation)
+            return nil unless OneApm::Agent.config[:'mongo.capture_queries']
+
+            result = { :operation => operation }
+
+            PLAINTEXT_KEYS.each do |key|
+              result[key] = statement[key] if statement.key?(key)
+            end
+
+            OBFUSCATE_KEYS.each do |key|
+              if statement.key?(key) && statement[key]
+                obfuscated = obfuscate(statement[key])
+                result[key] = obfuscated if obfuscated
+              end
+            end
+            result
+          end
+
+          def self.obfuscate(statement)
+            statement = Obfuscator.obfuscate_statement(statement) if OneApm::Agent.config[:'mongo.obfuscate_queries']
+            statement
+          end
+        end
+      end
+    end
+  end
+end

data/lib/one_apm/agent/event/event_listener.rb

@@ -0,0 +1,40 @@
+# encoding: utf-8
+
+module OneApm::Agent
+  class EventListener
+
+    attr_accessor :runaway_threshold
+
+    def initialize
+      @events = {}
+      @runaway_threshold = 100
+    end
+
+    def subscribe(event, &handler)
+      @events[event] ||= []
+      @events[event] << handler
+      check_for_runaway_subscriptions(event)
+    end
+
+    def check_for_runaway_subscriptions(event)
+      count = @events[event].size
+      OneApm::Agent.logger.debug("Run-away event subscription on #{event}? Subscribed #{count}") if count > @runaway_threshold
+    end
+
+    def clear
+      @events.clear
+    end
+
+    def notify(event, *args)
+      return unless @events.has_key?(event)
+
+      @events[event].each do |handler|
+        begin
+          handler.call(*args)
+        rescue => err
+          OneApm::Agent.logger.debug("Failure during notify for #{event}", err)
+        end
+      end
+    end
+  end
+end

data/lib/one_apm/agent/event/event_loop.rb

@@ -0,0 +1,191 @@
+# encoding: utf-8
+
+require 'thread'
+
+module OneApm
+  module Agent
+    class EventLoop
+      class Timer
+
+        attr_reader :next_fire_time, :event, :interval, :last_fired_at
+
+        def initialize(interval, event, repeat=false)
+          @interval      = interval
+          @event         = event
+          @repeat        = repeat
+          @started_at    = Time.now
+          @last_fired_at = nil
+          reschedule
+        end
+
+        def reschedule
+          @next_fire_time = calculate_next_fire_time
+        end
+
+        def advance(amount)
+          @next_fire_time -= amount
+        end
+
+        def last_interval_start
+          @last_fired_at || @started_at
+        end
+
+        def calculate_next_fire_time
+          now = Time.now
+          return now if @interval == 0
+          fire_time = @last_fired_at || now
+          while fire_time <= now
+            fire_time += @interval
+          end
+          fire_time
+        end
+
+        def set_fired_time
+          @last_fired_at = Time.now
+        end
+
+        def due?(now=Time.now)
+          now >= @next_fire_time
+        end
+
+        def finished?
+          !@repeat && @last_fired_at
+        end
+
+      end
+
+      def initialize
+        @self_pipe_rd, @self_pipe_wr = IO.pipe
+        @event_queue = Queue.new
+        @stopped     = false
+        @timers      = {}
+
+        @subscriptions = Hash.new { |h,k| h[k] = [] }
+        @subscriptions[:__add_timer] << Proc.new { |t| set_timer(t) }
+        @subscriptions[:__add_event] << Proc.new { |e, blk| @subscriptions[e] << blk }
+      end
+
+      def set_timer(timer)
+        existing_timer = @timers[timer.event]
+
+        if existing_timer
+          elapsed_interval = Time.now - existing_timer.last_interval_start
+          timer.advance(elapsed_interval)
+        end
+
+        @timers[timer.event] = timer
+
+        fire_timer(timer)
+      end
+
+      def next_timeout
+        return nil if @timers.empty?
+        timeout = @timers.values.map(&:next_fire_time).min - Time.now
+        timeout < 0 ? 0 : timeout
+      end
+
+      def stopped?
+        @stopped
+      end
+
+      def stop
+        @stopped = true
+        wakeup
+      end
+
+      def run
+        OneApm::Agent.logger.debug "Running event loop"
+        while !stopped?
+          run_once
+        end
+      end
+
+      def run_once(nonblock=false)
+        wait_to_run(nonblock)
+
+        prune_timers
+        fire_timers
+
+        until @event_queue.empty?
+          evt, args = @event_queue.pop
+          dispatch_event(evt, args)
+          reschedule_timer_for_event(evt)
+        end
+      end
+
+      def wait_to_run(nonblock)
+        timeout = nonblock ? 0 : next_timeout
+        ready = IO.select([@self_pipe_rd], nil, nil, timeout)
+
+        if ready && ready[0] && ready[0][0] && ready[0][0] == @self_pipe_rd
+          @self_pipe_rd.read(1)
+        end
+      end
+
+      def fire_timers
+        @timers.each do |event, timer|
+          fire_timer(timer)
+        end
+      end
+
+      def fire_timer(timer)
+        if timer.due?
+          @event_queue << [timer.event]
+          timer.set_fired_time
+        end
+      end
+
+      def prune_timers
+        @timers.delete_if { |e, t| t.finished? }
+      end
+
+      def dispatch_event(event, args)
+        OneApm::Agent.logger.debug("EventLoop: Dispatching event '#{event}' with #{@subscriptions[event].size} callback(s).")
+
+        errors = []
+        @subscriptions[event].each do |s|
+          begin
+            s.call(*args)
+          rescue OneApm::ForceRestartException, OneApm::ForceDisconnectException
+            raise
+          rescue => e
+            errors << e
+          end
+        end
+
+        if !errors.empty?
+          OneApm::Agent.logger.error "#{errors.size} error(s) running task for event '#{event}' in Agent Event Loop:", *errors
+        end
+      end
+
+      def reschedule_timer_for_event(e)
+        @timers[e].reschedule if @timers[e]
+      end
+
+      def on(event, &blk)
+        fire(:__add_event, event, blk)
+      end
+
+      def fire(event, *args)
+        @event_queue << [event, args]
+        wakeup
+      end
+
+      def fire_every(interval, event)
+        OneApm::Agent.logger.debug "Firing event #{event} every #{interval} seconds."
+        fire(:__add_timer, Timer.new(interval, event, true))
+      end
+
+      def fire_after(interval, event)
+        OneApm::Agent.logger.debug "Firing event #{event} after #{interval} seconds."
+        fire(:__add_timer, Timer.new(interval, event, false))
+      end
+
+      def wakeup
+        @self_pipe_wr.write_nonblock '.'
+      rescue Errno::EAGAIN
+        OneApm::Agent.logger.debug "Failed to wakeup event loop"
+      end
+    end
+  end
+end

data/lib/one_apm/agent/event/worker_loop.rb

@@ -0,0 +1,97 @@
+# encoding: utf-8
+
+require 'thread'
+module OneApm
+  module Agent
+
+    # A worker loop executes a set of registered tasks on a single thread.
+    # A task is a proc or block with a specified call period in seconds.
+    class WorkerLoop
+
+      attr_accessor :period, :propagate_errors
+      attr_reader :iterations
+
+      # Optional argument :duration (in seconds) for how long the worker loop runs
+      # or :limit (integer) for max number of iterations
+      def initialize(opts={})
+        @should_run = true
+        @next_invocation_time = Time.now
+        @period = 60.0
+        @duration = opts[:duration] if opts[:duration]
+        @limit = opts[:limit] if opts[:limit]
+        @iterations = 0
+        @propagate_errors = opts.fetch(:propagate_errors, false)
+      end
+
+      # Reset state that is changed by running the worker loop
+      def setup(period, task)
+        @task = task
+        @period = period if period
+        @should_run = true
+        @iterations = 0
+
+        now = Time.now
+        @deadline = now + @duration if @duration
+        @next_invocation_time = (now + @period)
+      end
+
+      # Run infinitely, calling the registered tasks at their specified
+      # call periods.  The caller is responsible for creating the thread
+      # that runs this worker loop.  This will run the task immediately.
+      def run(period=nil, &block)
+        setup(period, block)
+        while keep_running? do
+          sleep_time = schedule_next_invocation
+          sleep(sleep_time) if sleep_time > 0
+          run_task if keep_running?
+          @iterations += 1
+        end
+      end
+
+      def schedule_next_invocation
+        now = Time.now
+        while @next_invocation_time <= now && @period > 0
+          @next_invocation_time += @period
+        end
+        @next_invocation_time - Time.now
+      end
+
+      # a simple accessor for @should_run
+      def keep_running?
+        @should_run && under_duration? && under_limit?
+      end
+
+      def under_duration?
+        !@deadline || Time.now < @deadline
+      end
+
+      def under_limit?
+        @limit.nil? || @iterations < @limit
+      end
+
+      # Sets @should_run to false. Returns false
+      def stop
+        @should_run = false
+      end
+
+      # Executes the block given to the worker loop, and handles many
+      # possible errors. Also updates the execution time so that the
+      # next run occurs on schedule, even if we execute at some odd time
+      def run_task
+        if @propagate_errors
+          @task.call
+        else
+          begin
+            @task.call
+          rescue OneApm::ForceRestartException, OneApm::ForceDisconnectException
+            # blow out the loop
+            raise
+          rescue => e
+            # Don't blow out the stack for anything that hasn't already propagated
+            ::OneApm::Agent.logger.error "Error running task in Agent Worker Loop:", e
+          end
+        end
+      end
+    end
+  end
+end

data/lib/one_apm/agent/harvester.rb

@@ -0,0 +1,48 @@
+# encoding: utf-8
+
+module OneApm
+  module Agent
+    class Harvester
+
+      attr_accessor :starting_pid
+
+      # Inject target for after_fork call to avoid spawning thread in tests
+      def initialize(events, after_forker = OneApm::Agent)
+        # We intentionally don't set our pid as started at this point.
+        # Startup routines must call mark_started when they consider us set!
+        @starting_pid = nil
+        @after_forker = after_forker
+
+        if events
+          events.subscribe(:start_transaction, &method(:on_transaction))
+        end
+      end
+
+      def on_transaction(*_)
+        return unless restart_in_children_enabled? && needs_restart? && harvest_thread_enabled?
+
+        restart_harvest_thread
+      end
+
+      def mark_started(pid = Process.pid)
+        @starting_pid = pid
+      end
+
+      def needs_restart?(pid = Process.pid)
+        @starting_pid != pid
+      end
+
+      def restart_in_children_enabled?
+        OneApm::Agent.config[:restart_thread_in_children]
+      end
+
+      def harvest_thread_enabled?
+        !OneApm::Agent.config[:disable_harvest_thread]
+      end
+
+      def restart_harvest_thread
+        @after_forker.after_fork(:force_reconnect => true)
+      end
+    end
+  end
+end