// This script tracks time between packet queue and xmit.
// The information is provided to userspace via procfs and are read
// using the stap_json PCP PMDA.

// ===========================
// ==== Tapset file start ====
// ===========================

global __json_prefix
global __json_metric_type, __json_metric_desc, __json_metric_pointer
global __json_array_metric_type, __json_array_metric_desc
global __json_array_metric_pointer
global __json_metric_output, __json_array_output

probe init
{
  __json_prefix = ""

  // If the user doesn't use any arrays, then the systemtap translator
  // doesn't know what type '__json_array_output' is. So, help it
  // out. 
  __json_array_output[""] = 1
  delete __json_array_output
}

function json_set_prefix:long(prefix:string)
{
  __json_prefix = prefix
}

function json_add_numeric_metric:long(name:string, description:string)
{
  if (name in __json_metric_type)
    error(sprintf("Metric '%s' already exists", name))
  __json_metric_type[name] = "integer"
  __json_metric_desc[name] = description
  # FIXME: Do we need to validate the name? For instance, what if it
  # had a double-quote character in it?
  __json_metric_pointer[name] = "/" . name
  return 0
}

function json_add_string_metric:long(name:string, description:string)
{
  if (name in __json_metric_type)
    error(sprintf("Metric '%s' already exists", name))
  __json_metric_type[name] = "string"
  __json_metric_desc[name] = description
  __json_metric_pointer[name] = "/" . name
  return 0
}

function json_add_array:long(name:string, description:string)
{
  if (name in __json_metric_type)
    error(sprintf("Metric '%s' already exists", name))
  __json_metric_type[name] = "array"
  __json_metric_desc[name] = description
  __json_metric_pointer[name] = "/" . name

  # Go ahead and add "__id", which is the array index.
  json_add_array_string_metric(name, "__id", "")
  return 0
}

function json_add_array_numeric_metric:long(array_name:string,
	metric_name:string, metric_description:string)
{
  if ([array_name, metric_name] in __json_array_metric_type)
    error(sprintf("Array metric '%s' already exists in array %s", metric_name,
		  array_name))
  __json_array_metric_pointer[array_name, metric_name]
    = sprintf("/%s", metric_name)
  __json_array_metric_type[array_name, metric_name] = "integer"
  __json_array_metric_desc[array_name, metric_name] = metric_description
  return 0
}

function json_add_array_string_metric:long(array_name:string,
	metric_name:string, metric_description:string)
{
  if ([array_name, metric_name] in __json_array_metric_type)
    error(sprintf("Array metric '%s' already exists in array %s", metric_name,
		  array_name))
  __json_array_metric_pointer[array_name, metric_name]
    = sprintf("/%s", metric_name)
  __json_array_metric_type[array_name, metric_name] = "string"
  __json_array_metric_desc[array_name, metric_name] = metric_description
  return 0
}

@define __json_output_metric_metadata(indent_str, name, pointer, type, description)
%(
  $value .= sprintf("%s{\n", @indent_str)
  $value .= sprintf("%s  \"name\": \"%s\",\n", @indent_str, @name)
  $value .= sprintf("%s  \"pointer\": \"%s\",\n", @indent_str, @pointer)
  $value .= sprintf("%s  \"type\": \"%s\"", @indent_str, @type)
  if (strlen(@description) > 0) {
    $value .= sprintf(",\n%s  \"description\": \"%s\"", @indent_str,
		      @description)
  }
  $value .= sprintf("\n%s}", @indent_str)
%)

@define __json_output_array_metadata(indent_str, name)
%(
  $value .= sprintf("%s\{\n", @indent_str)
  $value .= sprintf("%s  \"name\": \"%s\",\n", @indent_str, @name)
  $value .= sprintf("%s  \"pointer\": \"%s\",\n", @indent_str,
		    __json_metric_pointer[@name])
  $value .= sprintf("%s  \"type\": \"array\",\n", @indent_str)
  if (strlen(__json_metric_desc[@name]) > 0)
    $value .= sprintf("%s  \"description\": \"%s\",\n", @indent_str,
		      __json_metric_desc[@name])
  $value .= sprintf("%s  \"index\": \"/__id\",\n", @indent_str)
  $value .= sprintf("%s  \"metrics\": [\n", @indent_str)
  __array_comma_needed = 0
  foreach ([__array_name, __metric_name] in __json_array_metric_type) {
    if (@name == __array_name) {
      # Skip '__id', the array index.
      if (__metric_name == "__id")
        continue

      if (__array_comma_needed)
        $value .= ",\n"
      __array_comma_needed = 1    

      __subindent_str = @indent_str . "    "
      __subpointer = __json_array_metric_pointer[__array_name, __metric_name]
      __subtype = __json_array_metric_type[__array_name, __metric_name]
      __subdesc = __json_array_metric_desc[__array_name, __metric_name]

      @__json_output_metric_metadata(__subindent_str, __metric_name,
				     __subpointer, __subtype, __subdesc)
    }
  }
  $value .= sprintf("\n%s  ]\n", @indent_str)
  $value .= sprintf("%s}", @indent_str)
%)

