#!/usr/bin/ruby
# Observe hardware ethernet statistics in readable form.
# (c) 2014 <abc@telekom.ru>
# License: GPL.

require 'getoptlong'
require 'set'
require 'pp'

@interval = 3    # update interval
@color    = true # red color for errors
@showall  = true # show bytes stat
@qlen     = true # show queue len stat

GetoptLong.new(
      ["--help",    "-h", GetoptLong::NO_ARGUMENT],
      ["--bytes",   "-b", GetoptLong::NO_ARGUMENT],
      ["--errors",  "-e", GetoptLong::NO_ARGUMENT],
      ["--nocolor", "-C", GetoptLong::NO_ARGUMENT],
      ["--qlen",    "-q", GetoptLong::NO_ARGUMENT]
).each do |opt, arg|
    case opt
    when '--help'
      puts "ethstat [-options] [seconds]"
      puts "  Periodically show ethernet stats in a table form."
      puts "  All output stats is per-second rates.  Options:"
      puts "    --bytes   -b  show bytes stat, not just packets."
      puts "    --errors  -e  show only errors stats."
      puts "    --nocolor -C  disable showing colors."
      puts "    --qlen    -q  disable showing active queue lengths."
      exit 0
    when '--bytes'
      @bytes = true
    when '--errors'
      @showall = false
    when '--nocolor'
      @color = false
    when '--qlen'
      @qlen = !@qlen
    end
end
if ARGV[0].to_f > 0
  @interval = ARGV[0].to_f
end

class String
  def to_nat
    (self =~ /^\d+$/)? self.to_i : self
  end
end

class Array
  def sub(older)
    zip(older).map {|a,b| a - b}
  end
  def add!(o)
    o.each_with_index {|a, i| self[i] += a}
  end
  def div(other)
    if other.kind_of? Array
      zip(other).map {|a,b| a.to_f / b.to_f rescue 0}
    else
      map{|e| e / other rescue 0}
    end
  end
  def mul(m)
    map{|e| e * m}
  end
  def to_nat
    map {|e| e.to_nat}
  end
  def sort_natural
    sort_by {|e| e.scan(/[a-z]+|\d+/).to_nat}
  end
  def bubble_up(v)
    insert(0, delete(v)) if include?(v)
  end
end

def stat_reinit
  @st ||= Hash.new {|h,k| h[k] = Hash.new}
  # [device_name, :queue_stat_name][:tag] => [q0, q2, ... per queue stat]

  @qstats = Hash.new {|h,k| h[k] = Set.new}
  # device_name => [:queue_stat_names...]
  # to list valid queue stats as secondary keys array

  @dstats = Hash.new {|h,k| h[k] = Set.new}
  # device_name => [:other_stat_names...]

  @stq = Set.new # short queue names
end

def st_record(key, newvals, norate = false)
  oldvals = @st[key][:o] = @st[key][:n]
  @st[key][:to] = @st[key][:tn]
  @st[key][:n] = newvals
  @st[key][:tn] = @now
  if oldvals
    if norate
      @st[key][:r] = newvals
    else
      @st[key][:t] = @st[key][:tn] - @st[key][:to]	# passed time
      @st[key][:d] = newvals.sub(oldvals)		# delta value
      @st[key][:r] = @st[key][:d].div(@st[key][:t])	# rate
    end

    dev, q = key
    if q.to_s =~ /^[tr]x_queue_/
      # global list of valid queues
      @qstats[dev].add(q)
    else
      # valid scalars
      @dstats[dev].add(q)
    end
  end
end

def net_devices_pci
  Dir['/sys/class/net/*'].reject do |f|
    f += "/device" unless File.symlink?(f)
    if File.symlink?(f)
      !(File.readlink(f) =~ %r{devices/pci})
    else
      false
    end
  end.map {|f| File.basename(f)}
end

