fluentd222 1.16.2-x86_64-linux

data/lib/fluent/counter/mutex_hash.rb

@@ -0,0 +1,163 @@
+#
+# Fluentd
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+require 'timeout'
+
+module Fluent
+  module Counter
+    class MutexHash
+      def initialize(data_store)
+        @mutex = Mutex.new
+        @data_store = data_store
+        @mutex_hash = {}
+        @thread = nil
+        @cleanup_thread = CleanupThread.new(@data_store, @mutex_hash, @mutex)
+      end
+
+      def start
+        @data_store.start
+        @cleanup_thread.start
+      end
+
+      def stop
+        @data_store.stop
+        @cleanup_thread.stop
+      end
+
+      def synchronize(*keys)
+        return if keys.empty?
+
+        locks = {}
+        loop do
+          @mutex.synchronize do
+            keys.each do |key|
+              mutex = @mutex_hash[key]
+              unless mutex
+                v = Mutex.new
+                @mutex_hash[key] = v
+                mutex = v
+              end
+
+              if mutex.try_lock
+                locks[key] = mutex
+              else
+                locks.each_value(&:unlock)
+                locks = {}          # flush locked keys
+                break
+              end
+            end
+          end
+
+          next if locks.empty?      # failed to lock all keys
+
+          locks.each do |(k, v)|
+            yield @data_store, k
+            v.unlock
+          end
+          break
+        end
+      end
+
+      def synchronize_keys(*keys)
+        return if keys.empty?
+        keys = keys.dup
+
+        while key = keys.shift
+          @mutex.lock
+
+          mutex = @mutex_hash[key]
+          unless mutex
+            v = Mutex.new
+            @mutex_hash[key] = v
+            mutex = v
+          end
+
+          if mutex.try_lock
+            @mutex.unlock
+            yield @data_store, key
+            mutex.unlock
+          else
+            # release global lock
+            @mutex.unlock
+            keys.push(key)          # failed lock, retry this key
+          end
+        end
+      end
+    end
+
+    class CleanupThread
+      CLEANUP_INTERVAL = 60 * 15 # 15 min
+
+      def initialize(store, mutex_hash, mutex)
+        @store = store
+        @mutex_hash = mutex_hash
+        @mutex = mutex
+        @thread = nil
+        @running = false
+      end
+
+      def start
+        @running = true
+        @thread = Thread.new do
+          while @running
+            sleep CLEANUP_INTERVAL
+            run_once
+          end
+        end
+      end
+
+      def stop
+        return unless @running
+        @running = false
+        begin
+          # Avoid waiting CLEANUP_INTERVAL
+          Timeout.timeout(1) do
+            @thread.join
+          end
+        rescue Timeout::Error
+          @thread.kill
+        end
+      end
+
+      private
+
+      def run_once
+        @mutex.synchronize do
+          last_cleanup_at = (Time.now - CLEANUP_INTERVAL).to_i
+          @mutex_hash.each do |(key, mutex)|
+            v = @store.get(key, raw: true)
+            next unless v
+            next if last_cleanup_at < v['last_modified_at'][0] # v['last_modified_at'] = [sec, nsec]
+            next unless mutex.try_lock
+
+            @mutex_hash[key] = nil
+            mutex.unlock
+
+            # Check that a waiting thread is in a lock queue.
+            # Can't get a lock here means this key is used in other places.
+            # So restore a mutex value to a corresponding key.
+            if mutex.try_lock
+              @mutex_hash.delete(key)
+              mutex.unlock
+            else
+              @mutex_hash[key] = mutex
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/fluent/counter/server.rb

