add headers reporting uncompressed size and doc count

Author: jsvd

lib/logstash/outputs/elasticsearch/http_client.rb

@@ -21,7 +21,8 @@ module LogStash; module Outputs; class ElasticSearch;
   # We wound up agreeing that a number greater than 10 MiB and less than 100MiB
   # made sense. We picked one on the lowish side to not use too much heap.
   TARGET_BULK_BYTES = 20 * 1024 * 1024 # 20MiB
-
+  EVENT_COUNT_HEADER = "X-Elastic-Event-Count".freeze
+  UNCOMPRESSED_LENGTH_HEADER = "X-Elastic-Uncompressed-Request-Length".freeze
 
   class HttpClient
     attr_reader :client, :options, :logger, :pool, :action_count, :recv_count
@@ -143,7 +144,11 @@ def bulk(actions)
                        :payload_size => stream_writer.pos,
                        :content_length => body_stream.size,
                        :batch_offset => (index + 1 - batch_actions.size))
-          bulk_responses << bulk_send(body_stream, batch_actions)
+          headers = {
+            EVENT_COUNT_HEADER => batch_actions.size,
+            UNCOMPRESSED_LENGTH_HEADER => stream_writer.pos
+          }
+          bulk_responses << bulk_send(body_stream, batch_actions, headers)
           body_stream.truncate(0) && body_stream.seek(0)
           stream_writer = gzip_writer(body_stream) if compression_level?
           batch_actions.clear
@@ -159,7 +164,14 @@ def bulk(actions)
                    :payload_size => stream_writer.pos,
                    :content_length => body_stream.size,
                    :batch_offset => (actions.size - batch_actions.size))
-      bulk_responses << bulk_send(body_stream, batch_actions) if body_stream.size > 0
+
+      if body_stream.size > 0
+        headers = {
+          EVENT_COUNT_HEADER => batch_actions.size,
+          UNCOMPRESSED_LENGTH_HEADER => stream_writer.pos
+        }
+        bulk_responses << bulk_send(body_stream, batch_actions, headers)
+      end
 
       body_stream.close unless compression_level?
       join_bulk_responses(bulk_responses)
@@ -179,8 +191,8 @@ def join_bulk_responses(bulk_responses)
       }
     end
 
-    def bulk_send(body_stream, batch_actions)
-      params = compression_level? ? {:headers => {"Content-Encoding" => "gzip"}} : {}
+    def bulk_send(body_stream, batch_actions, headers = {})
+      params = compression_level? ? {:headers => headers.merge("Content-Encoding" => "gzip") } : { :headers => headers }
 
       begin
         response = @pool.post(@bulk_path, params, body_stream.string)

spec/unit/outputs/elasticsearch/http_client_spec.rb

@@ -270,6 +270,85 @@
 
        end
      end
+    context "the 'user-agent' header" do
+      let(:pool) { double("pool") }
+      let(:compression_level) { 6 }
+      let(:base_options) { super().merge( :client_settings => {:compression_level => compression_level}) }
+      let(:actions) { [
+        ["index", {:_id=>nil, :_index=>"logstash"}, {"message_1"=> message_1}],
+        ["index", {:_id=>nil, :_index=>"logstash"}, {"message_2"=> message_2}],
+        ["index", {:_id=>nil, :_index=>"logstash"}, {"message_3"=> message_3}],
+      ]}
+      let(:message_1) { "hello" }
+      let(:message_2_size) { 1_000 }
+      let(:message_2) { SecureRandom.alphanumeric(message_2_size / 2 ) * 2 }
+      let(:message_3_size) { 1_000 }
+      let(:message_3) { "m" * message_3_size }
+      let(:messages_size) { message_1.size + message_2.size + message_3.size }
+      let(:action_overhead) { 42 + 16 + 2 } # header plus doc key size plus new line overhead per action
+
+      let(:response) do
+        response = double("response")
+        allow(response).to receive(:code).and_return(response)
+        allow(response).to receive(:body).and_return({"errors" => false}.to_json)
+        response
+      end
+
+      before(:each) do
+        subject.instance_variable_set("@pool", pool)
+      end
+
+      it "carries bulk request's uncompressed size" do
+        expect(pool).to receive(:post) do |path, params, body|
+          headers = params.fetch(:headers, {})
+          expect(headers["X-Elastic-Event-Count"]).to be(3)
+          expect(headers["X-Elastic-Uncompressed-Request-Length"]).to eq messages_size + (action_overhead * 3)
+          expect(headers["X-Elastic-Uncompressed-Request-Length"]).to be > body.size
+        end.and_return(response)
+
+        subject.send(:bulk, actions)
+      end
+      context "without compression" do
+        let(:compression_level) { 0 }
+        it "carries bulk request's uncompressed size" do
+          expect(pool).to receive(:post) do |path, params, body|
+            headers = params.fetch(:headers, {})
+            expect(headers["X-Elastic-Event-Count"]).to be(3)
+            expect(headers["X-Elastic-Uncompressed-Request-Length"]).to eq messages_size + (action_overhead * 3)
+            expect(headers["X-Elastic-Uncompressed-Request-Length"]).to eq body.size
+          end.and_return(response)
+          subject.send(:bulk, actions)
+        end
+      end
+
+      context "with compressed messages over 20MB" do
+        let(:message_2_size) { 21_000_000 }
+        it "carries bulk request's uncompressed size" do
+          # only the first, tiny, message is sent first
+          expect(pool).to receive(:post) do |path, params, body|
+            headers = params.fetch(:headers, {})
+            expect(headers["X-Elastic-Uncompressed-Request-Length"]).to be == message_1.size + action_overhead
+            expect(headers["X-Elastic-Event-Count"]).to be(1)
+          end.and_return(response)
+
+          # huge message_2 is sent afterwards alone
+          expect(pool).to receive(:post) do |path, params, body|
+            headers = params.fetch(:headers, {})
+            expect(headers["X-Elastic-Uncompressed-Request-Length"]).to be >= message_2.size + action_overhead
+            expect(headers["X-Elastic-Event-Count"]).to be(1)
+          end.and_return(response)
+
+          # finally medium message_3 is sent alone as well
+          expect(pool).to receive(:post) do |path, params, body|
+            headers = params.fetch(:headers, {})
+            expect(headers["X-Elastic-Uncompressed-Request-Length"]).to be >= message_3.size + action_overhead
+            expect(headers["X-Elastic-Event-Count"]).to be(1)
+          end.and_return(response)
+
+          subject.send(:bulk, actions)
+        end
+      end
+    end
   end
 
   describe "sniffing" do