# Licensed to Elasticsearch B.V. under one or more contributor
# license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright
# ownership. Elasticsearch B.V. licenses this file to you 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.

# Logic from elasticsearch/core/src/main/java/org/elasticsearch/monitor/os/OsProbe.java
# Move to ruby to remove any existing dependency
module LogStash module Instrument module PeriodicPoller
  class Cgroup
    include LogStash::Util::Loggable
    class Override
      attr_reader :key, :value

      def initialize(key)
        @key = key
        @value = java.lang.System.getProperty(@key)
      end

      def nil?
        value.nil?
      end

      def override(other)
        nil? ? other : value
      end
    end

    ## `/proc/self/cgroup` contents look like this
    # cgroupv1 (per-controller hierarchies):
    #   5:cpu,cpuacct:/
    #   4:cpuset:/
    #   2:net_cls,net_prio:/
    #   0::/user.slice/user-1000.slice/session-932.scope
    # cgroupv2 (unified hierarchy):
    #   0::/path
    #
    # For v1, we find the controller and path from lines matching N:controller:/path.
    # We assume these symlinks exist:
    #   `/sys/fs/cgroup/cpu` -> `/sys/fs/cgroup/cpu,cpuacct`
    #   `/sys/fs/cgroup/cpuacct` -> `/sys/fs/cgroup/cpu,cpuacct`
    # For v2, the 0::/path line identifies the cgroup path under the unified
    # hierarchy at /sys/fs/cgroup. Data files (cpu.stat, cpu.max) live there.

    CGROUP_FILE = "/proc/self/cgroup"
    CPUACCT_DIR = "/sys/fs/cgroup/cpuacct"
    CPU_DIR = "/sys/fs/cgroup/cpu"
    CRITICAL_PATHS = [CGROUP_FILE, CPUACCT_DIR, CPU_DIR]
    CGROUP_V2_DIR = "/sys/fs/cgroup"

    CONTROLLER_CPUACCT_LABEL = "cpuacct"
    CONTROLLER_CPU_LABEL = "cpu"

    class CGroupResources
      CONTROL_GROUP_RE = Regexp.compile("\\d+:([^:,]+(?:,[^:,]+)?):(/.*)")
      CONTROLLER_SEPARATOR_RE = ","

      def cgroup_available?
        # don't cache to ivar, in case the files are mounted after logstash starts??
        CRITICAL_PATHS.all? {|path| ::File.exist?(path)}
      end

      def controller_groups
        response = {}
        IO.readlines(CGROUP_FILE).each do |line|
          matches = CONTROL_GROUP_RE.match(line)
          next if matches.nil?
          # multiples controls, same hierarchy
          controllers = matches[1].split(CONTROLLER_SEPARATOR_RE)
          controllers.each do |controller|
            case controller
            when CONTROLLER_CPU_LABEL
              response[controller] = CpuResource.new(matches[2])
            when CONTROLLER_CPUACCT_LABEL
              response[controller] = CpuAcctResource.new(matches[2])
            else
              response[controller] = UnimplementedResource.new(controller, matches[2])
            end
          end
        end
        response
      end
    end

    class CGroupV2Resources
      CGROUP_V2_RE = Regexp.compile("^0::(/.*)")

      def cgroup_available?
        path = resolve_v2_path
        return false if path.nil?
        cpu_override = Override.new("ls.cgroup.cpu.path.override")
        resolved = cpu_override.override(path)
        ::File.exist?(::File.join(CGROUP_V2_DIR, resolved, "cpu.stat"))
      end

      def controller_groups
        path = resolve_v2_path
        return {} if path.nil?
        {
          CONTROLLER_CPU_LABEL => CpuResourceV2.new(path),
          CONTROLLER_CPUACCT_LABEL => CpuAcctResourceV2.new(path)
        }
      end

      private

      def resolve_v2_path
        @v2_path ||= read_v2_path
      end

      def read_v2_path
        return nil unless ::File.exist?(CGROUP_FILE)
        IO.readlines(CGROUP_FILE).each do |line|
          match = CGROUP_V2_RE.match(line.strip)
          return match[1] if match
        end
        nil
      end
    end

    module ControllerResource
      attr_reader :base_path, :override, :offset_path

      def implemented?
        true
      end
      private
      def common_initialize(base, override_key, original_path)
        @base_path = base
        # override is needed here for the logging statements
        @override = Override.new(override_key)
        @offset_path = @override.override(original_path)
        @procs = {}
        @procs[:read_int] = lambda {|path| IO.readlines(path).first.to_i }
        @procs[:read_lines] = lambda {|path| IO.readlines(path) }
      end

      def call_if_file_exists(call_key, file, not_found_value)
        path = ::File.join(@base_path, @offset_path, file)
        if ::File.exist?(path)
          @procs[call_key].call(path)
        else
          message = "File #{path} cannot be found, "
          if override.nil?
            message.concat("try providing an override '#{override.key}' in the Logstash JAVA_OPTS environment variable")
          else
            message.concat("even though the '#{override.key}' override is: '#{override.value}'")
          end
          logger.debug(message)
          not_found_value
        end
      end
    end

    class CpuAcctResource
      include LogStash::Util::Loggable
      include ControllerResource
      def initialize(original_path)
        common_initialize(CPUACCT_DIR, "ls.cgroup.cpuacct.path.override", original_path)
      end

      def to_hash
        {:control_group => offset_path, :usage_nanos => cpuacct_usage}
      end
      private
      def cpuacct_usage
        call_if_file_exists(:read_int, "cpuacct.usage", -1)
      end
    end

    class CpuAcctResourceV2
      include LogStash::Util::Loggable
      include ControllerResource
      def initialize(original_path)
        common_initialize(CGROUP_V2_DIR, "ls.cgroup.cpuacct.path.override", original_path)
      end

      def to_hash
        {:control_group => offset_path, :usage_nanos => cpuacct_usage}
      end
      private
      def cpuacct_usage
        lines = call_if_file_exists(:read_lines, "cpu.stat", [])
        lines.each do |line|
          fields = line.split(/\s+/)
          return fields[1].to_i * 1000 if fields.first == "usage_usec"
        end
        -1
      end
    end

    class CpuResource
      include LogStash::Util::Loggable
      include ControllerResource
      def initialize(original_path)
        common_initialize(CPU_DIR, "ls.cgroup.cpu.path.override", original_path)
      end

      def to_hash
        {
          :control_group => offset_path,
          :cfs_period_micros => cfs_period_us,
          :cfs_quota_micros => cfs_quota_us,
          :stat => build_cpu_stats_hash
        }
      end
      private
      def cfs_period_us
        call_if_file_exists(:read_int, "cpu.cfs_period_us", -1)
      end

      def cfs_quota_us
        call_if_file_exists(:read_int, "cpu.cfs_quota_us", -1)
      end

      def build_cpu_stats_hash
        stats = CpuStats.new
        lines = call_if_file_exists(:read_lines, "cpu.stat", [])
        stats.update(lines)
        stats.to_hash
      end
    end

    class CpuResourceV2
      include LogStash::Util::Loggable
      include ControllerResource
      def initialize(original_path)
        common_initialize(CGROUP_V2_DIR, "ls.cgroup.cpu.path.override", original_path)
      end

      def to_hash
        {
          :control_group => offset_path,
          :cfs_period_micros => cfs_period_us,
          :cfs_quota_micros => cfs_quota_us,
          :stat => build_cpu_stats_hash
        }
      end
      private
      def cfs_period_us
        read_cpu_max[1]
      end

      def cfs_quota_us
        read_cpu_max[0]
      end

      def read_cpu_max
        @cpu_max ||= begin
          line = call_if_file_exists(:read_lines, "cpu.max", []).first
          if line.nil?
            [-1, -1]
          else
            parts = line.split(/\s+/)
            quota = parts[0] == "max" ? -1 : parts[0].to_i
            period = parts[1].to_i
            [quota, period]
          end
        end
      end

      def build_cpu_stats_hash
        stats = CpuStatsV2.new
        lines = call_if_file_exists(:read_lines, "cpu.stat", [])
        stats.update(lines)
        stats.to_hash
      end
    end

    class UnimplementedResource
      attr_reader :controller, :original_path

      def initialize(controller, original_path)
        @controller, @original_path = controller, original_path
      end

      def implemented?
        false
      end
    end

    class CpuStats
      def initialize
        @number_of_elapsed_periods = -1
        @number_of_times_throttled = -1
        @time_throttled_nanos = -1
      end

      def update(lines)
        lines.each do |line|
          fields = line.split(/\s+/)
          next unless fields.size > 1
          case fields.first
          when "nr_periods" then @number_of_elapsed_periods = fields[1].to_i
          when "nr_throttled" then @number_of_times_throttled = fields[1].to_i
          when "throttled_time" then @time_throttled_nanos = fields[1].to_i
          end
        end
      end

      def to_hash
        {
          :number_of_elapsed_periods => @number_of_elapsed_periods,
          :number_of_times_throttled => @number_of_times_throttled,
          :time_throttled_nanos => @time_throttled_nanos
        }
      end
    end

    class CpuStatsV2
      def initialize
        @number_of_elapsed_periods = -1
        @number_of_times_throttled = -1
        @time_throttled_nanos = -1
      end

      def update(lines)
        lines.each do |line|
          fields = line.split(/\s+/)
          next unless fields.size > 1
          case fields.first
          when "nr_periods" then @number_of_elapsed_periods = fields[1].to_i
          when "nr_throttled" then @number_of_times_throttled = fields[1].to_i
          when "throttled_usec" then @time_throttled_nanos = fields[1].to_i * 1000
          end
        end
      end

      def to_hash
        {
          :number_of_elapsed_periods => @number_of_elapsed_periods,
          :number_of_times_throttled => @number_of_times_throttled,
          :time_throttled_nanos => @time_throttled_nanos
        }
      end
    end

    CGROUP_RESOURCES = CGroupResources.new
    CGROUP_V2_RESOURCES = CGroupV2Resources.new

    class << self
      def get_all
        resolve_resources unless @resolved
        return nil if @active_resources.nil?

        groups = @active_resources.controller_groups

        if groups.empty?
          logger.debug("#{@active_label}: no controllers found") unless @logged_empty
          @logged_empty = true
          return
        end

        cgroups_stats = {}
        groups.each do |name, controller|
          next unless controller.implemented?
          cgroups_stats[name.to_sym] = controller.to_hash
        end
        cgroups_stats.empty? ? nil : cgroups_stats
      rescue => e
        logger.debug("Error, cannot retrieve #{@active_label} cgroups information", :exception => e.class.name, :message => e.message, :backtrace => e.backtrace.take(4)) if logger.debug?
        nil
      end

      def get
        get_all
      end

      private

      def resolve_resources
        @resolved = true
        if CGROUP_RESOURCES.cgroup_available?
          @active_resources = CGROUP_RESOURCES
          @active_label = "cgroupv1"
          logger.debug("using cgroupv1")
        elsif CGROUP_V2_RESOURCES.cgroup_available?
          @active_resources = CGROUP_V2_RESOURCES
          @active_label = "cgroupv2"
          logger.debug("using cgroupv2")
        else
          @active_resources = nil
          logger.debug("no cgroup support detected (neither v1 nor v2)")
        end
      end
    end
  end
end end end
