Adds temp file buffer used during network downtime

Author: MonishkaDas

lib/logstash/outputs/kusto.rb

@@ -77,8 +77,13 @@ class LogStash::Outputs::Kusto < LogStash::Outputs::Base
   default :codec, 'json_lines'
 
   def register
-    # Initialize the custom buffer with size and interval
-    @buffer = LogStash::Outputs::CustomSizeBasedBuffer.new(@max_size, @max_interval) do |events|
+    # Set buffer_file to a valid file path
+    buffer_file = File.join(Dir.pwd, 'buffer', 'kusto_buffer.dat')
+  
+    # Ensure the buffer directory exists
+    FileUtils.mkdir_p(File.dirname(buffer_file))
+    # Initialize the custom buffer with size, interval, and buffer file
+    @buffer = LogStash::Outputs::CustomSizeBasedBuffer.new(@max_size, @max_interval, buffer_file) do |events|
       flush_buffer(events)
     end
 

lib/logstash/outputs/kusto/custom_size_based_buffer.rb

@@ -1,13 +1,17 @@
 require 'logger'
 require 'thread'
+require 'fileutils'
 
 module LogStash
   module Outputs
     class CustomSizeBasedBuffer
-      def initialize(max_size_mb, max_interval, &flush_callback)
+      def initialize(max_size_mb, max_interval, buffer_file, &flush_callback)
+        raise ArgumentError, "buffer_file cannot be nil" if buffer_file.nil?
+
         @buffer_config = {
           max_size: max_size_mb * 1024 * 1024, # Convert MB to bytes
           max_interval: max_interval,
+          buffer_file: buffer_file,
           logger: Logger.new(STDOUT)
         }
         @buffer_state = {
@@ -25,6 +29,7 @@ def initialize(max_size_mb, max_interval, &flush_callback)
         @shutdown = false
         @pending_mutex = Mutex.new
         @flush_mutex = Mutex.new
+        load_buffer_from_file
         @buffer_config[:logger].info("CustomSizeBasedBuffer initialized with max_size: #{max_size_mb} MB, max_interval: #{max_interval} seconds")
       end
 
@@ -46,6 +51,7 @@ def shutdown
         @shutdown = true
         @buffer_state[:timer].kill
         buffer_flush(final: true)
+        clear_file_buffer
       end
 
       private
@@ -69,30 +75,49 @@ def buffer_flush(options = {})
         items_flushed = 0
 
         begin
+          outgoing_items = []
+          outgoing_size = 0
+
           @pending_mutex.synchronize do
             return 0 if @buffer_state[:pending_size] == 0
 
             time_since_last_flush = Time.now.to_i - @buffer_state[:last_flush]
 
-            return 0 if !force && @buffer_state[:pending_size] < @buffer_config[:max_size] && time_since_last_flush < @buffer_config[:max_interval]
+            if !force && @buffer_state[:pending_size] < @buffer_config[:max_size] && time_since_last_flush < @buffer_config[:max_interval]
+              return 0
+            end
 
             if force
               @buffer_config[:logger].info("Time-based flush triggered after #{@buffer_config[:max_interval]} seconds")
             elsif @buffer_state[:pending_size] >= @buffer_config[:max_size]
               @buffer_config[:logger].info("Size-based flush triggered at #{@buffer_state[:pending_size]} bytes was reached")
+            else
+              @buffer_config[:logger].info("Flush triggered without specific condition")
             end
 
             outgoing_items = @buffer_state[:pending_items].dup
             outgoing_size = @buffer_state[:pending_size]
             buffer_initialize
+          end
 
+          begin
             @flush_callback.call(outgoing_items) # Pass the list of events to the callback
+            clear_flushed_buffer_states(outgoing_items) unless ::File.zero?(@buffer_config[:buffer_file]) # Clear the flushed items from the file
+          rescue => e
+            @buffer_config[:logger].error("Flush failed: #{e.message}")
+            # Save the items to the file buffer in case of failure
+            @pending_mutex.synchronize do
+              @buffer_state[:pending_items] = outgoing_items + @buffer_state[:pending_items]
+              @buffer_state[:pending_size] += outgoing_size
+              save_buffer_to_file
+            end
+            raise e
+          end
 
-            @buffer_state[:last_flush] = Time.now.to_i
-            @buffer_config[:logger].info("Flush completed. Flushed #{outgoing_items.size} events, #{outgoing_size} bytes")
+          @buffer_state[:last_flush] = Time.now.to_i
+          @buffer_config[:logger].info("Flush completed. Flushed #{outgoing_items.size} events, #{outgoing_size} bytes")
 
-            items_flushed = outgoing_items.size
-          end
+          items_flushed = outgoing_items.size
         ensure
           @flush_mutex.unlock
         end
@@ -104,6 +129,78 @@ def buffer_initialize
         @buffer_state[:pending_items] = []
         @buffer_state[:pending_size] = 0
       end
+
+      def clear_flushed_buffer_states(flushed_items)
+        remaining_buffer_states = []
+        ::File.foreach(@buffer_config[:buffer_file]) do |line|
+          begin
+            buffer_state = Marshal.load(line)
+            buffer_state[:pending_items] -= flushed_items
+            buffer_state[:pending_size] = buffer_state[:pending_items].sum(&:bytesize)
+            remaining_buffer_states << buffer_state unless buffer_state[:pending_items].empty?
+          rescue ArgumentError => e
+            @buffer_config[:logger].error("Failed to load buffer state: #{e.message}")
+            next
+          end
+        end
+      
+        ::File.open(@buffer_config[:buffer_file], 'w') do |file|
+          remaining_buffer_states.each do |state|
+            file.write(Marshal.dump(state) + "\n")
+          end
+        end
+      end
+
+      def save_buffer_to_file
+        buffer_state_copy = @buffer_state.dup
+        buffer_state_copy.delete(:timer) # Exclude the Thread object from serialization
+      
+        ::FileUtils.mkdir_p(::File.dirname(@buffer_config[:buffer_file])) # Ensure directory exists
+        ::File.open(@buffer_config[:buffer_file], 'a') do |file|
+          file.write(Marshal.dump(buffer_state_copy) + "\n")
+        end
+        @buffer_config[:logger].info("Saved buffer state to file")
+      end
+      
+      def load_buffer_from_file
+        ::FileUtils.mkdir_p(::File.dirname(@buffer_config[:buffer_file])) # Ensure directory exists
+        ::File.open(@buffer_config[:buffer_file], 'a') {} # Create the file if it doesn't exist
+      
+        if ::File.file?(@buffer_config[:buffer_file]) && !::File.zero?(@buffer_config[:buffer_file])
+          begin
+            @pending_mutex.synchronize do
+              buffer_states = []
+              ::File.foreach(@buffer_config[:buffer_file]) do |line|
+                buffer_states << Marshal.load(line)
+              end
+              @buffer_state = buffer_states.reduce do |acc, state|
+                acc[:pending_items].concat(state[:pending_items])
+                acc[:pending_size] += state[:pending_size]
+                acc
+              end
+              @buffer_state[:timer] = Thread.new do
+                loop do
+                  sleep(@buffer_config[:max_interval])
+                  buffer_flush(force: true)
+                end
+              end
+              # Ensure the buffer does not flush immediately upon loading
+              @buffer_state[:last_flush] = Time.now.to_i
+            end
+            @buffer_config[:logger].info("Loaded buffer state from file")
+          rescue => e
+            @buffer_config[:logger].error("Failed to load buffer from file: #{e.message}")
+            buffer_initialize
+          end
+        else
+          buffer_initialize
+        end
+      end
+
+      def clear_file_buffer
+        ::File.open(@buffer_config[:buffer_file], 'w') {} # Truncate the file
+        @buffer_config[:logger].info("File buffer cleared on shutdown")
+      end
     end
   end
 end

lib/logstash/outputs/kusto/ingestor.rb

@@ -126,10 +126,15 @@ def upload_async(data)
 
       @workers_pool.post do
         LogStash::Util.set_thread_name("Kusto to ingest data")
-        upload(data)
+        begin
+          upload(data)
+        rescue => e
+          @logger.error('Error during async upload.', exception: e.class, message: e.message, backtrace: e.backtrace)
+          raise e
+        end
       end
     rescue Exception => e
-      @logger.error('StandardError.', exception: e.class, message: e.message, backtrace: e.backtrace)
+      @logger.error('StandardError in upload_async.', exception: e.class, message: e.message, backtrace: e.backtrace)
       raise e
     end
 
@@ -150,21 +155,30 @@ def upload(data)
       #   local_ingestion_properties.addJsonMappingName(json_mapping)
       # end
 
+
       if data.size > 0
-      data_source_info = Java::com.microsoft.azure.kusto.ingest.source.StreamSourceInfo.new(java.io.ByteArrayInputStream.new(data.to_java_bytes))
-      @kusto_client.ingestFromStream(data_source_info, @ingestion_properties)
-    else
-      @logger.warn("Data is empty and is not ingested.")
-    end
-    @logger.debug("Data sent to Kusto.")
-  rescue => e
-      # When the retry limit is reached or another error happen we will wait and retry.
+        data_source_info = Java::com.microsoft.azure.kusto.ingest.source.StreamSourceInfo.new(java.io.ByteArrayInputStream.new(data.to_java_bytes))
+        @kusto_client.ingestFromStream(data_source_info, @ingestion_properties)
+      else
+        @logger.warn("Data is empty and is not ingested.")
+      end
+      @logger.debug("Data sent to Kusto.")
+    rescue => e
+      # When the retry limit is reached or another error happens we will wait and retry.
       #
-      # Thread might be stuck here, but I think its better than losing anything
-      # its either a transient errors or something bad really happened.
+      # Thread might be stuck here, but I think it's better than losing anything
+      # it's either a transient error or something bad really happened.
       @logger.error('Uploading failed, retrying.', exception: e.class, message: e.message, backtrace: e.backtrace)
-      sleep RETRY_DELAY_SECONDS
-      retry
+      retry_count = 0
+      max_retries = 5
+      begin
+        sleep (2 ** retry_count) * RETRY_DELAY_SECONDS
+        retry_count += 1
+        retry if retry_count <= max_retries
+      rescue => retry_error
+        @logger.error('Retrying failed.', exception: retry_error.class, message: retry_error.message, backtrace: retry_error.backtrace)
+        raise retry_error
+      end
     end
 
     def stop

spec/outputs/kusto_spec.rb

@@ -50,13 +50,23 @@
 
       events = [LogStash::Event.new("message" => "test1"), LogStash::Event.new("message" => "test2")]
       encoded_events = events.map { |e| [e, e.to_json] }
-      kusto.multi_receive_encoded(encoded_events)
 
+      # Temporarily disable automatic flushing for the test
       buffer = kusto.instance_variable_get(:@buffer)
-      expect(buffer.instance_variable_get(:@buffer_state)[:pending_items].size).to eq(2)
+      allow(buffer).to receive(:buffer_flush)
+
+      # Clear the buffer before the test
+      buffer.instance_variable_set(:@buffer_state, { pending_items: [], pending_size: 0, last_flush: Time.now.to_i })
+
+      kusto.multi_receive_encoded(encoded_events)
+
+      pending_items = buffer.instance_variable_get(:@buffer_state)[:pending_items]
+      RSpec.configuration.reporter.message("Pending items in buffer: #{pending_items.inspect}")
+
+      expect(pending_items.size).to eq(2)
       RSpec.configuration.reporter.message("Completed test: processes events and adds them to the buffer")
     end
-  
+
     it 'handles errors during event processing' do
       RSpec.configuration.reporter.message("Running test: handles errors during event processing")
       kusto = described_class.new(options)
@@ -102,9 +112,14 @@
 
       events = [LogStash::Event.new("message" => "test1")]
       encoded_events = events.map { |e| [e, e.to_json] }
-      expect(kusto.instance_variable_get(:@ingestor)).to receive(:upload_async).with(anything)
+      # Ensure upload_async is called only once
+      expect(kusto.instance_variable_get(:@ingestor)).to receive(:upload_async).with(anything).once
       kusto.multi_receive_encoded(encoded_events)
-      kusto.flush_buffer(encoded_events) # Pass the encoded events
+      
+      # Trigger the buffer flush manually
+      buffer = kusto.instance_variable_get(:@buffer)
+      buffer.send(:buffer_flush, force: true)
+
       RSpec.configuration.reporter.message("Completed test: flushes the buffer when max_size is reached")
     end
 
@@ -119,7 +134,7 @@
       sleep(2) # Wait for the interval to pass
 
       expect(kusto.instance_variable_get(:@ingestor)).to receive(:upload_async).with(anything)
-      kusto.flush_buffer(encoded_events) # Pass the encoded events
+      kusto.flush_buffer(encoded_events) # Pass the encoded events here
       RSpec.configuration.reporter.message("Completed test: flushes the buffer when max_interval is reached")
     end
 
@@ -136,7 +151,7 @@
       sleep(2)
 
       expect(kusto.instance_variable_get(:@ingestor)).to receive(:upload_async).with(anything)
-      kusto.flush_buffer(encoded_events) # Pass the encoded events
+      kusto.flush_buffer(encoded_events) # Pass the encoded events here
       RSpec.configuration.reporter.message("Completed test: eventually flushes without receiving additional events based on max_interval")
     end
   end