scripts/sysmon/proc_metrics.py

# Copyright 2017 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Process metrics."""

from __future__ import absolute_import

from functools import partial
import logging

import psutil  # pylint: disable=import-error

from chromite.lib import metrics


logger = logging.getLogger(__name__)

_count_metric = metrics.GaugeMetric(
    'proc/count',
    description='Number of processes currently running.')
_cpu_percent_metric = metrics.GaugeMetric(
    'proc/cpu_percent',
    description='CPU usage percent of processes.')


def collect_proc_info():
  collector = _ProcessMetricsCollector()
  collector.collect()


class _ProcessMetricsCollector(object):
  """Class for collecting process metrics."""

  def __init__(self):
    self._metrics = [
        _ProcessMetric('autoserv',
                       test_func=_is_parent_autoserv),
        _ProcessMetric('curl',
                       test_func=partial(_is_process_name, 'curl')),
        _ProcessMetric('getty',
                       test_func=partial(_is_process_name, 'getty')),
        _ProcessMetric('gs_archive_server',
                       test_func=partial(_is_python_module,
                                         'gs_archive_server')),
        _ProcessMetric('gs_offloader',
                       test_func=partial(_is_process_name, 'gs_offloader.py')),
        _ProcessMetric('gsutil',
                       test_func=_is_gsutil),
        _ProcessMetric('java',
                       test_func=partial(_is_process_name, 'java')),
        _ProcessMetric('lxc-attach',
                       test_func=partial(_is_process_name, 'lxc-attach')),
        _ProcessMetric('lxc-start',
                       test_func=partial(_is_process_name, 'lxc-start')),
        _ProcessMetric('sshd',
                       test_func=partial(_is_process_name, 'sshd')),
        _ProcessMetric('swarming_bot',
                       test_func=_is_swarming_bot),
        _ProcessMetric('sysmon',
                       test_func=partial(_is_python_module,
                                         'chromite.scripts.sysmon')),
    ]
    self._other_metric = _ProcessMetric('other')

  def collect(self):
    for proc in psutil.process_iter():
      self._collect_proc(proc)
    self._flush()

  def _collect_proc(self, proc):
    for metric in self._metrics:
      if metric.add(proc):
        break
    else:
      self._other_metric.add(proc)

  def _flush(self):
    for metric in self._metrics:
      metric.flush()
    self._other_metric.flush()


class _ProcessMetric(object):
  """Class for gathering process metrics."""

  def __init__(self, process_name, test_func=lambda proc: True):
    """Initialize instance.

    process_name is used to identify the metric stream.

    test_func is a function called
    for each process.  If it returns True, the process is counted.  The
    default test is to count every process.
    """
    self._fields = {
        'process_name': process_name,
    }
    self._test_func = test_func
    self._count = 0
    self._cpu_percent = 0

  def add(self, proc):
    """Do metric collection for the given process.

    Returns True if the process was collected.
    """
    if not self._test_func(proc):
      return False
    self._count += 1
    self._cpu_percent += proc.cpu_percent()
    return True

  def flush(self):
    """Finish collection and send metrics."""
    _count_metric.set(self._count, fields=self._fields)
    self._count = 0
    _cpu_percent_metric.set(int(round(self._cpu_percent)), fields=self._fields)
    self._cpu_percent = 0


def _is_parent_autoserv(proc):
  """Return whether proc is a parent (not forked) autoserv process."""
  return _is_autoserv(proc) and not _is_autoserv(proc.parent())


def _is_autoserv(proc):
  """Return whether proc is an autoserv process."""
  # This relies on the autoserv script being run directly.  The script should
  # be named autoserv exactly and start with a shebang that is /usr/bin/python,
  # NOT /bin/env
  return _is_process_name('autoserv', proc)


def _is_python_module(module, proc):
  """Return whether proc is a process running a Python module."""
  cmdline = proc.cmdline()
  return (cmdline and
          cmdline[0].endswith('python') and
          cmdline[1:3] == ['-m', module])


def _is_process_name(name, proc):
  """Return whether process proc is named name."""
  return proc.name() == name


def _is_swarming_bot(proc):
  """Return whether proc is a Swarming bot.

  A swarming bot process is like '/usr/bin/python3.8 <bot-zip-path> start_bot'.
  """
  cmdline = proc.cmdline()
  return (len(cmdline) == 3 and
          cmdline[0].split('/')[-1].startswith('python') and
          cmdline[2] == 'start_bot')


def _is_gsutil(proc):
  """Return whether proc is gsutil."""
  cmdline = proc.cmdline()
  return (len(cmdline) >= 2 and
          cmdline[0] == 'python' and
          cmdline[1].endswith('gsutil'))