logstash-output-loki 1.0.0 → 1.0.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9dd08c392ec88148a6f410e159169b3a19f877a25174fa3f43b2fdeb7fe16128
-  data.tar.gz: 34aaf03576d8fbbb97aab02013113254734984ee1744c2ab4e8438bd51657569
+  metadata.gz: 5d3ae7289992fc6a6b98db3f5f1d929041656f2cd1e7e53915d48df9a8bf216b
+  data.tar.gz: 21e0b37676d7f6d0bca071bd2644a8ce96370e1c62efbf26becb3bce1de63f53
 SHA512:
-  metadata.gz: 3520ab5618092976cfe771ad4f9ddea4aca3b96efe75d322c2f505e7fd6e642f3cee4775eb1ef0514c9db360a3c36b138e9f1426f6de9b625945cafdafb0b27d
-  data.tar.gz: 6c7bb11d79fe0828e629bc27771a4f63584ea76af2ddb69eac950e5c7e3b115b62b674be28e8a828c97d8d3d851406748f0c83d49d9147080f0cc593ad1a076f
+  metadata.gz: 74e562693367ffcae0a7c639185ca75a984a156a7b7d8c7c61626b17ea01e20470901a2847230aff5f89d7b800f14cdc35ed82be93dd865af91d54548332dbf2
+  data.tar.gz: 9caf43738978ad056df2661460452c7059bd8a4c646f6e9949b3c7ee492b57ec1d0fa9320df02748f928fe4b002404ec2cfa75c0e9f4d78465606d38088dc8d7

data/Gemfile

@@ -2,10 +2,13 @@ source 'https://rubygems.org'
 
 gemspec
 
-logstash_path = ENV["LOGSTASH_PATH"] || "logstash-libs"
-use_logstash_source = ENV["LOGSTASH_SOURCE"] && ENV["LOGSTASH_SOURCE"].to_s == "1"
+logstash_path = ENV["LOGSTASH_PATH"] || "./logstash"
 
-if Dir.exist?(logstash_path) && use_logstash_source
+if Dir.exist?(logstash_path)
   gem 'logstash-core', :path => "#{logstash_path}/logstash-core"
   gem 'logstash-core-plugin-api', :path => "#{logstash_path}/logstash-core-plugin-api"
+else
+  raise 'missing logstash vendoring'
 end
+
+gem "webmock", "~> 3.8"

data/README.md

@@ -1,6 +1,6 @@
-# Loki Logstash Output Plugin
+# Contributing to Loki Logstash Output Plugin
 
-Logstash plugin to send logstash aggregated logs to Loki.
+For information about hwo to use this plugin see this [documentation](../../docs/clients/logstash/README.md).
 
 ## Install dependencies
 
@@ -41,11 +41,11 @@ git clone git@github.com:elastic/logstash.git
 cd logstash
 git checkout tags/v7.6.2
 export LOGSTASH_PATH=`pwd`
-export LOGSTASH_SOURCE="1"
-export GEM_PATH=$LOGSTASH_PATH/vendor/bundle/
-export GEM_HOME=$LOGSTASH_PATH/vendor/bundle/
+export GEM_PATH=$LOGSTASH_PATH/vendor/bundle/jruby/2.5.0
+export GEM_HOME=$LOGSTASH_PATH/vendor/bundle/jruby/2.5.0
+./gradlew assemble
 cd ..
