Test multicast reception via Ruby

mcast-data-test.rb

#!/usr/bin/env ruby
# - @oogali

require 'rubygems'
require 'socket'
require 'ipaddr'
require 'getoptlong'

MAX_RUNTIME_DURATION = 30 # seconds

def usage
  puts
  puts "#{__FILE__} [--bind ADDRESS] [--threshold THRESHOLD] --group GROUP:PORT"
  puts
  puts "-b, --bind\t\tAddress to bind our multicast socket to"
  puts "-t, --threshold\tMinimum # of packets we expect to receive in #{MAX_RUNTIME_DURATION} seconds"
  puts "-g, --group\t\tMulticast group we should join (IP:PORT)"
  puts

  exit 1
end

options = {
  :addr => nil,
  :port => nil,
  :bindaddr => nil,
  :threshold => 10
}

opts = GetoptLong.new(
  [ '--group', '-g', GetoptLong::REQUIRED_ARGUMENT ],
  [ '--bind', '-b', GetoptLong::OPTIONAL_ARGUMENT ],
  [ '--threshold', '-t', GetoptLong::OPTIONAL_ARGUMENT ]
)

opts.each do |opt, arg|
  case opt
    when '--group'
      options[:addr], options[:port] = arg.split(':', 2)
    when '--bind'
      options[:bindaddr] = arg
    when '--threshold'
      options[:threshold] = arg.to_i
  end
end

if options[:threshold] <= 0
  puts 'You must specify a non-zero packet count threshold.'
  exit 1
end

usage unless options[:addr] and options[:port]

count = 0
ip = IPAddr.new(options[:addr]).hton + IPAddr.new(options[:bindaddr] || '0.0.0.0').hton

begin
  now = Time.now
  puts "Started: #{now.strftime '%Y-%m-%d %H:%M:%S'}"
  puts

  sock = UDPSocket.open
  sock.setsockopt Socket::SOL_SOCKET, Socket::SO_REUSEADDR, true
  sock.setsockopt Socket::IPPROTO_IP, Socket::IP_ADD_MEMBERSHIP, ip
  sock.bind options[:bindaddr] || Socket::INADDR_ANY, options[:port]

  loop do
    begin
      if Time.now - now > MAX_RUNTIME_DURATION
        puts "ERROR: Execution time limit reached before we hit our threshold"
        exit 1
      end

      msg, info = sock.recvfrom_nonblock(1024)
    rescue Errno::EAGAIN, Errno::ECONNABORTED, Errno::EINTR, Errno::EWOULDBLOCK
      sleep 0.005
      retry
    end

    puts "msg from #{info[2]} (#{'%15s' % info[3]} / #{'%5d' % info[1]}) len #{msg.size}"
    count += 1

    if count >= options[:threshold]
      puts
      puts "Reached our threshold successfully (in #{Time.now - now} seconds)"
      exit 0
    end
  end
ensure
  sock.close
end