fluent-plugin-quota-throttle 0.0.2

data/lib/fluent/plugin/matcher.rb

@@ -0,0 +1,61 @@
+##
+# This module is responsible for matching records to quotas.
+module Matcher
+
+  ##
+  # MatchHelper class is responsible for matching records to quotas.
+  # Methods:
+  #   +get_quota+: Takes a list of keys and returns the quota that maximally matches
+  #   +matching_score+: Calculates the matching score between two hashes.
+  class MatchHelper
+
+    def initialize(processed_quotas,default_quota)
+      @quotas = processed_quotas
+      @default_quota = default_quota
+    end
+
+    # Takes a list of keys and returns the quota that maximally matches
+    # If no quota matches, returns the default quota
+    # Params:
+    #   +record+: (Hash) A hash of keys and values to match against the quotas
+    # Returns:
+    #   +quota+: (Quota Class)The quota that maximally matches the record
+    def get_quota(record)
+
+      max_score = 0
+      quota_to_return = @default_quota
+      if @quotas.nil?
+        return @default_quota
+      end
+      @quotas.each do |quota|
+        score = matching_score(quota.match_by, record)
+        if score > max_score
+          max_score = score
+          quota_to_return = quota
+        end
+      end
+      quota_to_return
+    end
+
+    private
+
+    # Calculates the matching score between two hashes.
+    # Params:
+    #   +match+: (Hash) A hash of keys and values to match against the record
+    #   +record+: (Hash) A hash of keys and values to match against the match
+    def matching_score(match, record)
+      score = 0
+      if match.nil? || record.nil?
+        return 0
+      end
+      match.each do |key, value|
+        if record.dig(*key) == value
+          score += 1
+        else
+          return 0
+        end
+      end
+      score
+    end
+  end
+end

data/lib/fluent/plugin/rate_limiter.rb

@@ -0,0 +1,113 @@
+##
+# Rate Limiter module, contains the rate limiting logic
+module RateLimiter
+
+  ##
+  # Bucket class, contains the rate limiting logic for each group
+  # Attributes:
+  #   +bucket_count+: Number of requests in the bucket
+  #   +bucket_last_reset+: Time when the bucket was last reset
+  #   +approx_rate_per_second+: Approximate rate of requests per second
+  #   +rate_last_reset+: Time when the rate was last reset
+  #   +curr_count+: Number of requests in the current second
+  #   +last_warning+: Time when the last warning was issued
+  #   +timeout_s+: Timeout for the bucket
+  #   +bucket_limit+: Maximum number of requests allowed in the bucket
+  #   +bucket_period+: Time period for the bucket
+  #   +rate_limit+: Maximum number of requests allowed per second
+  class Bucket
+    attr_accessor :bucket_count, :bucket_last_reset, :approx_rate_per_second, :rate_last_reset, :curr_count, :last_warning
+    attr_reader :bucket_limit, :bucket_period, :rate_limit, :timeout_s, :group
+    def initialize( group, bucket_limit, bucket_period)
+      now = Time.now
+      @group = group
+      @bucket_count = 0
+      @bucket_last_reset = now
+      @approx_rate_per_second = 0
+      @rate_last_reset = now
+      @curr_count = 0
+      @last_warning = nil
+      @bucket_limit = bucket_limit
+      @bucket_period = bucket_period
+      @rate_limit = bucket_limit/bucket_period
+      @timeout_s = 2*bucket_period
+    end
+
+    # Checks if the bucket is free or full
+    # Returns:
+    #   +true+ if the bucket is free
+    #   +false+ if the bucket is full
+    def allow
+      if @bucket_limit == -1
+        return true
+      end
+      now = Time.now
+      @curr_count += 1
+      time_lapsed = now - @rate_last_reset
+
+      if time_lapsed.to_i >= 1
+        @approx_rate_per_second = @curr_count / time_lapsed
+        @rate_last_reset = now
+        @curr_count = 0
+      end
+
+      if now.to_i / @bucket_period > @bucket_last_reset.to_i / @bucket_period
+        reset_bucket
+      end
+
+      if @bucket_count == -1 or @bucket_count > @bucket_limit
+        @bucket_count = -1
+        return false
+      else
+        @bucket_count += 1
+        true
+      end
+    end
+
+    # Checks if bucket is expired
+    # Returns:
+    #   +true+ if the bucket is expired
+    #   +false+ if the bucket is not expired
+    def expired
+      now = Time.now
+      now.to_i - @rate_last_reset.to_i > @timeout_s
+    end
+
+    private
+
+    # Resets the bucket when the window moves to the next time period
+    def reset_bucket
+      now = Time.now
+      unless @bucket_count == -1 && @approx_rate_per_second > @rate_limit
+        @bucket_count = 0
+        @bucket_last_reset = now
+      end
+    end
+  end
+
+  ##
+  # BucketStore class, organizes the all the group buckets
+  # Attributes:
+  #   +buckets+: Hash containing all the group buckets
+  class BucketStore
+    def initialize
+      @buckets = {}
+    end
+
+    # Gets the bucket for the group
+    # Arguments:
+    #   +group+: Group for which the bucket is required
+    #   +quota+: Quota object containing the bucket size and duration
+    def get_bucket(group, quota)
+      @buckets[group] = @buckets.delete(group) || Bucket.new( group, quota.bucket_size, quota.duration)
+    end
+
+    # Cleans the buckets that have expired
+    def clean_buckets
+      lru_group, lru_bucket = @buckets.first
+      if !lru_group.nil? && lru_bucket.expired
+        @buckets.delete(lru_group)
+      end
+    end
+  end
+end