def ethtool_grab_stat(dev = nil)
  unless dev
    net_devices_pci.each {|e| ethtool_grab_stat(e)}
    return
  end
  h = Hash.new {|k,v| k[v] = Array.new}
  t = `ethtool -S #{dev} 2>/dev/null`
  return if t == ''
  t.split("\n").map { |e|
    e.split(':')
  }.reject { |e|
    !e[1] # no value
  }.each { |a,b|
    a.strip!
    b = b.strip.to_i
    if a =~ /^.x_queue_(\d+)_/
      qt = a.split('_')
      qnum = qt[2].to_i
      qt.delete_at(2)
      qnam = [qt[0], qnum].join('-') # short queue name (rx-1)
      @stq.add(qnam)
      qt = qt.join('_').to_sym # queue name w/o number (rx_queue_packets)
      h[qt][qnum] = b
    else
      h[a.to_sym] << b
    end
  }
  h.each_pair {|k,v| st_record([dev, k], v)}
end

def ethtool_grab_regs(dev = nil)
  unless dev
    net_devices_pci.each {|e| ethtool_grab_regs(e)}
    return
  end
  t = `ethtool -d #{dev} 2>/dev/null`
  return if t == ''
  # 0x02808: RDLEN  (Receive desc length)                 0x00010000
  # 0x03810: TDH         (Transmit desc head)             0x000008D0
  # 0x02810: RDH0        (Rx descriptor head0)            0x00000144
  h = Hash.new {|k,v| k[v] = Hash.new}
  t.split("\n").map do |e|
    r, v = e.scan(%r{^0x[0-9A-Z]+: ([A-Z0-9]+)\s.*\s(0x[0-9A-Z]+)$}).flatten
    next unless r =~ %r{^[TR]D(H|T|LEN)\d+$}
    dir = r[/^./].to_sym # direction
    qn  = r[/\d+/].to_i	 # queue number
    rn  = r[/^.D([HTL])/, 1].to_sym # register short name (H, T, L)
    v   = v.hex

    v /= 16 if rn == :L # descriptor len is always 16 (bytes)
    x = h[[dir, qn]] ||= {}
    x[rn] = v
  end
  # calculate queue lengths
  y = Hash.new {|k,v| k[v] = Array.new}
  h.each do |k, v|
    dir, qn = k
    next if v[:L] == 0
    k = (dir == :R)? "rx" : "tx"
    k = (k + "_queue_qlen").to_sym
    b = v[:T] - v[:H]
    b += v[:L] if b < 0
    # make receive queue negative
    b = b + 1 - v[:L] if dir == :R
    y[k][qn.to_i] = b
  end
  y.each_pair {|k,v| st_record([dev, k], v, :norate)}
end

# sorted list of ethernets
def st_eths
  @st.keys.map{|a,b| a}.uniq.sort_natural
end

# red everything except plain stat
def colorize_stat(st, txt, val = nil)
  if st == "bytes" || st == "packets" || val == 0 || !@color
    txt
  elsif st =~ /_err|[rt]x_no_|drop|restart|out_of_/
    "\e[1;31m#{txt}\e[m"
  else
    "\e[1;30m#{txt}\e[m"
  end
end

def print_value_unformatted(qst, v)
  unless v
    # empty cell
    print " %s" % '-'
    return
  end
  print colorize_stat(qst,
    if v == 0
      " %d" % v
    elsif v < 0.1
      " %.2f" % v
    elsif v < 10
      " %.1f" % v
    elsif v < 1e6
      " %d" % v
    elsif v < 1e8
      " %dK" % (v / 1e3)
    elsif v < 1e11
      " %dM" % (v / 1e6)
    elsif
      " %dG" % (v / 1e9)
    end, v)
end
def print_value(qst, v)
  unless v
    # empty cell
    print " %6s" % v
    return
  end
  print colorize_stat(qst,
    if v == 0
      " %6d" % v
    elsif v < 0.1
      " %6.2f" % v
    elsif v < 10
      " %6.1f" % v
    elsif v < 1e6
      " %6d" % v
    elsif v < 1e8
      " %5dK" % (v / 1e3)
    elsif v < 1e11
      " %5dM" % (v / 1e6)
    elsif
      " %5dG" % (v / 1e9)
    end, v)
end