@@ -0,0 +1,273 @@
+#
+# Fluentd
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+require 'cool.io'
+require 'fluent/counter/base_socket'
+require 'fluent/counter/validator'
+require 'fluent/counter/store'
+require 'fluent/counter/mutex_hash'
+
+module Fluent
+  module Counter
+    class Server
+      DEFAULT_ADDR = '127.0.0.1'
+      DEFAULT_PORT = 24321
+
+      def initialize(name, opt = {})
+        raise 'Counter server name is invalid' unless Validator::VALID_NAME.match?(name)
+        @name = name
+        @opt = opt
+        @addr = @opt[:addr] || DEFAULT_ADDR
+        @port = @opt[:port] || DEFAULT_PORT
+        @loop = @opt[:loop] || Coolio::Loop.new
+        @log =  @opt[:log] || $log
+
+        @store = Fluent::Counter::Store.new(opt)
+        @mutex_hash = MutexHash.new(@store)
+
+        @server = Coolio::TCPServer.new(@addr, @port, Handler, method(:on_message))
+        @thread = nil
+        @running = false
+      end
+
+      def start
+        @server.attach(@loop)
+        @thread = Thread.new do
+          @running = true
+          @loop.run(0.5)
+          @running = false
+        end
+        @log.debug("starting counter server #{@addr}:#{@port}")
+        @mutex_hash.start
+        self
+      end
+
+      def stop
+        # This `sleep` for a test to wait for a `@loop` to begin to run
+        sleep 0.1
+        @server.close
+        @loop.stop if @running
+        @mutex_hash.stop
+        @thread.join if @thread
+        @log.debug("calling stop in counter server #{@addr}:#{@port}")
+      end
+
+      def on_message(data)
+        errors = Validator.request(data)
+        unless errors.empty?
+          return { 'id' => data['id'], 'data' => [], 'errors' => errors }
+        end
+
+        result = safe_run do
+          send(data['method'], data['params'], data['scope'], data['options'])
+        end
+        result.merge('id' => data['id'])
+      rescue => e
+        @log.error e.to_s
+      end
+
+      private
+
+      def establish(params, _scope, _options)
+        validator = Fluent::Counter::ArrayValidator.new(:empty, :scope)
+        valid_params, errors = validator.call(params)
+        res = Response.new(errors)
+
+        if scope = valid_params.first
+          new_scope = "#{@name}\t#{scope}"
+          res.push_data new_scope
+          @log.debug("Establish new key: #{new_scope}")
+        end
+
+        res.to_hash
+      end
+
+      def init(params, scope, options)
+        validator = Fluent::Counter::HashValidator.new(:empty, :name, :reset_interval)
+        valid_params, errors = validator.call(params)
+        res = Response.new(errors)
+        key_hash = valid_params.reduce({}) do |acc, vp|
+          acc.merge(Store.gen_key(scope, vp['name']) => vp)
+        end
+
+        do_init = lambda do |store, key|
+          begin
+            param = key_hash[key]
+            v = store.init(key, param, ignore: options['ignore'])
+            @log.debug("Create new key: #{param['name']}")
+            res.push_data v
+          rescue => e
+            res.push_error e
+          end
+        end
+
+        if options['random']
+          @mutex_hash.synchronize_keys(*(key_hash.keys), &do_init)
+        else
+          @mutex_hash.synchronize(*(key_hash.keys), &do_init)
+        end
+
+        res.to_hash
+      end
+
+      def delete(params, scope, options)
+        validator = Fluent::Counter::ArrayValidator.new(:empty, :key)
+        valid_params, errors = validator.call(params)
+        res = Response.new(errors)
+        keys = valid_params.map { |vp| Store.gen_key(scope, vp) }
+
+        do_delete = lambda do |store, key|
+          begin
+            v = store.delete(key)
+            @log.debug("delete a key: #{key}")
+            res.push_data v
+          rescue => e
+            res.push_error e
+          end
+        end
+
+        if options['random']
+          @mutex_hash.synchronize_keys(*keys, &do_delete)
+        else
+          @mutex_hash.synchronize(*keys, &do_delete)
+        end
+
+        res.to_hash
+      end
+
+      def inc(params, scope, options)
+        validate_param = [:empty, :name, :value]
+        validate_param << :reset_interval if options['force']
+        validator = Fluent::Counter::HashValidator.new(*validate_param)
+        valid_params, errors = validator.call(params)
+        res = Response.new(errors)
+        key_hash = valid_params.reduce({}) do |acc, vp|
+          acc.merge(Store.gen_key(scope, vp['name']) => vp)
+        end
+
+        do_inc = lambda do |store, key|
+          begin
+            param = key_hash[key]
+            v = store.inc(key, param, force: options['force'])
+            @log.debug("Increment #{key} by #{param['value']}")
+            res.push_data v
+          rescue => e
+            res.push_error e
+          end
+        end
+
+        if options['random']
+          @mutex_hash.synchronize_keys(*(key_hash.keys), &do_inc)
+        else
+          @mutex_hash.synchronize(*(key_hash.keys), &do_inc)
+        end
+
+        res.to_hash
+      end
+
+      def reset(params, scope, options)
+        validator = Fluent::Counter::ArrayValidator.new(:empty, :key)
+        valid_params, errors = validator.call(params)
+        res = Response.new(errors)
+        keys = valid_params.map { |vp| Store.gen_key(scope, vp) }
+
+        do_reset = lambda do |store, key|
+          begin
+            v = store.reset(key)
+            @log.debug("Reset #{key}'s' counter value")
+            res.push_data v
+          rescue => e
+            res.push_error e
+          end
+        end
+
+        if options['random']
+          @mutex_hash.synchronize_keys(*keys, &do_reset)
+        else
+          @mutex_hash.synchronize(*keys, &do_reset)
+        end
+
+        res.to_hash
+      end
+
+      def get(params, scope, _options)
+        validator = Fluent::Counter::ArrayValidator.new(:empty, :key)
+        valid_params, errors = validator.call(params)
+        res = Response.new(errors)
+
+        keys = valid_params.map { |vp| Store.gen_key(scope, vp) }
+        keys.each do |key|
+          begin
+            v = @store.get(key, raise_error: true)
+            @log.debug("Get counter value: #{key}")
+            res.push_data v
+          rescue => e
+            res.push_error e
+          end
+        end
+        res.to_hash
+      end
+
+      def safe_run
+        yield
+      rescue => e
+        {
+          'errors' => [InternalServerError.new(e).to_hash],
+          'data' => []
+        }
+      end
+
+      class Response
+        def initialize(errors = [], data = [])
+          @errors = errors
+          @data = data
+        end
+
+        def push_error(error)
+          @errors << error
+        end
+
+        def push_data(data)
+          @data << data
+        end
+
+        def to_hash
+          if @errors.empty?
+            { 'data' => @data }
+          else
+            errors = @errors.map do |e|
+              error = e.respond_to?(:to_hash) ? e : InternalServerError.new(e.to_s)
+              error.to_hash
+            end
+            { 'data' => @data, 'errors' => errors }
+          end
+        end
+      end
+    end
+
+    class Handler < Fluent::Counter::BaseSocket
+      def initialize(io, on_message)
+        super(io)
+        @on_message = on_message
+      end
+
+      def on_message(data)
+        res = @on_message.call(data)
+        packed_write res if res
+      end
+    end
+  end
+end