@define json_output_metadata
%(
  # Note: This is the "pretty-printed" version, intended to be read by
  # humans. We could remove the whitespace and newlines if we wanted
  # to make the output shorter (but less readable).
  #
  # Note 2: Note that we have to break this long string into more than
  # 1 assignment since we're bumping up against MAXSTRINGLEN. Procfs
  # $value can hold more than MAXSTRINGLEN because of the
  # '.maxsize(N)' parameter.
  $value = "{\n"

  if (__json_prefix != "")
    $value .= sprintf("  \"prefix\": \"%s\",\n", __json_prefix)

  $value .= "  \"metrics\": [\n"

  __comma_needed = 0
  foreach (__name in __json_metric_type) {
    if (__comma_needed)
      $value .= ",\n"
    __comma_needed = 1    

    if (__json_metric_type[__name] != "array") {
      @__json_output_metric_metadata("    ", __name,
				     __json_metric_pointer[__name],
				     __json_metric_type[__name],
				     __json_metric_desc[__name])
    }
    else {
      @__json_output_array_metadata("    ", __name)
    }
  }

  $value .=
    "\n"
    "  ]\n"
    "}\n"
  # FIXME: optional prefix value would go right after the metrics array
%)

# NOTE: This is the "pretty-printed" version of the data, intended
# to be read by humans. We could remove the whitespace and newlines
# if we wanted to make the output shorter (but less readable).
@define json_output_data_start
%(
  __comma_needed = 0
  $value =
    "{\n"
%)

# Make sure we don't try to output the same metric twice in the same
# data fetch.
@define __json_output_check(name)
%(
  if (@name in __json_metric_output)
    error(sprintf("Metric '%s' already output", @name))
  __json_metric_output[@name] = 1
%)

# Make sure we don't try to output the same array index twice in the same
# data fetch.
@define __json_output_array_check(array_index)
%(
  if (@array_index in __json_array_output)
    error(sprintf("Array index '%s' already output for array metric %s",
		  @array_index, __json_array_started))
  __json_array_output[@array_index] = 1
%)

# Output a string value.
@define json_output_string_value(name, value)
%(
  @__json_output_check(@name)
  @__json_output_array_end
  if (__comma_needed)
    $value .= ",\n"
  __comma_needed = 1    
  $value .= sprintf("  \"%s\": \"%s\"", @name, @value)
%)

# Output a numeric value.
@define json_output_numeric_value(name, value)
%(
  @__json_output_check(@name)
  @__json_output_array_end
  if (__comma_needed)
    $value .= ",\n"
  __comma_needed = 1    
  $value .= sprintf("  \"%s\": %d", @name, @value)
%)

# Output a string value for an array.
@define json_output_array_string_value(array_name, array_index, metric_name, value)
%(
  @__json_output_array_start(@array_name, @array_index)
  if (__comma_needed)
    $value .= ",\n"
  __comma_needed = 1    
  $value .= sprintf("      \"%s\": \"%s\"", @metric_name, @value)
%)

# Output a numeric value for an array.
@define json_output_array_numeric_value(array_name, array_index, metric_name,
					value)
%(
  @__json_output_array_start(@array_name, @array_index)
  if (__comma_needed)
    $value .= ",\n"
  __comma_needed = 1    
  $value .= sprintf("      \"%s\": %d", @metric_name, @value)
%)

