#! /usr/bin/env stap

global devices, reads, writes

/* kernel commit f3fa33acca9f0058157214800f68b10d8e71ab7a */
@define __disk_name(__request)
%(
    ( @choose_defined(@__request->rq_disk->disk_name, @__request->q->disk->disk_name ) )
%)

/* data collection: SCSI disk */
%{
#include <linux/blkdev.h>
#include <linux/blk-mq.h>
%}

function get_nr_sectors:long(rq:long)
%{ /* pure */
  STAP_RETVALUE = blk_rq_sectors((const struct request *)(long)STAP_ARG_rq);
%}


probe sd_prep_fn = module("sd_mod").function("sd_prep_fn") !,
		   kernel.function("sd_prep_fn")
{
  device = kernel_string(@__disk_name($rq))
  sector_size = @cast($q->queuedata, "scsi_device", "kernel<scsi/scsi_device.h>")->sector_size
  nr_sectors = get_nr_sectors($rq)

  _cmd_flags = $rq->cmd_flags
}

function scsi_cmd_to_rq(scmd) {
    return scmd - @cast_module_sizeof("kernel", "request")
}

probe sd_init_command = module("sd_mod").function("sd_init_command") !,
			kernel.function("sd_init_command")
{
    sector_size = @choose_defined($cmd, $SCpnt)->device->sector_size
    # Kernel commits aa8e25e5006aac52c943c84e9056ab488630ee19 2266a2def97ce11ec979b6c58a1b637a16eca7dd
    if (@defined(@choose_defined($cmd, $SCpnt)->request)) {
      device = kernel_string(@__disk_name(@choose_defined($cmd, $SCpnt)->request))
      nr_sectors = get_nr_sectors(@choose_defined($cmd, $SCpnt)->request)
      _cmd_flags = @choose_defined($cmd, $SCpnt)->request->cmd_flags
    } else {
      device = kernel_string(@__disk_name(@cast(scsi_cmd_to_rq(@choose_defined($cmd, $SCpnt)), "request", "kernel")))
      nr_sectors = get_nr_sectors(scsi_cmd_to_rq(@choose_defined($cmd, $SCpnt)))
      _cmd_flags = @cast(scsi_cmd_to_rq(@choose_defined($cmd, $SCpnt)), "request", "kernel")->cmd_flags
  }
}

probe sd_prep_fn !, sd_init_command
{
  devices[device] = 1
  if (_cmd_flags & 1)
    writes[device] <<< nr_sectors * sector_size
  else
    reads[device] <<< nr_sectors * sector_size
}

/* data collection: SCSI tape */
probe module("st").function("st_do_scsi").call !,
      kernel.function("st_do_scsi").call
{
  # Kernel commit 45938335d0a9773d65a82a7ca722bb76e4b997a8
  device = kernel_string(@choose_defined($STp->disk->disk_name, $STp->name))
  devices[device] = 1
  if ($direction)
    writes[device] <<< $bytes
  else
    reads[device] <<< $bytes
}
/*
 * Handle inlined version of function separately, since on some kernels
 * (3.10.0-196.el7.ppc64) the variables aren't available (bad
 * debuginfo). We want to handle this separately from the '.call'
 * version, so that we'll know if the variable names change.
 */
probe module("st").function("st_do_scsi").inline !,
      kernel.function("st_do_scsi").inline ?
{
  if (@defined($STp) && @defined($direction) && @defined($bytes)) {
    device = kernel_string($STp->disk->disk_name)
    devices[device] = 1
    if ($direction)
      writes[device] <<< $bytes
    else
      reads[device] <<< $bytes
  }
}


/* reporting */
global blksize = 512
global hdrcount
probe timer.s($1)
{
  if ((hdrcount % 10) == 0)
    printf("%9s %9s %9s %9s %9s %9s\n",
           "Device:", "tps", "blk_read/s", "blk_wrtn/s", "blk_read", "blk_wrtn")
  hdrcount++

  foreach (dev in devices) {
    rdcount = @count(reads[dev])
    wrcount = @count(writes[dev])
    tps = (rdcount + wrcount) * 100 / $1
    rdblkcount = rdcount ? @sum(reads[dev]) / blksize : 0
    wrblkcount =wrcount ? @sum(writes[dev]) / blksize : 0
    rdblkrate = rdblkcount * 100 / $1
    wrblkrate = wrblkcount * 100 / $1
    printf("%9s %6d.%02d %6d.%02d %6d.%02d %9d %9d\n",
           dev, tps / 100, tps % 100,
           rdblkrate / 100, rdblkrate % 100,
           wrblkrate / 100, wrblkrate % 100,
           rdblkcount, wrblkcount)
  }
  printf("\n")
  delete devices
  delete reads
  delete writes
}