#!/usr/bin/env python
# Copyright (C) 2019 The Android Open Source Project
#
# 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.

import argparse
import logging
import os
import sys
import subprocess
import tempfile

ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ADB_PATH = os.path.join(ROOT_DIR, 'buildtools/android_sdk/platform-tools/adb')

TEMPLATED_PERFETTO_CFG = '''
buffers {{
  size_kb: 65536
  fill_policy: RING_BUFFER
}}
data_sources {{
  config {{
    name: "linux.ftrace"
    target_buffer: 0
    ftrace_config {{
      ftrace_events: "sched_switch"
      buffer_size_kb: {buffer_size_kb}
      drain_period_ms: {drain_period_ms}
    }}
  }}
}}
duration_ms: 15000
'''

PERFETTO_PER_CPU_BUFFER_SIZE_KB_PARAMS = [
  128,
  256,
  512,
  1 * 1024, # 1 MB
  2 * 1024, # 2 MB
  4 * 1024, # 4 MB
]

PERFETTO_DRAIN_RATE_MS_PARAMS = [
  100,
  240,
  500,
  1 * 1000, # 1s
  2 * 1000, # 2s
  5 * 1000, # 5s
]

BUSY_THREADS_NUM_THREADS_PARAMS = [
  8,
  32,
  128,
]

BUSY_THREADS_DUTY_CYCLE_PARAMS = [
  10,
  100,
]

BUSY_THREADS_PERIOD_US_PARAMS = [
  500,
  1 * 1000,  # 1 ms
  10 * 1000, # 10 ms
]

TRACE_PROCESSOR_QUERY = """
SELECT
  a.value as num_sched,
  b.value as num_overrun
FROM (
  SELECT COUNT(*) as value
  FROM sched
) as a, (
  SELECT SUM(value) as value
  FROM stats
  WHERE name = 'ftrace_cpu_overrun_end'
) as b
"""

def AdbArgs(*args):
  cmd = [ADB_PATH] + list([str(x) for x in args])
  logging.debug('> adb ' + ' '.join([str(x) for x in args]))
  return cmd

def AdbCall(*args):
  return subprocess.check_output(AdbArgs(*args)).decode('utf-8').rstrip()

def SingleTraceRun(out_dir,
                   prio_name,
                   buffer_size_kb,
                   drain_rate_ms,
                   num_threads,
                   duty_cycle,
                   period_us):
  busy_threads_args = AdbArgs(
    'shell',
    '/data/local/tmp/busy_threads',
    '--threads={}'.format(num_threads),
    '--duty_cycle={}'.format(duty_cycle),
    '--period_us={}'.format(period_us)
  )
  perfetto_args = AdbArgs(
    'shell',
    'perfetto',
    '--txt',
    '-c',
    '-',
    '-o',
    '-'
  )

  # Create a file object to read the trace into.
  with tempfile.NamedTemporaryFile() as trace_file:
    logging.info("Starting trace with parameters ({}, {}, {}, {}, {}, {})"
      .format(prio_name,
              buffer_size_kb,
              drain_rate_ms,
              num_threads,
              duty_cycle,
              period_us))

    # Start the busy threads running.
    busy_threads_handle = subprocess.Popen(busy_threads_args)

    # Start the Perfetto trace.
    perfetto_handle = subprocess.Popen(
      perfetto_args, stdin=subprocess.PIPE, stdout=trace_file)

    # Create the config with the parameters
    config = TEMPLATED_PERFETTO_CFG.format(
      buffer_size_kb=buffer_size_kb, drain_period_ms=drain_rate_ms)

    # Send the config to the Perfetto binary and wait for response.
    perfetto_handle.stdin.write(config.encode())
    perfetto_handle.stdin.close()
    perfetto_ret = perfetto_handle.wait()

    # Stop busy threads from running.
    busy_threads_handle.terminate()

    # Return any errors from Perfetto.
    if perfetto_ret:
      raise subprocess.CalledProcessError(
        cmd=perfetto_args, returncode=perfetto_ret)

    # TODO(lalitm): allow trace processor to take the query file from stdin
    # to prevent this hack from being required.
    with tempfile.NamedTemporaryFile() as trace_query_file:
      trace_query_file.file.write(TRACE_PROCESSOR_QUERY.encode())
      trace_query_file.file.flush()

      # Run the trace processor on the config.
      tp_path = os.path.join(out_dir, 'trace_processor_shell')
      tp_out = subprocess.check_output(
        [tp_path, '-q', trace_query_file.name, trace_file.name])

      # Get the CSV output from trace processor (stripping the header).
      [num_sched, num_overrun] = str(tp_out).split('\n')[1].split(',')

    # Print the row to stdout.
    sys.stdout.write('"{}",{},{},{},{},{},{},{}\n'
      .format(prio_name,
              buffer_size_kb,
              drain_rate_ms,
              num_threads,
              duty_cycle,
              period_us,
              num_sched,
              num_overrun))