data/test/config_files/filter_plugin_test.yml

@@ -0,0 +1,28 @@
+quotas:
+  - name: quota1
+    description: first quota
+    group_by:
+      - group1.a
+    match_by:
+      group1.a: value1
+    bucket_size: 5
+    duration: 60s
+    action: drop
+  - name: quota2
+    description: second quota
+    group_by:
+      - group1.a
+      - group1.b
+    match_by:
+      group1.a: value2
+      group1.b: value3
+    bucket_size: 6
+    duration: 120s
+    action: reemit
+default:
+  description: default quota
+  group_by:
+    - group1.a
+  bucket_size: 4
+  duration: 180s
+  action: reemit

data/test/config_files/matcher_test.yml

@@ -0,0 +1,39 @@
+quotas:
+- name: quota1
+  description: first quota
+  group_by:
+  - group1.a
+  match_by:
+    group1.a: value1
+  bucket_size: 100
+  duration: 60s
+  action: drop
+- name: quota2
+  description: second quota
+  group_by:
+  - group1.a
+  - group1.b
+  match_by:
+    group1.a: value2
+    group1.b: value3
+  bucket_size: 200
+  duration: 120s
+  action: reemit
+- name: quota3
+  description: third quota
+  group_by:
+  - group2
+  - group3
+  match_by:
+      group2: value2
+      group3: value3
+  bucket_size: 300
+  duration: 180s
+  action: drop
+default:
+  description: default quota
+  group_by:
+    - group1.a
+  bucket_size: 300
+  duration: 180s
+  action: reemit

data/test/config_files/parser_test.yml

@@ -0,0 +1,28 @@
+quotas:
+- name: quota1
+  description: first quota
+  group_by:
+  - group1.a
+  match_by:
+    group1.a: value1
+  bucket_size: 100
+  duration: 1m
+  action: drop
+- name: quota2
+  description: second quota
+  group_by:
+  - group1.a
+  - group1.b
+  match_by:
+    group1.a: value2
+    group1.b: value3
+  bucket_size: 200
+  duration: 120s
+  action: reemit
+default:
+  description: default quota
+  group_by:
+    - group1.a
+  bucket_size: 300
+  duration: 3m
+  action: reemit

data/test/fluent/plugin/filter_quota_throttle_test.rb

