Merge pull request #54 from mattrayner/mattrayner/48_expose-request-id-within-requests

[#48] Expose request id within parent app when using middleware

.travis.yml

@@ -6,6 +6,7 @@ rvm:
   - 2.2.0
   - 2.3.3
   - 2.4.0
+  - 2.5.0
   - ruby-head
 
 cache: bundler

CHANGELOG.md

@@ -3,7 +3,10 @@
 This file needs to be updated with every significant pull request. It is used to write down release notes.
 
 ## Version 0.5.6
-
+* Expose request id to parent Rack application when using `ApplicationInsights::Rack::TrackRequest` middleware through `env['ApplicationInsights.request.id']`.
+* Implement operation context functionality for `ApplicationInsights::Rack::TrackRequest`.
+* Add functionality to accept a Request-Id header for telemetry correlation.
+* Add operation context to request tracking middleware.
 
 ## Version 0.5.5
 * Add some basic logging when failed to send telemetry to the server
@@ -11,4 +14,4 @@ This file needs to be updated with every significant pull request. It is used to
 
 ## Version 0.5.4
 
-Changelog started after this release. 
+Changelog started after this release. 

README.md

@@ -145,3 +145,17 @@ use ApplicationInsights::Rack::TrackRequest, '<YOUR INSTRUMENTATION KEY GOES HER
 # For rails, suggest to set up this middleware in application.rb so that unhandled exceptions from controllers are also collected
 config.middleware.use 'ApplicationInsights::Rack::TrackRequest', '<YOUR INSTRUMENTATION KEY GOES HERE>', <buffer size>
 ```
+
+#### Rerieving the Request-Id value from ApplicationInsights ####
+```ruby
+# from time to time you may need to access a request's id from within your app
+application_insights_request_id = env['ApplicationInsights.request.id']
+
+# this can be used for a number of different purposes, including telemetry correlation
+uri = URI('http://api.example.com/search/?q=test')
+
+req = Net::HTTP::Get.new(uri)
+req['Request-Id'] = "#{application_insights_request_id}1" if application_insights_request_id
+
+Net::HTTP.start(uri.hostname, uri.port) { |http| http.request(req) }
+```

application_insights.gemspec

@@ -25,4 +25,6 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency 'redcarpet', '~> 3.2.2'
   spec.add_development_dependency 'rack', '>= 1.0.0'
   spec.add_development_dependency 'test-unit', '~> 3.0.8'
+  spec.add_development_dependency 'mocha', '~> 1.5.0'
+
 end

lib/application_insights/channel/contracts/operation.rb

@@ -13,7 +13,7 @@ module ApplicationInsights::Channel::Contracts
       root_id: 'ai.operation.rootId',
       synthetic_source: 'ai.operation.syntheticSource',
       is_synthetic: 'ai.operation.isSynthetic',
-      correlation_vector: "ai.operation.correlationVector"
+      correlation_vector: 'ai.operation.correlationVector'
     )
   end
 end

lib/application_insights/channel/queue_base.rb

@@ -13,8 +13,7 @@ module ApplicationInsights
       def initialize(sender)
         @queue = Queue.new
         @max_queue_length = 500
-        @sender = sender
-        @sender.queue = self if sender
+        self.sender = sender
       end
 
       # The maximum number of items that will be held by the queue before the
@@ -27,6 +26,15 @@ module ApplicationInsights
       # @return [SenderBase] the sender object.
       attr_reader :sender
 
+      # Change the sender that is associated with this queue.
+      # @param [SenderBase] sender the sender object.
+      # @return [SenderBase] the sender object.
+      def sender=(sender)
+        @sender = sender
+        @sender.queue = self if sender
+        @sender
+      end
+
       # Adds the passed in item object to the queue and calls {#flush} if the
       # size of the queue is larger than {#max_queue_length}. This method does
       # nothing if the passed in item is nil.

lib/application_insights/rack/track_request.rb

@@ -1,4 +1,5 @@
 require 'rack'
+require 'securerandom'
 require_relative '../channel/contracts/request_data'
 require_relative '../telemetry_client'
 
@@ -19,11 +20,23 @@ module ApplicationInsights
         @instrumentation_key = instrumentation_key
         @buffer_size = buffer_size
         @send_interval = send_interval
+
+        @sender = Channel::AsynchronousSender.new
+        @sender.send_interval = @send_interval
+        queue = Channel::AsynchronousQueue.new @sender
+        queue.max_queue_length = @buffer_size
+        @channel = Channel::TelemetryChannel.new nil, queue
+
+        @client = TelemetryClient.new @instrumentation_key, @channel
       end
 
       # Track requests and send data to Application Insights asynchronously.
       # @param [Hash] env the rack environment.
       def call(env)
+        # Build a request ID, incorporating one from our request if one exists.
+        request_id = request_id_header(env['HTTP_REQUEST_ID'])
+        env['ApplicationInsights.request.id'] = request_id
+
         start = Time.now
         begin
           status, headers, response = @app.call(env)
@@ -33,28 +46,17 @@ module ApplicationInsights
         end
         stop = Time.now
 
-        unless @client
-          sender = @sender || Channel::AsynchronousSender.new
-          sender.send_interval = @send_interval
-          queue = Channel::AsynchronousQueue.new sender
-          queue.max_queue_length = @buffer_size
-          channel = Channel::TelemetryChannel.new nil, queue
-
-          @client = TelemetryClient.new @instrumentation_key, channel
-        end
-
-        request = ::Rack::Request.new env
-        id = rand(16**32).to_s(16)
         start_time = start.iso8601(7)
         duration = format_request_duration(stop - start)
         success = status.to_i < 400
-        options = {
-          :name => "#{request.request_method} #{request.path}",
-          :http_method => request.request_method,
-          :url => request.url
-        }
 
-        @client.track_request id, start_time, duration, status, success, options
+        request = ::Rack::Request.new env
+        options = options_hash(request)
+
+        data = request_data(request_id, start_time, duration, status, success, options)
+        context = telemetry_context(request_id, env['HTTP_REQUEST_ID'])
+
+        @client.channel.write data, context, start_time
 
         if exception
           @client.track_exception exception, handled_at: 'Unhandled'
@@ -67,7 +69,10 @@ module ApplicationInsights
       private
 
       def sender=(sender)
-        @sender = sender if sender.is_a? Channel::AsynchronousSender
+        if sender.is_a? Channel::AsynchronousSender
+          @sender = sender
+          @client.channel.queue.sender = @sender
+        end
       end
 
       def client
@@ -82,6 +87,68 @@ module ApplicationInsights
 
         Time.at(duration_seconds).gmtime.strftime("00.%H:%M:%S.%7N")
       end
+
+      def request_id_header(request_id)
+        valid_request_id_header = valid_request_id(request_id)
+
+        length = valid_request_id_header ? 5 : 10
+        id = SecureRandom.base64(length)
+
+        if valid_request_id_header
+          request_id_has_end = %w[. _].include?(request_id[-1])
+          request_id << '.' unless request_id_has_end
+
+          return "#{request_id}#{id}_"
+        end
+
+        "|#{id}."
+      end
+
+      def valid_request_id(request_id)
+        request_id && request_id[0] == '|'
+      end
+
+      def operation_id(id)
+        # Returns the root ID from the '|' to the first '.' if any.
+        root_start = id[0] == '|' ? 1 : 0
+
+        root_end = id.index('.')
+        root_end = root_end ? root_end - 1 : id.length - root_start
+
+        id[root_start..root_end]
+      end
+
+      def options_hash(request)
+        {
+            name: "#{request.request_method} #{request.path}",
+            http_method: request.request_method,
+            url: request.url
+        }
+      end
+
+      def request_data(request_id, start_time, duration, status, success, options)
+        Channel::Contracts::RequestData.new(
+            :id => request_id || 'Null',
+            :start_time => start_time || Time.now.iso8601(7),
+            :duration => duration || '0:00:00:00.0000000',
+            :response_code => status || 200,
+            :success => success == nil ? true : success,
+            :name => options[:name],
+            :http_method => options[:http_method],
+            :url => options[:url],
+            :properties => options[:properties] || {},
+            :measurements => options[:measurements] || {}
+        )
+      end
+
+      def telemetry_context(request_id, request_id_header)
+        context = Channel::TelemetryContext.new
+        context.instrumentation_key = @instrumentation_key
+        context.operation.id = operation_id(request_id)
+        context.operation.parent_id = request_id_header
+
+        context
+      end
     end
   end
 end

lib/application_insights/version.rb

@@ -1,3 +1,3 @@
 module ApplicationInsights
-  VERSION = '0.5.6'
+  VERSION = '0.5.6'.freeze
 end

test/application_insights/rack/test_track_request.rb

@@ -1,4 +1,5 @@
 require 'test/unit'
+require 'mocha/test_unit'
 require 'rack/mock'
 require_relative '../mock_sender'
 require_relative '../../../lib/application_insights/rack/track_request'
@@ -20,8 +21,12 @@ class TestTrackRequest < Test::Unit::TestCase
     track_request = TrackRequest.new app, instrumentation_key, 500, 1
     track_request.send(:sender=, sender)
     start_time = Time.now
+
+    SecureRandom.expects(:base64).with(10).returns('y0NM2eOY/fnQPw==')
     result = track_request.call(env)
-    assert_equal app.call(env), result
+
+    app_result = app.call(env)
+    assert_equal app_result, result
     sleep(sender.send_interval)
 
     assert_equal 1, sender.buffer.count
@@ -89,20 +94,14 @@ class TestTrackRequest < Test::Unit::TestCase
     send_interval = 5
     track_request = TrackRequest.new app, 'key', buffer_size, send_interval
     client = track_request.send(:client)
-    # test lazy initialization
-    assert_nil client
+    # test client initialization
+    assert_equal ApplicationInsights::TelemetryClient, client.class
 
     track_request.call(env)
     client = track_request.send(:client)
     channel = client.channel
     assert_equal buffer_size, channel.queue.max_queue_length
     assert_equal send_interval, channel.sender.send_interval
-
-    track_request.call(env)
-    client2 = track_request.send(:client)
-    channel2 = client2.channel
-    assert_same client, client2
-    assert_same channel, channel2
   end
 
   def test_format_request_duration_less_than_a_day
@@ -139,4 +138,45 @@ class TestTrackRequest < Test::Unit::TestCase
     assert_equal 0, match['second'].to_i
     assert_equal 0, match['fraction'].to_i
   end
+
+  def test_request_id_is_generated_correctly
+    app = Proc.new {|env| [200, {"Content-Type" => "text/html"}, ["Hello Rack!"]]}
+    url = "http://localhost:8080/foo?a=b"
+    http_method = 'PUT'
+    env = Rack::MockRequest.env_for(url, :method => http_method)
+    instrumentation_key = 'key'
+    sender = MockAsynchronousSender.new
+    track_request = TrackRequest.new app, instrumentation_key, 500, 0
+    track_request.send(:sender=, sender)
+
+    # ignores ids that don't begin with | (16 chars)
+    env['HTTP_REQUEST_ID'] = 'ab456_1.ea6741a'
+    SecureRandom.expects(:base64).with(10).returns('y0NM2eOY/fnQPw==')
+    track_request.call(env)
+    assert_equal '|y0NM2eOY/fnQPw==.', env['ApplicationInsights.request.id']
+
+    # appends to ids with a dot (8 chars)
+    env['HTTP_REQUEST_ID'] = '|1234.'
+    SecureRandom.expects(:base64).with(5).returns('eXsMFHs=')
+    track_request.call(env)
+    assert_equal '|1234.eXsMFHs=_', env['ApplicationInsights.request.id']
+
+    # appends to ids with an underscore (8 chars)
+    env['HTTP_REQUEST_ID'] = '|1234_'
+    SecureRandom.expects(:base64).with(5).returns('eXsMFHs=')
+    track_request.call(env)
+    assert_equal '|1234_eXsMFHs=_', env['ApplicationInsights.request.id']
+
+    # appends a dot if neither a dot or underscore are present (8 chars)
+    env['HTTP_REQUEST_ID'] = '|ab456_1.ea6741a'
+    SecureRandom.expects(:base64).with(5).returns('eXsMFHs=')
+    track_request.call(env)
+    assert_equal '|ab456_1.ea6741a.eXsMFHs=_', env['ApplicationInsights.request.id']
+
+    # generates a stand-alone id if one is not provided (16 chars)
+    env.delete('HTTP_REQUEST_ID')
+    SecureRandom.expects(:base64).with(10).returns('y0NM2eOY/fnQPw==')
+    track_request.call(env)
+    assert_equal '|y0NM2eOY/fnQPw==.', env['ApplicationInsights.request.id']
+  end
 end