def SinglePriorityRun(out_dir, prio_name):
  for buffer_size_kb in PERFETTO_PER_CPU_BUFFER_SIZE_KB_PARAMS:
    for drain_rate_ms in PERFETTO_DRAIN_RATE_MS_PARAMS:
      for num_threads in BUSY_THREADS_NUM_THREADS_PARAMS:
        for duty_cycle in BUSY_THREADS_DUTY_CYCLE_PARAMS:
          for period_us in BUSY_THREADS_PERIOD_US_PARAMS:
            SingleTraceRun(
              out_dir,
              prio_name,
              buffer_size_kb,
              drain_rate_ms,
              num_threads,
              duty_cycle,
              period_us
            )

def CycleTracedAndProbes():
  AdbCall('shell', 'stop', 'traced')
  AdbCall('shell', 'stop', 'traced_probes')
  AdbCall('shell', 'start', 'traced')
  AdbCall('shell', 'start', 'traced_probes')
  AdbCall('shell', 'sleep', '5')
  traced_pid = AdbCall('shell', 'pidof', 'traced')
  probes_pid = AdbCall('shell', 'pidof', 'traced_probes')
  assert(traced_pid is not None and probes_pid is not None)
  return (traced_pid, probes_pid)

def Main():
  parser = argparse.ArgumentParser()
  parser.add_argument('linux_out_dir', help='out/android/')
  parser.add_argument('android_out_dir', help='out/android/')
  args = parser.parse_args()

  # Root ourselves on the device.
  logging.info('Waiting for device and rooting ...')
  AdbCall('wait-for-device')
  AdbCall('root')
  AdbCall('wait-for-device')

  # Push busy threads to device
  busy_threads_path = os.path.join(args.android_out_dir, 'busy_threads')
  AdbCall('shell', 'rm', '-rf', '/data/local/tmp/perfetto_load_test')
  AdbCall('shell', 'mkdir', '/data/local/tmp/perfetto_load_test')
  AdbCall('push', busy_threads_path, '/data/local/tmp/perfetto_load_test/')

  # Stop and start traced and traced_probes
  (traced_pid, probes_pid) = CycleTracedAndProbes()

  # Print the header for csv.
  sys.stdout.write('"{}","{}","{}","{}","{}","{}","{}","{}"\n'
      .format('prio_name',
              'buffer_size_kb',
              'drain_rate_ms',
              'num_threads',
              'duty_cycle',
              'period_us',
              'num_sched',
              'num_overrun'))

  # First, do a single run in all configurations without changing prio.
  SinglePriorityRun(args.linux_out_dir, 'Default')

  # Stop and start traced and traced_probes
  (traced_pid, probes_pid) = CycleTracedAndProbes()

  # Setup the nice values and check them.
  AdbCall('shell', 'renice', '-n', '-19', '-p', traced_pid)
  AdbCall('shell', 'renice', '-n', '-19', '-p', probes_pid)
  logging.debug(AdbCall('shell', 'cat', '/proc/{}/stat'.format(traced_pid)))
  logging.debug(AdbCall('shell', 'cat', '/proc/{}/stat'.format(probes_pid)))

  # Do the run.
  SinglePriorityRun(args.linux_out_dir, '-19 Nice')

  # Stop and start traced and traced_probes
  (traced_pid, probes_pid) = CycleTracedAndProbes()

  # Then do a run with FIFO scheduling for traced and traced_probes.
  AdbCall('shell', 'chrt', '-f', '-p', traced_pid, '99')
  AdbCall('shell', 'chrt', '-f', '-p', probes_pid, '99')
  logging.debug(AdbCall('shell', 'chrt', '-p', traced_pid))
  logging.debug(AdbCall('shell', 'chrt', '-p', probes_pid))

  # Do the run.
  SinglePriorityRun(args.linux_out_dir, 'FIFO')

  # Stop and start traced and traced_probes
  (traced_pid, probes_pid) = CycleTracedAndProbes()

  # Cleanup any pushed files, priorities etc.
  logging.info("Cleaning up test")
  AdbCall('shell', 'rm', '-rf', '/data/local/tmp/perfetto_load_test')
  logging.debug(AdbCall('shell', 'cat', '/proc/{}/stat'.format(traced_pid)))
  logging.debug(AdbCall('shell', 'cat', '/proc/{}/stat'.format(probes_pid)))
  AdbCall('unroot')

  return 0

if __name__ == '__main__':
  logging.basicConfig(level=logging.INFO)
  sys.exit(Main())