# Handle the details of starting the output of an array.
@define __json_output_array_start(array_name, array_index)
%(
  if (__json_array_started != @array_name) {
    @__json_output_check(@array_name)
    @__json_output_array_end
    if (__comma_needed)
      $value .= ",\n"
    __comma_needed = 1    
    $value .= sprintf("  \"%s\": [\n", @array_name)
    $value .= "    {\n"
    __json_array_started = @array_name
  }
  if (__json_array_index_started != @array_index) {
    @__json_output_array_check(@array_index)
    if (__json_array_index_started != "") {
      $value .=
        "\n"
	"    },\n"
	"    {\n"
    }
    __json_array_index_started = @array_index
    $value .= sprintf("      \"__id\": \"%s\"", @array_index)
  }
%)

# Handle the details of finishing the output of an array.
@define __json_output_array_end
%(
  if (__json_array_started != "") {
    $value .=
      "\n"
      "    }\n"
      "  ]"
    __json_array_started = ""
    __json_array_index_started = ""
    delete __json_array_output
  }
%)

# Finish outputting data.
@define json_output_data_end
%(
  @__json_output_array_end
  $value .=
    "\n"
    "}\n"
  __comma_needed = 0
  delete __json_metric_output
%)

probe procfs("metadata").read.maxsize(8192)
{
  @json_output_metadata
}

probe json_data = procfs("data").read.maxsize(8192)
{
}

// =========================
// ==== Tapset file end ====
// =========================


global net_devices
global read_count

probe json_data
{
  @json_output_data_start
  @json_output_string_value("xstring", "testing, 1, 2, 3")
  @json_output_numeric_value("read_count", read_count)
  read_count++

  foreach (dev in net_devices) {
    if (@count(skb_queue_t[dev])) {
      @json_output_array_numeric_value("net_xmit_data", dev, "xmit_count",
				       @sum(skb_queue_t[dev]))
      @json_output_array_numeric_value("net_xmit_data", dev, "xmit_latency",
				       @count(skb_queue_t[dev]))
    }
    else {
      @json_output_array_numeric_value("net_xmit_data", dev, "xmit_count", 0)
      @json_output_array_numeric_value("net_xmit_data", dev, "xmit_latency", 0)
    }
  }

  # Add dummy values just to test the interface
  @json_output_array_numeric_value("dummy_array", "0", "dummy1", 1)
  @json_output_array_string_value("dummy_array", "0", "dummy2", "abc")
  @json_output_array_numeric_value("dummy_array", "1", "dummy1", 2)
  @json_output_array_string_value("dummy_array", "1", "dummy2", "def")
  @json_output_array_numeric_value("dummy_array", "2", "dummy1", 3)
  @json_output_array_string_value("dummy_array", "2", "dummy2", "ghi")
  @json_output_string_value("dummy2", "dummy2")
  @json_output_data_end
}

// Set up the metrics
probe begin
{
  // fallback instance device "eth0" if none specified
  if (argc == 0) {
    argv[1] = "eth0"
    argc = 1
  }
  
  // remember all the network devices
  for (i = 1; i <= argc; i++) {
    dev = argv[i]
    net_devices[dev] = i - 1
  }

  // Set the prefix to be used instead of the module name.
  json_set_prefix("net_xmit")

  // Add the metrics
  json_add_string_metric("xstring", "Test string")
  json_add_numeric_metric("read_count", "Times values read")

  json_add_array("net_xmit_data", "Network transmit data indexed by ethernet device")
  json_add_array_numeric_metric("net_xmit_data", "xmit_count", "number of packets for xmit device")
  json_add_array_numeric_metric("net_xmit_data", "xmit_latency", "sum of latency for xmit device")

  // Add some dummy metrics just to test the tapset.
  json_add_array("dummy_array", "")
  json_add_array_numeric_metric("dummy_array", "dummy1", "")
  json_add_array_string_metric("dummy_array", "dummy2", "")
  json_add_string_metric("dummy2", "Test string")
}

// probes to track the information

global skb_queue_start_t, skb_queue_t

probe kernel.trace("net_dev_queue") {
  skb_queue_start_t[$skb] = gettimeofday_ns();
}

probe kernel.trace("net_dev_start_xmit"), kernel.trace("net_dev_xmit") {
  t = gettimeofday_ns();
  st = skb_queue_start_t[$skb]
  if (st){
    skb_queue_t[kernel_string($dev->name)] <<< t - st
    delete skb_queue_start_t[$skb]
  }
}