@@ -0,0 +1,85 @@
+require_relative '../../helper'
+
+class QuotaThrottleFilterTest < Minitest::Test
+  include Fluent::Test::Helpers
+
+  def setup
+    Fluent::Test.setup
+  end
+
+  CONFIG = %[
+    path test/config_files/filter_plugin_test.yml
+    warning_delay 2m
+    enable_metrics false
+  ]
+
+  def create_driver(conf = CONFIG)
+    Fluent::Test::Driver::Filter.new(Fluent::Plugin::QuotaThrottleFilter).configure(conf)
+  end
+
+  def test_configure
+    d = create_driver
+    assert_equal "test/config_files/filter_plugin_test.yml", d.instance.path
+    assert_equal 120, d.instance.warning_delay
+  end
+
+  def test_filter
+    d = create_driver
+    d.run(default_tag: 'test') do
+      10.times do
+        d.feed("group1" => { "a" => "value1" , "b" => "value2" })
+        d.feed("group1" => { "a" => "value2" , "b" => "value3" })
+        d.feed("group1" => { "a" => "value2" , "b" => "value2" })
+        d.feed("group1" => { "a" => "value3" , "b" => "value2" })
+      end
+    end
+    events = d.filtered_records
+    assert_equal 23, events.length
+  end
+
+  # Due to the way the Driver is implemented, both metrics tests cannot be same time because the registry of metrics is not cleared between tests
+  # def test_metrics_without_labels
+  #   modified_config = CONFIG.sub("enable_metrics false", "enable_metrics true")
+  #   d = create_driver(modified_config)
+  #   d.run(default_tag: 'test') do
+  #     10.times do
+  #       d.feed("group1" => { "a" => "value1" , "b" => "value2" })
+  #       d.feed("group1" => { "a" => "value2" , "b" => "value3" })
+  #       d.feed("group1" => { "a" => "value2" , "b" => "value2" })
+  #       d.feed("group1" => { "a" => "value3" , "b" => "value2" })
+  #     end
+  #   end
+  #   assert_equal 10, d.instance.registry.get(:fluentd_quota_throttle_input).get(labels: {quota: 'quota1'})
+  #   assert_equal 4, d.instance.registry.get(:fluentd_quota_throttle_exceeded).get(labels: {quota: 'quota1'})
+  #   assert_equal 10, d.instance.registry.get(:fluentd_quota_throttle_input).get(labels: {quota: 'quota2'})
+  #   assert_equal 3, d.instance.registry.get(:fluentd_quota_throttle_exceeded).get(labels: {quota: 'quota2'})
+  #   assert_equal 20, d.instance.registry.get(:fluentd_quota_throttle_input).get(labels: {quota: 'default'})
+  #   assert_equal 10, d.instance.registry.get(:fluentd_quota_throttle_exceeded).get(labels: {quota: 'default'})
+  # end
+  def test_metrics_with_labels
+    labels = %[
+        <labels>
+          source $.group1.a
+          dummy d1
+        </labels>
+    ]
+    modified_config = CONFIG.sub("enable_metrics false", "enable_metrics true" + labels)
+    d = create_driver(modified_config)
+    d.run(default_tag: 'test') do
+      10.times do
+        d.feed("group1" => { "a" => "value1" , "b" => "value2" })
+        d.feed("group1" => { "a" => "value2" , "b" => "value3" })
+        d.feed("group1" => { "a" => "value2" , "b" => "value2" })
+        d.feed("group1" => { "a" => "value3" , "b" => "value2" })
+      end
+    end
+    assert_equal 10, d.instance.registry.get(:fluentd_quota_throttle_input).get(labels: {source: "value1", quota: 'quota1', dummy: 'd1'})
+    assert_equal 4, d.instance.registry.get(:fluentd_quota_throttle_exceeded).get(labels: {source: "value1", quota: 'quota1', dummy: 'd1'})
+    assert_equal 10, d.instance.registry.get(:fluentd_quota_throttle_input).get(labels: {source: "value2", quota: 'quota2', dummy: 'd1'})
+    assert_equal 3, d.instance.registry.get(:fluentd_quota_throttle_exceeded).get(labels: {source: "value2", quota: 'quota2', dummy: 'd1'})
+    assert_equal 10, d.instance.registry.get(:fluentd_quota_throttle_input).get(labels: {source: "value2", quota: 'default', dummy: 'd1'})
+    assert_equal 5, d.instance.registry.get(:fluentd_quota_throttle_exceeded).get(labels: {source: "value2", quota: 'default', dummy: 'd1'})
+    assert_equal 10, d.instance.registry.get(:fluentd_quota_throttle_input).get(labels: {source: "value3", quota: 'default', dummy: 'd1'})
+    assert_equal 5, d.instance.registry.get(:fluentd_quota_throttle_exceeded).get(labels: {source: "value3", quota: 'default', dummy: 'd1'})
+  end
+end

data/test/helper.rb

@@ -0,0 +1,10 @@
+$LOAD_PATH.unshift File.expand_path("../test", __dir__)
+require 'fluent/test'
+require 'fluent/test/driver/filter'
+require 'fluent/test/helpers'
+require 'yaml'
+require 'minitest/autorun'
+require_relative '../lib/fluent/plugin/config_parser'
+require_relative '../lib/fluent/plugin/matcher'
+require_relative '../lib/fluent/plugin/rate_limiter'
+require_relative '../lib/fluent/plugin/filter_quota_throttle'