-ruby -S bundle install --path
+ruby -S bundle install
 ruby -S bundle exec rake vendor
 ```
 
@@ -55,7 +55,15 @@ ruby -S bundle exec rake vendor
 
 ### Test
 
-`bundle exec rspec`
+`ruby -S bundle exec rspec`
+
+Alternatively if you don't want to install JRuby. Enter inside logstash-loki container.
+
+```bash
+docker build -t logstash-loki ./
+docker run -v  `pwd`/spec:/home/logstash/spec -it --rm --entrypoint /bin/sh logstash-loki
+bundle exec rspec
+```
 
 ## Install plugin to local logstash
 

data/lib/logstash/outputs/loki.rb

@@ -1,5 +1,7 @@
 # encoding: utf-8
 require "logstash/outputs/base"
+require "logstash/outputs/loki/entry"
+require "logstash/outputs/loki/batch"
 require "logstash/namespace"
 require 'net/http'
 require 'concurrent-edge'
@@ -8,9 +10,7 @@ require 'uri'
 require 'json'
 
 class LogStash::Outputs::Loki < LogStash::Outputs::Base
-  require 'logstash/outputs/loki/batch'
-  require 'logstash/outputs/loki/entry'
-
+  include Loki
   config_name "loki"
 
   ## 'A single instance of the Output will be shared among the pipeline worker threads'
@@ -39,12 +39,6 @@ class LogStash::Outputs::Loki < LogStash::Outputs::Base
   ## 'Interval in seconds to wait before pushing a batch of records to loki. Defaults to 1 second'
   config :batch_wait, :validate => :number, :default => 1, :required => false
 
-  ## 'Array of label names to include in all logstreams'
-  config :include_labels, :validate => :array, :default => [], :required => true
-
-  ## 'Extra labels to add to all log streams'
-  config :external_labels, :validate => :hash,  :default => {}, :required => false
-
   ## 'Log line field to pick from logstash. Defaults to "message"'
   config :message_field, :validate => :string, :default => "message", :required => false
 
@@ -57,6 +51,7 @@ class LogStash::Outputs::Loki < LogStash::Outputs::Base
   ## 'Backoff configuration. Maximum number of retries to do'
   config :retries, :validate => :number, :default => 10, :required => false
 
+  attr_reader :batch
   public
   def register
     @uri = URI.parse(@url)
@@ -64,22 +59,16 @@ class LogStash::Outputs::Loki < LogStash::Outputs::Base
       raise LogStash::ConfigurationError, "url parameter must be valid HTTP, currently '#{@url}'"
     end
 
-    if @include_labels.empty?
-      raise LogStash::ConfigurationError, "include_labels should contain atleast one label, currently '#{@include_labels}'"
-    end
-
     if @min_delay > @max_delay
       raise LogStash::ConfigurationError, "Min delay should be less than Max delay, currently 'Min delay is #{@min_delay} and Max delay is #{@max_delay}'"
     end
 
     @logger.info("Loki output plugin", :class => self.class.name)
 
-    # intialize channels
+    # initialize channels
     @Channel = Concurrent::Channel
     @entries = @Channel.new
-
-    # excluded message and timestamp from labels
-    @exclude_labels = ["message", "@timestamp"]
+    @stop = @Channel.new
 
     # create nil batch object.
     @batch = nil
@@ -130,8 +119,8 @@ class LogStash::Outputs::Loki < LogStash::Outputs::Base
   end
 
   def run()
-	  min_wait_checkfrequency = 1/1000 #1 millisecond
-	  max_wait_checkfrequency = @batch_wait
+	  min_wait_checkfrequency = 1/100 #1 millisecond
+	  max_wait_checkfrequency = @batch_wait / 10
 	  if max_wait_checkfrequency < min_wait_checkfrequency
 		  max_wait_checkfrequency = min_wait_checkfrequency
     end
@@ -139,30 +128,21 @@ class LogStash::Outputs::Loki < LogStash::Outputs::Base
     @max_wait_check = Concurrent::Channel.tick(max_wait_checkfrequency)
     loop do
       Concurrent::Channel.select do |s|
+        s.take(@stop) {
+          return
+        }
         s.take(@entries) { |e|
-          if @batch.nil?
-            @batch = Batch.new(e)
-            next
-          end
-
-          line = e.entry['line']
-          if @batch.size_bytes_after(line) > @batch_size
-            @logger.debug("Max batch_size is reached. Sending batch to loki")
-            send(@tenant_id, @batch)
-            @batch = Batch.new(e)
-            next
-          end
-          @batch.add(e)
+         if !add_entry_to_batch(e)
+           @logger.debug("Max batch_size is reached. Sending batch to loki")
+           send(@batch)
+           @batch = Batch.new(e)
+         end
         }
         s.take(@max_wait_check) {
           # Send batch if max wait time has been reached
-          if !@batch.nil?
-            if @batch.age() < @batch_wait
-              next
-            end
-
+          if is_batch_expired
             @logger.debug("Max batch_wait time is reached. Sending batch to loki")
-            send(@tenant_id, @batch)
+            send(@batch)
             @batch = nil
           end
         }
@@ -170,85 +150,62 @@ class LogStash::Outputs::Loki < LogStash::Outputs::Base
     end
   end
 
-  ## Receives logstash events
-  public
-  def receive(event)
-    labels = {}
-    event_hash = event.to_hash
-    lbls = handle_labels(event_hash, labels, "")
+  # add an entry to the current batch return false if the batch is full
+  # and the entry can't be added.
+  def add_entry_to_batch(e)
+    line = e.entry['line']
+    # we don't want to send empty lines.
+    return true if line.to_s.strip.empty?
 
-    data_labels, entry_hash = build_entry(lbls, event)
-    @entries << Entry.new(data_labels, entry_hash)
+    if @batch.nil?
+      @batch = Batch.new(e)
+      return true
+    end
 
+    if @batch.size_bytes_after(line) > @batch_size
+      return false
+    end
+    @batch.add(e)
+    return true
   end
 
-  def close
-    @logger.info("Closing loki output plugin. Flushing all pending batches")
-    send(@tenant_id, @batch) if !@batch.nil?
-    @entries.close
-    @max_wait_check.close if !@max_wait_check.nil?
+  def is_batch_expired
+    return !@batch.nil? && @batch.age() >= @batch_wait
   end
 
-  def build_entry(lbls, event)
-    labels = lbls.merge(@external_labels)
-    entry_hash = {
-      "ts" => event.get("@timestamp").to_i * (10**9),
-      "line" => event.get(@message_field).to_s
-    }
-    return labels, entry_hash
+  ## Receives logstash events
+  public
+  def receive(event)
+    @entries << Entry.new(event, @message_field)
   end
 
-  def handle_labels(event_hash, labels, parent_key)
-    event_hash.each{ |key,value|
-      if !@exclude_labels.include?(key)
-        if value.is_a?(Hash)
-          if parent_key != ""
-            handle_labels(value, labels, parent_key + "_" + key)
-          else
-            handle_labels(value, labels, key)
-          end
-        else
-          if parent_key != ""
-            labels[parent_key + "_" + key] = value.to_s
-          else
-            labels[key] = value.to_s
-          end
-        end
-      end
-    }
-    return extract_labels(labels)
-  end
+  def close
+    @entries.close
+    @max_wait_check.close if !@max_wait_check.nil?
+    @stop << true # stop will block until it's accepted by the worker.
 
-  def extract_labels(extracted_labels)
-    labels = {}
-    extracted_labels.each { |key, value|
-      if @include_labels.include?(key)
-        key = key.gsub("@", '')
-        labels[key] = value
-      end
-    }
-    return labels
+    # if by any chance we still have a forming batch, we need to send it.
+    send(@batch) if !@batch.nil?
+    @batch = nil
   end
 
-  def send(tenant_id, batch)
-    payload = build_payload(batch)
-    res = loki_http_request(tenant_id, payload, @min_delay, @max_delay, @retries)
-
+  def send(batch)
+    payload = batch.to_json
+    res = loki_http_request(payload)
     if res.is_a?(Net::HTTPSuccess)
       @logger.debug("Successfully pushed data to loki")
-      return
     else
-      @logger.error("failed to write post to ", :uri => @uri, :code => res.code, :body => res.body, :message => res.message) if !res.nil?
-      @logger.debug("Payload object ", :payload => payload)
+      @logger.debug("failed payload", :payload => payload)
     end
   end
 
-  def loki_http_request(tenant_id, payload, min_delay, max_delay, retries)
+  def loki_http_request(payload)
     req = Net::HTTP::Post.new(
       @uri.request_uri
     )
     req.add_field('Content-Type', 'application/json')
-    req.add_field('X-Scope-OrgID', tenant_id) if tenant_id
+    req.add_field('X-Scope-OrgID', @tenant_id) if @tenant_id
+    req['User-Agent']= 'loki-logstash'
     req.basic_auth(@username, @password) if @username
     req.body = payload
 
@@ -256,53 +213,28 @@ class LogStash::Outputs::Loki < LogStash::Outputs::Base
 
     @logger.debug("sending #{req.body.length} bytes to loki")
     retry_count = 0
-    delay = min_delay
+    delay = @min_delay
     begin
-      res = Net::HTTP.start(@uri.host, @uri.port, **opts) { |http| http.request(req) }
-    rescue Net::HTTPTooManyRequests, Net::HTTPServerError, Errno::ECONNREFUSED => e
-      unless retry_count < retries
-        @logger.error("Error while sending data to loki. Tried #{retry_count} times\n. :error => #{e}")
-        return res
-      end
-
+      res = Net::HTTP.start(@uri.host, @uri.port, **opts) { |http|
+        http.request(req)
+      }
+      return res if !res.nil? && res.code.to_i != 429 && res.code.to_i.div(100) != 5
+      raise StandardError.new res
+    rescue StandardError => e
       retry_count += 1
-      @logger.warn("Trying to send again. Attempt number: #{retry_count}. Retrying in #{delay}s")
-      sleep delay
-
-      if (delay * 2 - delay) > max_delay
-        delay = delay
+      @logger.warn("Failed to send batch attempt: #{retry_count}/#{@retries}", :error_inspect => e.inspect, :error => e)
+      if retry_count < @retries
+        sleep delay
+        if (delay * 2 - delay) > @max_delay
+          delay = delay
+        else
+          delay = delay * 2
+        end
+        retry
       else
-        delay = delay * 2
+        @logger.error("Failed to send batch", :error_inspect => e.inspect, :error => e)
+        return res
       end
-
-      retry
-    rescue StandardError => e
-      @logger.error("Error while connecting to loki server ", :error_inspect => e.inspect, :error => e)
-      return res
     end
-    return res
-  end
-
-  def build_payload(batch)
-    payload = {}
-    payload['streams'] = []
-    batch.streams.each { |labels, stream|
-        stream_obj = get_stream_obj(stream)
-        payload['streams'].push(stream_obj)
-    }
-    return payload.to_json
-  end
-
-  def get_stream_obj(stream)
-    stream_obj = {}
-    stream_obj['stream'] = stream['labels']
-    stream_obj['values'] = []
-    values = []
-    stream['entries'].each { |entry|
-        values.push(entry['ts'].to_s)
-        values.push(entry['line'])
-    }
-    stream_obj['values'].push(values)
-    return stream_obj
   end
 end

data/lib/logstash/outputs/loki/batch.rb

@@ -1,47 +1,63 @@
 require 'time'
 
-module LogStash
-    module Outputs
-        class Loki 
-            class Batch
-                attr_reader :streams
-                def initialize(e)
-                    @bytes = 0
-                    @createdAt = Time.now
-                    @streams = {}
-                    add(e)
-                end 
-
-                def size_bytes()
-                    return @bytes
-                end
-
-                def add(e)
-                    @bytes = @bytes + e.entry['line'].length
-
-                    # Append the entry to an already existing stream (if any)
-                    labels = e.labels.to_s
-                    if @streams.has_key?(labels)
-                        stream = @streams[labels]
-                        stream['entries'] = stream['entries'] + e.entry
-                        return
-                    else  
-                        # Add the entry as a new stream
-                        @streams[labels] = {
-                            "labels" => e.labels,
-                            "entries" => [e.entry],
-                        }
-                    end
-                end 
-
-                def size_bytes_after(line)
-                    return @bytes + line.length
-                end 
-
-                def age()
-                    return Time.now - @createdAt
-                end
-            end 
-        end 
-    end 
-end 
+module Loki
+    class Batch
+        attr_reader :streams
+        def initialize(e)
+            @bytes = 0
+            @createdAt = Time.now
+            @streams = {}
+            add(e)
+        end
+
+        def size_bytes
+            return @bytes
+        end
+
+        def add(e)
+            @bytes = @bytes + e.entry['line'].length
+
+            # Append the entry to an already existing stream (if any)
+            labels = e.labels.sort.to_h
+            labelkey = labels.to_s
+            if @streams.has_key?(labelkey)
+                stream = @streams[labelkey]
+                stream['entries'].append(e.entry)
+                return
+            else
+                # Add the entry as a new stream
+                @streams[labelkey] = {
+                    "labels" => labels,
+                    "entries" => [e.entry],
+                }
+            end
+        end
+
+        def size_bytes_after(line)
+            return @bytes + line.length
+        end
+
+        def age()
+            return Time.now - @createdAt
+        end
+
+        def to_json
+            streams = []
+            @streams.each { |_ , stream|
+                streams.append(build_stream(stream))
+            }
+            return {"streams"=>streams}.to_json
+        end
+
+        def build_stream(stream)
+            values = []
+            stream['entries'].each { |entry|
+                values.append([entry['ts'].to_s, entry['line']])
+            }
+            return {
+                'stream'=>stream['labels'],
+                'values' => values
+            }
+        end
+    end
+end

data/lib/logstash/outputs/loki/entry.rb

@@ -1,13 +1,25 @@
-module LogStash
-    module Outputs
-        class Loki 
-            class Entry
-                attr_reader :labels, :entry
-                def initialize(labels, entry)
-                    @labels   = labels
-                    @entry = entry
-                end 
-            end 
-        end 
-    end 
-end 
+module Loki
+    def to_ns(s)
+        (s.to_f * (10**9)).to_i
+    end
+    class Entry
+        include Loki
+        attr_reader :labels, :entry
+        def initialize(event,message_field)
+            @entry = {
+                "ts" => to_ns(event.get("@timestamp")),
+                "line" => event.get(message_field).to_s
+            }
+            event = event.clone()
+            event.remove(message_field)
+            event.remove("@timestamp")
+
+            @labels = {}
+            event.to_hash.each { |key,value|
+                next if key.start_with?('@')
+                next if value.is_a?(Hash)
+                @labels[key] = value.to_s
+            }
+        end
+    end
+end

data/logstash-output-loki.gemspec

@@ -1,8 +1,8 @@
 Gem::Specification.new do |s|
   s.name = 'logstash-output-loki'
-  s.version         = '1.0.0'
-  s.authors = ['Aditya C S']
-  s.email   = ['aditya.gnu@gmail.com']
+  s.version         = '1.0.1'
+  s.authors = ['Aditya C S','Cyril Tovena']
+  s.email   = ['aditya.gnu@gmail.com','cyril.tovena@grafana.com']
 
   s.summary       = 'Output plugin to ship logs to a Grafana Loki server'
   s.description   = 'Output plugin to ship logs to a Grafana Loki server'
@@ -11,7 +11,7 @@ Gem::Specification.new do |s|
   s.require_paths = ["lib"]
 
   # Files
-  s.files = Dir['lib/**/*','spec/**/*','vendor/**/*','*.gemspec','*.md','CONTRIBUTORS','Gemfile','LICENSE','NOTICE.TXT']
+  s.files = Dir['lib/**/*','spec/**/*','vendor/**/*','*.gemspec','*.md','CONTRIBUTORS','Gemfile']
    # Tests
   s.test_files = s.files.grep(%r{^(test|spec|features)/})
 

data/spec/outputs/loki/entry_spec.rb

@@ -0,0 +1,59 @@
+# encoding: utf-8
+require "logstash/devutils/rspec/spec_helper"
+require "logstash/outputs/loki"
+require "logstash/codecs/plain"
+require "logstash/event"
+require "net/http"
+include Loki
+
+describe Loki::Entry do
+  context 'test entry generation' do
+      let (:event) {
+        LogStash::Event.new(
+            {
+              'message' => 'hello',
+              '@metadata' => {'foo'=>'bar'},
+              '@version' => '1',
+              'foo' => 5,
+              'agent' => 'filebeat',
+              'log' => {
+                   'file' =>
+                        {'@path' => '/path/to/file.log'},
+                   },
+              'host' => '172.0.0.1',
+              '@timestamp' => Time.now
+            }
+        )
+      }
+
+      it 'labels extracted should not contains object and metadata or timestamp' do
+        entry = Entry.new(event,"message")
+        expect(entry.labels).to eql({ 'agent' => 'filebeat', 'host' => '172.0.0.1', 'foo'=>'5'})
+        expect(entry.entry['ts']).to eql to_ns(event.get("@timestamp"))
+        expect(entry.entry['line']).to eql 'hello'
+      end
+  end
+
+  context 'test batch generation with label order' do
+    let (:entries)  {[
+        Entry.new(LogStash::Event.new({"message"=>"foobuzz","buzz"=>"bar","cluster"=>"us-central1","@timestamp"=>Time.at(1)}),"message"),
+        Entry.new(LogStash::Event.new({"log"=>"foobar","bar"=>"bar","@timestamp"=>Time.at(2)}),"log"),
+        Entry.new(LogStash::Event.new({"cluster"=>"us-central1","message"=>"foobuzz","buzz"=>"bar","@timestamp"=>Time.at(3)}),"message"),
+
+    ]}
+    let (:expected) {
+      {"streams" => [
+          {"stream"=> {"buzz"=>"bar","cluster"=>"us-central1"}, "values" => [[to_ns(Time.at(1)).to_s,"foobuzz"],[to_ns(Time.at(3)).to_s,"foobuzz"]]},
+          {"stream"=> {"bar"=>"bar"}, "values"=>[[to_ns(Time.at(2)).to_s,"foobar"]]},
+      ] }
+    }
+
+    it 'to_json' do
+      @batch = Loki::Batch.new(entries.first)
+      entries.drop(1).each { |e| @batch.add(e)}
+      expect(JSON.parse(@batch.to_json)).to eql expected
+    end
+  end
+
+
+end

data/spec/outputs/loki_spec.rb

@@ -4,9 +4,12 @@ require "logstash/outputs/loki"
 require "logstash/codecs/plain"
 require "logstash/event"
 require "net/http"
+require 'webmock/rspec'
+include Loki
 
 describe LogStash::Outputs::Loki do
-  let (:simple_loki_config) {{'url' => 'http://localhost:3100', 'include_labels' => ["test_key", "other_key"], 'external_labels' => {"test" => "value"}}}
+
+  let (:simple_loki_config) { {'url' => 'http://localhost:3100'} }
 
   context 'when initializing' do
     it "should register" do
@@ -14,100 +17,231 @@ describe LogStash::Outputs::Loki do
       expect { loki.register }.to_not raise_error
     end
 
-    it 'should populate loki config with default or intialized values' do
+    it 'should populate loki config with default or initialized values' do
       loki = LogStash::Outputs::Loki.new(simple_loki_config)
       expect(loki.url).to eql 'http://localhost:3100'
       expect(loki.tenant_id).to eql nil
       expect(loki.batch_size).to eql 102400
       expect(loki.batch_wait).to eql 1
-      expect(loki.include_labels).to eql ["test_key", "other_key"]
-      expect(loki.external_labels).to include("test" => "value")
     end
   end
 
-  context 'test labels' do
-      let (:simple_loki_config) {{'url' => 'http://localhost:3100', 'include_labels' => ["@version", 'log_file_@path', 'host']}}
-      let (:event) { LogStash::Event.new({'message' => 'hello', '@version' => '1', 'agent' => 'filebeat', 'log' => {'file' => {'@path' => '/path/to/file.log'}}, 'host' => '172.0.0.1',
-                                      '@timestamp' => LogStash::Timestamp.now}) }
-      let(:loki) { LogStash::Plugin.lookup("output", "loki").new(simple_loki_config) }
+  context 'when adding en entry to the batch' do
+    let (:simple_loki_config) {{'url' => 'http://localhost:3100'}}
+    let (:entry) {Entry.new(LogStash::Event.new({"message"=>"foobuzz","buzz"=>"bar","cluster"=>"us-central1","@timestamp"=>Time.at(1)}),"message")}
+    let (:lbs) { {"buzz"=>"bar","cluster"=>"us-central1"}.sort.to_h}
 
-      before do
-        loki.register
-        loki.close
-      end
+    it 'should not add empty line' do
+      plugin = LogStash::Plugin.lookup("output", "loki").new(simple_loki_config)
+      emptyEntry = Entry.new(LogStash::Event.new({"message"=>"foobuzz","buzz"=>"bar","cluster"=>"us-central1","@timestamp"=>Time.at(1)}),"foo")
+      expect(plugin.add_entry_to_batch(emptyEntry)).to eql true
+      expect(plugin.batch).to eql nil
+    end
 
-      it 'labels extracted should have only included labels' do
-        labels = {}
-        event_hash = event.to_hash
-        expected_labels = {"version" => "1", "host" => "172.0.0.1", "log_file_path" => '/path/to/file.log'}
-        expect(loki.handle_labels(event_hash, labels, "")).to eql expected_labels
-      end
+    it 'should add entry' do
+      plugin = LogStash::Plugin.lookup("output", "loki").new(simple_loki_config)
+      expect(plugin.batch).to eql nil
+      expect(plugin.add_entry_to_batch(entry)).to eql true
+      expect(plugin.add_entry_to_batch(entry)).to eql true
+      expect(plugin.batch).not_to be_nil
+      expect(plugin.batch.streams.length).to eq 1
+      expect(plugin.batch.streams[lbs.to_s]['entries'].length).to eq 2
+      expect(plugin.batch.streams[lbs.to_s]['labels']).to eq lbs
+      expect(plugin.batch.size_bytes).to eq 14
+    end
+
+    it 'should not add if full' do
+      plugin = LogStash::Plugin.lookup("output", "loki").new(simple_loki_config.merge!({'batch_size'=>10}))
+      expect(plugin.batch).to eql nil
+      expect(plugin.add_entry_to_batch(entry)).to eql true # first entry is fine.
+      expect(plugin.batch).not_to be_nil
+      expect(plugin.batch.streams.length).to eq 1
+      expect(plugin.batch.streams[lbs.to_s]['entries'].length).to eq 1
+      expect(plugin.batch.streams[lbs.to_s]['labels']).to eq lbs
+      expect(plugin.batch.size_bytes).to eq 7
+      expect(plugin.add_entry_to_batch(entry)).to eql false # second entry goes over the limit.
+      expect(plugin.batch).not_to be_nil
+      expect(plugin.batch.streams.length).to eq 1
+      expect(plugin.batch.streams[lbs.to_s]['entries'].length).to eq 1
+      expect(plugin.batch.streams[lbs.to_s]['labels']).to eq lbs
+      expect(plugin.batch.size_bytes).to eq 7
+    end
   end
 
-  context 'validate entries' do
-    let(:timestamp) {LogStash::Timestamp.now}
-    let (:simple_loki_config) {{'url' => 'http://localhost:3100', 'include_labels' => ["version", "host", "test"], 'external_labels' => {"test" => "value"}}}
-    let (:event) { LogStash::Event.new({'message' => 'hello', '@version' => '1', 'agent' => 'filebeat', 'host' => '172.0.0.1',
-                                    '@timestamp' => timestamp}) }
-    let(:loki) { LogStash::Plugin.lookup("output", "loki").new(simple_loki_config) }
+  context 'batch expiration' do
+    let (:entry) {Entry.new(LogStash::Event.new({"message"=>"foobuzz","buzz"=>"bar","cluster"=>"us-central1","@timestamp"=>Time.at(1)}),"message")}
 
-    before do
-      loki.register
-      loki.close
+    it 'should not expire if empty' do
+      loki = LogStash::Outputs::Loki.new(simple_loki_config.merge!({'batch_wait'=>0.5}))
+      sleep(1)
+      expect(loki.is_batch_expired).to be false
     end
-
-    it 'validate expected entries are added to entries stream' do
-      labels = {"version" => "1", "host" => "172.0.0.1"}
-      expected_labels = {"version" => "1", "host" => "172.0.0.1", "test" => "value"}
-      expected_entry_hash = {
-        "ts" => timestamp.to_i * (10**9),
-        "line" => "hello".to_s
-      }
-      expected_labels_and_entry_hash = [{"version" => "1", "host" => "172.0.0.1", "test" => "value"}, expected_entry_hash]
-      expect(loki.build_entry(labels, event)).to eq(expected_labels_and_entry_hash)
+    it 'should not expire batch if not old' do
+      loki = LogStash::Outputs::Loki.new(simple_loki_config.merge!({'batch_wait'=>0.5}))
+      expect(loki.add_entry_to_batch(entry)).to eql true
+      expect(loki.is_batch_expired).to be false
+    end
+    it 'should expire if old' do
+      loki = LogStash::Outputs::Loki.new(simple_loki_config.merge!({'batch_wait'=>0.5}))
+      expect(loki.add_entry_to_batch(entry)).to eql true
+      sleep(1)
+      expect(loki.is_batch_expired).to be true
     end
   end
 
-  context 'test http requests' do
-    let (:simple_loki_config) {{'url' => 'http://localhost:3100', 'include_labels' => ["@version", "host", "test"],}}
-    let (:event) { LogStash::Event.new({'message' => 'hello', '@version' => '1', 'host' => '172.0.0.1',
-                                      '@timestamp' => LogStash::Timestamp.now}) }
-    let(:loki) { LogStash::Plugin.lookup("output", "loki").new(simple_loki_config) }
+  context 'channel' do
+    let (:event) {LogStash::Event.new({"message"=>"foobuzz","buzz"=>"bar","cluster"=>"us-central1","@timestamp"=>Time.at(1)})}
 
-    before do
+    it 'should send entry if batch size reached with no tenant' do
+      loki = LogStash::Outputs::Loki.new(simple_loki_config.merge!({'batch_wait'=>0.5,'batch_size'=>10}))
       loki.register
+      sent = Concurrent::Channel.new(capacity: 3)
+      allow(loki).to receive(:send) do |batch|
+        Thread.new do
+          sent << batch
+        end
+      end
+      loki.receive(event)
+      loki.receive(event)
       loki.close
+      ~sent
+      ~sent
     end
+    it 'should send entry while closing' do
+      loki = LogStash::Outputs::Loki.new(simple_loki_config.merge!({'batch_wait'=>10,'batch_size'=>10}))
+      loki.register
+      sent = Concurrent::Channel.new(capacity: 3)
+      allow(loki).to receive(:send) do | batch|
+        Thread.new  do
+          sent << batch
+        end
+      end
+      loki.receive(event)
+      loki.close
+      ~sent
+    end
+    it 'should send entry when batch is expiring' do
+      loki = LogStash::Outputs::Loki.new(simple_loki_config.merge!({'batch_wait'=>0.5,'batch_size'=>10}))
+      loki.register
+      sent = Concurrent::Channel.new(capacity: 3)
+      allow(loki).to receive(:send) do | batch|
+        Thread.new  do
+          sent << batch
+        end
+      end
+      loki.receive(event)
+      ~sent
+      expect(loki.batch).to be_nil
+      loki.close
+    end
+  end
 
-    it 'test http requests and raise_error when requests are not successful' do
-      labels = {}
-      event_hash = event.to_hash
-      lbls = loki.handle_labels(event_hash, labels, "")
-      entry_hash = {
-        "ts" => event.get("@timestamp").to_i * (10**9),
-        "line" => event.get("message").to_s
-      }
-      e = LogStash::Outputs::Loki::Entry.new(lbls, entry_hash)
-      batch = LogStash::Outputs::Loki::Batch.new(e)
-      payload = loki.build_payload(batch)
-
-      # response should be nil on connection error
-      expect(loki.loki_http_request("fake", payload, 1, 2, 3)).to eql nil
-
-      success = Net::HTTPSuccess.new(1.0, 200, 'OK')
-      allow(loki).to receive(:loki_http_request) { success }
-      allow(success).to receive(:payload).and_return('fake body')
-      expect(loki.loki_http_request("fake", batch, 1, 300, 10).class).to eql Net::HTTPSuccess
+  context 'http requests' do
+    let (:entry) {Entry.new(LogStash::Event.new({"message"=>"foobuzz","buzz"=>"bar","cluster"=>"us-central1","@timestamp"=>Time.at(1)}),"message")}
 
-      too_many_requests = Net::HTTPTooManyRequests.new(1.0, 429, 'OK')
-      allow(loki).to receive(:loki_http_request) { too_many_requests }
-      allow(too_many_requests).to receive(:payload).and_return('fake body')
-      expect(loki.loki_http_request("fake", batch, 1, 300, 10).class).to eql Net::HTTPTooManyRequests
+    it 'should send credentials' do
+      conf = {
+          'url'=>'http://localhost:3100/loki/api/v1/push',
+          'username' => 'foo',
+          'password' => 'bar',
+          'tenant_id' => 't'
+      }
+      loki = LogStash::Outputs::Loki.new(conf)
+      loki.register
+      b = Batch.new(entry)
+      post = stub_request(:post, "http://localhost:3100/loki/api/v1/push").with(
+          basic_auth: ['foo', 'bar'],
+          body: b.to_json,
+          headers:{
+              'Content-Type' => 'application/json' ,
+              'User-Agent' => 'loki-logstash',
+              'X-Scope-OrgID'=>'t',
+              'Accept'=>'*/*',
+              'Accept-Encoding'=>'gzip;q=1.0,deflate;q=0.6,identity;q=0.3',
+          }
+      )
+      loki.send(b)
+      expect(post).to have_been_requested.times(1)
+    end
 
-      server_error = Net::HTTPServerError.new(1.0, 429, 'OK')
-      allow(loki).to receive(:loki_http_request) { server_error }
-      allow(server_error).to receive(:payload).and_return('fake body')
-      expect(loki.loki_http_request("fake", batch, 1, 300, 10).class).to eql Net::HTTPServerError
+    it 'should not send credentials' do
+      conf = {
+          'url'=>'http://foo.com/loki/api/v1/push',
+      }
+      loki = LogStash::Outputs::Loki.new(conf)
+      loki.register
+      b = Batch.new(entry)
+      post = stub_request(:post, "http://foo.com/loki/api/v1/push").with(
+          body: b.to_json,
+          headers:{
+              'Content-Type' => 'application/json' ,
+              'User-Agent' => 'loki-logstash',
+              'Accept'=>'*/*',
+              'Accept-Encoding'=>'gzip;q=1.0,deflate;q=0.6,identity;q=0.3',
+          }
+      )
+      loki.send(b)
+      expect(post).to have_been_requested.times(1)
+    end
+    it 'should retry 500' do
+      conf = {
+          'url'=>'http://foo.com/loki/api/v1/push',
+          'retries' => 3,
+      }
+      loki = LogStash::Outputs::Loki.new(conf)
+      loki.register
+      b = Batch.new(entry)
+      post = stub_request(:post, "http://foo.com/loki/api/v1/push").with(
+          body: b.to_json,
+      ).to_return(status: [500, "Internal Server Error"])
+      loki.send(b)
+      loki.close
+      expect(post).to have_been_requested.times(3)
+    end
+    it 'should retry 429' do
+      conf = {
+          'url'=>'http://foo.com/loki/api/v1/push',
+          'retries' => 2,
+      }
+      loki = LogStash::Outputs::Loki.new(conf)
+      loki.register
+      b = Batch.new(entry)
+      post = stub_request(:post, "http://foo.com/loki/api/v1/push").with(
+          body: b.to_json,
+          ).to_return(status: [429, "stop spamming"])
+      loki.send(b)
+      loki.close
+      expect(post).to have_been_requested.times(2)
+    end
+    it 'should not retry 400' do
+      conf = {
+          'url'=>'http://foo.com/loki/api/v1/push',
+          'retries' => 11,
+      }
+      loki = LogStash::Outputs::Loki.new(conf)
+      loki.register
+      b = Batch.new(entry)
+      post = stub_request(:post, "http://foo.com/loki/api/v1/push").with(
+          body: b.to_json,
+          ).to_return(status: [400, "bad request"])
+      loki.send(b)
+      loki.close
+      expect(post).to have_been_requested.times(1)
+    end
+    it 'should retry exception' do
+      conf = {
+          'url'=>'http://foo.com/loki/api/v1/push',
+          'retries' => 11,
+      }
+      loki = LogStash::Outputs::Loki.new(conf)
+      loki.register
+      b = Batch.new(entry)
+      post = stub_request(:post, "http://foo.com/loki/api/v1/push").with(
+          body: b.to_json,
+          ).to_raise("some error").then.to_return(status: [200, "fine !"])
+      loki.send(b)
+      loki.close
+      expect(post).to have_been_requested.times(2)
     end
   end
 end

metadata

@@ -1,14 +1,15 @@
 --- !ruby/object:Gem::Specification
 name: logstash-output-loki
 version: !ruby/object:Gem::Version
-  version: 1.0.0
+  version: 1.0.1
 platform: ruby
 authors:
 - Aditya C S
+- Cyril Tovena
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2020-07-13 00:00:00.000000000 Z
+date: 2020-07-16 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   requirement: !ruby/object:Gem::Requirement
@@ -75,6 +76,7 @@ dependencies:
 description: Output plugin to ship logs to a Grafana Loki server
 email:
 - aditya.gnu@gmail.com
+- cyril.tovena@grafana.com
 executables: []
 extensions: []
 extra_rdoc_files: []
@@ -85,6 +87,7 @@ files:
 - lib/logstash/outputs/loki/batch.rb
 - lib/logstash/outputs/loki/entry.rb
 - logstash-output-loki.gemspec
+- spec/outputs/loki/entry_spec.rb
 - spec/outputs/loki_spec.rb
 homepage: https://github.com/grafana/loki/
 licenses:
@@ -112,4 +115,5 @@ signing_key:
 specification_version: 4
 summary: Output plugin to ship logs to a Grafana Loki server
 test_files:
+- spec/outputs/loki/entry_spec.rb
 - spec/outputs/loki_spec.rb