data/lib/fluent/counter/store.rb

@@ -0,0 +1,205 @@
+#
+# Fluentd
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+#
+
+require 'fluent/config'
+require 'fluent/counter/error'
+require 'fluent/plugin/storage_local'
+require 'fluent/time'
+
+module Fluent
+  module Counter
+    class Store
+      def self.gen_key(scope, key)
+        "#{scope}\t#{key}"
+      end
+
+      def initialize(opt = {})
+        @log = opt[:log] || $log
+
+        # Notice: This storage is not be implemented auto save.
+        @storage = Plugin.new_storage('local', parent: DummyParent.new(@log))
+        conf = if opt[:path]
+                 {'persistent' => true, 'path' => opt[:path] }
+               else
+                 {'persistent' => false }
+               end
+        @storage.configure(Fluent::Config::Element.new('storage', {}, conf, []))
+      end
+
+      # This class behaves as a configurable plugin for using in storage (OwnedByMixin).
+      class DummyParent
+        include Configurable
+
+        attr_reader :log
+
+        def initialize(log)
+          @log = log
+        end
+
+        def plugin_id
+          'dummy_parent_store'
+        end
+
+        def plugin_id_configured?
+          false
+        end
+
+        # storage_local calls PluginId#plugin_root_dir
+        def plugin_root_dir
+          nil
+        end
+      end
+
+      def start
+        @storage.load
+      end
+
+      def stop
+        @storage.save
+      end
+
+      def init(key, data, ignore: false)
+        ret = if v = get(key)
+                raise InvalidParams.new("#{key} already exists in counter") unless ignore
+                v
+              else
+                @storage.put(key, build_value(data))
+              end
+
+        build_response(ret)
+      end
+
+      def get(key, raise_error: false, raw: false)
+        ret = if raise_error
+                @storage.get(key) or raise UnknownKey.new("`#{key}` doesn't exist in counter")
+              else
+                @storage.get(key)
+              end
+        if raw
+          ret
+        else
+          ret && build_response(ret)
+        end
+      end
+
+      def key?(key)
+        !!@storage.get(key)
+      end
+
+      def delete(key)
+        ret = @storage.delete(key) or raise UnknownKey.new("`#{key}` doesn't exist in counter")
+        build_response(ret)
+      end
+
+      def inc(key, data, force: false)
+        value = data.delete('value')
+        init(key, data) if !key?(key) && force
+        v = get(key, raise_error: true, raw: true)
+        valid_type!(v, value)
+
+        v['total'] += value
+        v['current'] += value
+        t = EventTime.now
+        v['last_modified_at'] = [t.sec, t.nsec]
+        @storage.put(key, v)
+
+        build_response(v)
+      end
+
+      def reset(key)
+        v = get(key, raise_error: true, raw: true)
+        success = false
+        old_data = v.dup
+        now = EventTime.now
+        last_reset_at = EventTime.new(*v['last_reset_at'])
+
+        #  Does it need reset?
+        if (last_reset_at + v['reset_interval']) <= now
+          success = true
+          v['current'] = initial_value(v['type'])
+          t = [now.sec, now.nsec]
+          v['last_reset_at'] = t
+          v['last_modified_at'] = t
+          @storage.put(key, v)
+        end
+
+        {
+          'elapsed_time' => now - last_reset_at,
+          'success' => success,
+          'counter_data' => build_response(old_data)
+        }
+      end
+
+      private
+
+      def build_response(d)
+        {
+          'name' => d['name'],
+          'total' => d['total'],
+          'current' => d['current'],
+          'type' => d['type'],
+          'reset_interval' => d['reset_interval'],
+          'last_reset_at' => EventTime.new(*d['last_reset_at']),
+        }
+      end
+
+      # value is Hash. value requires these fileds.
+      # :name, :total, :current, :type, :reset_interval, :last_reset_at, :last_modified_at
+      def build_value(data)
+        type = data['type'] || 'numeric'
+        now = EventTime.now
+        t = [now.sec, now.nsec]
+
+        v = initial_value(type)
+
+        data.merge(
+          'type' => type,
+          'last_reset_at' => t,
+          'last_modified_at' => t,
+          'current' => v,
+          'total' => v,
+        )
+      end
+
+      def initial_value(type)
+        case type
+        when 'numeric', 'integer' then 0
+        when 'float' then 0.0
+        else raise InvalidParams.new('`type` should be integer, float, or numeric')
+        end
+      end
+
+      def valid_type!(v, value)
+        type = v['type']
+        return unless (type != 'numeric') && (type_str(value) != type)
+        raise InvalidParams.new("`type` is #{type}. You should pass #{type} value as a `value`")
+      end
+
+      def type_str(v)
+        case v
+        when Integer
+          'integer'
+        when Float
+          'float'
+        when Numeric
+          'numeric'
+        else
+          raise InvalidParams.new("`type` should be integer, float, or numeric")
+        end
+      end
+    end
+  end
+end