data/test/modules/config_parser_test.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+require_relative '../helper'
+
+class ParserTest < Minitest::Test
+  def setup
+    config_file_path = Dir.pwd+"/test/config_files/parser_test.yml"
+    config_parser = ConfigParser::Configuration.new(config_file_path)
+    @quotas = config_parser.quotas
+    @default_quota = config_parser.default_quota
+  end
+
+  def test_parse_quotas
+    # Check if the quotas are parsed correctly
+    # TODO 1: Add wrong configuration files and check if it raises errors
+    refute_nil @quotas
+    assert_equal 2, @quotas.length
+
+    # Check if the quotas are parsed correctly
+    assert_equal "quota1", @quotas[0].name
+    assert_equal "quota2", @quotas[1].name
+    assert_equal "first quota", @quotas[0].desc
+    assert_equal "second quota", @quotas[1].desc
+    assert_equal [["group1", "a"]], @quotas[0].group_by
+    assert_equal [["group1", "a"], ["group1", "b"]], @quotas[1].group_by
+    assert_equal 100, @quotas[0].bucket_size
+    assert_equal 200, @quotas[1].bucket_size
+    assert_equal 60, @quotas[0].duration
+    assert_equal 120, @quotas[1].duration
+    assert_equal "drop", @quotas[0].action
+    assert_equal "reemit", @quotas[1].action
+    assert_equal ({["group1", "a"] => "value1"}), @quotas[0].match_by
+    assert_equal ({["group1", "a"] => "value2", ["group1", "b"] => "value3"}), @quotas[1].match_by
+    assert_equal "default", @default_quota.name
+    assert_equal "default quota", @default_quota.desc
+    assert_equal [["group1", "a"]], @default_quota.group_by
+    assert_equal [], @default_quota.match_by
+    assert_equal 300, @default_quota.bucket_size
+    assert_equal 180, @default_quota.duration
+    assert_equal "reemit", @default_quota.action
+  end
+
+end

data/test/modules/matcher_test.rb

@@ -0,0 +1,67 @@
+require_relative '../helper'
+
+class MatcherTest < Minitest::Test
+  def setup
+    config_file_path = Dir.pwd+"/test/config_files/matcher_test.yml"
+    config_parser = ConfigParser::Configuration.new(config_file_path)
+    quotas = config_parser.quotas
+    @default_quota = config_parser.default_quota
+    @match_helper = Matcher::MatchHelper.new(quotas,@default_quota)
+  end
+
+  def test_get_quota
+    # Check if the correct quota is retrieved from above definition
+    # UT 1: Subset of groups match fully
+    keys = { "group1" => { "a" => "value1" , "b" => "value2" } }
+    quota = @match_helper.get_quota(keys)
+    assert_equal "quota1", quota.name
+
+    # UT 2: All groups match fully
+    keys = { "group1" => {"a" => "value2" , "b" => "value3" } }
+    quota = @match_helper.get_quota(keys)
+    assert_equal "quota2", quota.name
+
+    # UT 3: Subset of group match partially
+    keys = { "group1" => { "a" => "value2" , "b" => "value2" } }
+    quota = @match_helper.get_quota(keys)
+    assert_equal @default_quota, quota
+
+    # UT 4: None of the group matches
+    keys = { "group1" => { "a" => "value3" , "b" => "value2" } }
+    quota = @match_helper.get_quota(keys)
+    assert_equal @default_quota, quota
+
+    # UT 5: Non-nested group matches
+    keys = { "group2" => "value2" , "group3" => "value3" }
+    quota = @match_helper.get_quota(keys)
+    assert_equal "quota3", quota.name
+
+    # UT 6: Non-nested group mismatches
+    keys = { "group2" => "value2" , "group3" => "value4" }
+    quota = @match_helper.get_quota(keys)
+    assert_equal @default_quota, quota
+  end
+
+  def test_matching_score
+    # Testing the private helper function
+    score = @match_helper.send(:matching_score, { "key1" => "value1" , "key2" => "value2" }, { "key1" => "value1" , "key2" => "value2" })
+    assert_equal 2, score
+
+    score = @match_helper.send(:matching_score, { "key1" => "value1" }, { "key2" => "value2" })
+    assert_equal 0, score
+
+    score = @match_helper.send(:matching_score, { "key1" => "value1" , "key2" => "value2" }, { "key1" => "value1" , "key2" => "value3" })
+    assert_equal 0, score
+
+    score = @match_helper.send(:matching_score, nil, { "key1" => "value1" })
+    assert_equal 0, score
+
+    score = @match_helper.send(:matching_score, { "key1" => "value1" }, nil)
+    assert_equal 0, score
+
+    score = @match_helper.send(:matching_score, nil, nil)
+    assert_equal 0, score
+  end
+
+
+end