def fix_packets_sum
  st_eths.each do |dev|
    [[:rx_packets, /^rx_.cast_packets$/],
     [:tx_packets, /^tx_.cast_packets$/]].each do |k, re|
      unless @dstats[dev].include?(k)
	# sum rx_ucast_packets + rx_mcast_packets + rx_bcast_packets into rx_packets
	a = [0]
	@dstats[dev].map{|e| e.to_s}.grep(re).each do |kst|
	  a.add!(@st[[dev, kst.to_sym]][:r])
	end
	@st[[dev, k]][:r] = a
      end
    end
  end
end

def eth_stat_print
  # output: header, ethN queues, ethN errors, ethN+1...
  eths = st_eths.sort_natural

  # Calculate and print header
  # Different eths could have different amount of queues.
  # Totals will be output on the first column before their (tx/rx) queues.
  # Also, if there is queue stat that have total (like drops),
  #   it should be output only once in queues table total column
  hn = Set.new(@stq)
  hn.add('rx')
  hn.add('tx')
  cols = {}
  hn.to_a.sort_natural.each_with_index {|e, i| cols[i] = e}
  # cols: column index to short queue name (like 'tx-1')

  print "%14s:" % @now.strftime('-%H:%M.%S- ')
  hn.size.times {|i| print " %6s" % cols[i]}
  puts " (#{@hostname})"
  hn = nil

  havelines = false
  eths.each do |dev|
    # first output queues stats, then errors

    # in what order we want to output queue stats:
    #  packets, drops, no bytes, others...
    ro = @qstats[dev]
    ro = ro.map{|e| e.to_s.sub(/^[tr]x_queue_/, '')}.uniq
    ro.bubble_up("drops")
    ro.bubble_up("packets")
    ro.bubble_up("bytes")
    ro.bubble_up("qlen")
    ro.insert(0, "drops")   if !ro.include?("drops")
    ro.insert(0, "packets") if !ro.include?("packets")
    ro.insert(0, "bytes")   if !ro.include?("bytes")
    ro.delete("bytes") unless @bytes
    # ro: array of queue stats names to output by row

    # output queue stats
    ro.each do |qst|

      # prepare row with numbers only
      # we can then skip this row if all values are 0
      r = []

      cols.each do |cnum, cname|
	if cname == "rx" || cname == "tx"
	  xqs = (qst == 'drops')? 'dropped' : qst
	  qq = "#{cname}_#{xqs}".to_sym
	  r[cnum] = @st[[dev, qq]][:r][0] rescue nil
	else
	  q, n = cname.split('-')
	  n = n.to_i
	  qq = "#{q}_queue_#{qst}".to_sym
	  if @st[[dev, qq]]
	    r[cnum] = @st[[dev, qq]][:r][n] rescue nil
	  end
	end
      end
      # skip row if it's empty
      next if r.map{|e| e.to_i}.reduce(:+) == 0

      # all good, print line
      print "%5s" % dev
      print colorize_stat(qst, " %8s:" % qst)
      r.each do |v|
	print_value(qst, v)
      end
      puts
      havelines = true
    end

    # output scalar (errors)
    r = [] # [[name, value]...]
    @dstats[dev].each do |sn|
      # skip redundant stat
      next if sn.to_s =~ /^(.x_(packets|bytes)|.x_long_byte_count|.x_(.cast|\d.*byte)_packets)$/
      next if sn.to_s =~ /^multicast$/ && @dstats[dev].include?(:rx_multicast)
      # skip non-errors
      next unless sn.to_s =~ /_err|[rt]x_no_|drop|restart|out_of_/ || @showall
      v = @st[[dev, sn]][:r].first
      r << [sn, v] if v > 0
    end
    next if r.empty?

    # all good, print line
    print "%5s" % dev
    r.each do |kst, v|
      print colorize_stat(kst, " %s:" % kst)
      print_value_unformatted(kst, v)
    end
    puts
    havelines = true

  end

  puts " zero stat from: " + eths.join(', ') unless havelines
end

@hostname = `hostname -s`.strip

loop do
  @now = Time.now
  stat_reinit
  ethtool_grab_stat
  fix_packets_sum
  ethtool_grab_regs if @qlen

  unless @notfirst
    @notfirst = true
    sleep [0.5, @interval].min
    next
  end

  puts
  eth_stat_print
  sleep @interval
end