data/test/modules/rate_limiter_test.rb

@@ -0,0 +1,72 @@
+require_relative '../helper'
+
+class TestBucket < Minitest::Test
+  def setup
+    # Initialize bucket with dummy parameters
+    @bucket = RateLimiter::Bucket.new( ["group"] ,10, 2)
+  end
+
+  def test_bucket_initialization
+    assert_equal 0, @bucket.bucket_count
+    assert @bucket.bucket_last_reset <= Time.now
+    assert_equal 0, @bucket.approx_rate_per_second
+    assert @bucket.rate_last_reset <= Time.now
+    assert_equal 0, @bucket.curr_count
+    assert_nil @bucket.last_warning
+    assert_equal 10, @bucket.instance_variable_get(:@bucket_limit)
+    assert_equal 2, @bucket.instance_variable_get(:@bucket_period)
+    assert_equal 5, @bucket.instance_variable_get(:@rate_limit)
+    assert_equal 4, @bucket.timeout_s
+  end
+
+  def test_bucket_allow_free
+    @bucket.allow
+    assert_equal 1, @bucket.bucket_count
+  end
+
+  def test_bucket_allow_full
+    11.times { @bucket.allow }
+    assert_equal false, @bucket.allow
+  end
+
+  def test_reset_bucket
+    @bucket.allow
+    @bucket.send(:reset_bucket)
+    assert_equal 0, @bucket.bucket_count
+  end
+
+  def test_expired
+    @bucket.allow
+    assert_equal false, @bucket.expired
+    sleep(5)
+    assert_equal true, @bucket.expired
+  end
+end
+
+class TestBucketStore < Minitest::Test
+  def setup
+    @bucket_store = RateLimiter::BucketStore.new
+    @quota = ConfigParser::Quota.new("Dummy", "dummy quota for testing",[["group1","a"],["group1","b"]], {["group1","a"] => "value1"}, 10, "2s", "reemit")
+  end
+
+  def test_get_bucket
+    group = "value1"
+    bucket = @bucket_store.get_bucket(group, @quota)
+    assert_instance_of RateLimiter::Bucket, bucket
+    assert_equal 10, bucket.instance_variable_get(:@bucket_limit)
+    assert_equal 2, bucket.instance_variable_get(:@bucket_period)
+  end
+
+  def test_clean_buckets
+    group1 = "value1"
+    @bucket_store.get_bucket(group1, @quota)
+    group2 = "value2"
+    @bucket_store.get_bucket(group2, @quota)
+    lru_group, lru_counter = @bucket_store.instance_variable_get(:@buckets).first
+    assert_equal group1, lru_group
+    sleep(5)
+    @bucket_store.clean_buckets
+    lru_group, lru_counter = @bucket_store.instance_variable_get(:@buckets).first
+    assert_equal group2, lru_group
+  end
+end

data/test/scripts/example.conf

@@ -0,0 +1,24 @@
+<source>
+  @type forward
+  port 24224
+</source>
+
+<source>
+    @type prometheus
+    port 24231
+</source>
+
+<filter test.**>
+  @log_level debug
+  @type quota_throttle
+  path example.yml
+  warning_delay 30
+  enable_metrics true
+  <labels>
+   user_id $.user_id
+  </labels>
+</filter>
+
+<match **>
+  @type stdout
+</match>

data/test/scripts/example.yml

@@ -0,0 +1,24 @@
+default:
+  action: drop
+  bucket_size: 6
+  duration: 5s
+  group_by:
+    - "user_id"
+
+quotas:
+  - name: "Quota1"
+    action: reemit
+    bucket_size: 8
+    duration: 5s
+    group_by:
+      - "user_id"
+    match_by:
+      user_id: "user1"
+  - name: "Quota2"
+    action: drop
+    bucket_size: 5
+    duration: 5s
+    group_by:
+      - "user_id"
+    match_by:
+      user_id: "user2"

data/test/scripts/test.sh

@@ -0,0 +1,10 @@
+for i in {1..20}; do
+	echo "{\"group\": \"group1\", \"user_id\": \"user1\", \"message\": \"Test log $i\"}" | fluent-cat test.default
+done
+for i in {1..20}; do
+	echo "{\"group\": \"group1\", \"user_id\": \"user3\", \"message\": \"Test log $i\"}" | fluent-cat test.default
+done
+for i in {1..20}; do
+	echo "{\"group\": \"group1\", \"user_id\": \"user2\", \"message\": \"Test log $i\"}" | fluent-cat test.default
+done
+