sadvertise

sadvertise.rb

#!/usr/bin/env ruby
require 'rubygems'
require 'optparse'
require 'ipaddr'
require 'socket'
require 'msgpack'
require 'yaml'

SA_ADDR = '224.0.0.123'
SA_PORT = '12121'

OPTIONS = {}
OPTIONS[:send_delay] = 10

class SASender
  def initialize(ifaddr = nil)
    @udp = UDPSocket.open()
    if ifaddr != nil
      mif = IPAddr.new(ifaddr, Socket::AF_INET).hton
      @udp.setsockopt(Socket::IPPROTO_IP, Socket::IP_MULTICAST_IF, mif)
    end
  end

  def send(data)
    @udp.send(data, 0, SA_ADDR, SA_PORT)
  end
end

class SAReceiver
  def initialize(ifaddr = nil)
    @udp = UDPSocket.open()
    @udp.bind(SA_ADDR, SA_PORT)

    mreq = IPAddr.new(SA_ADDR).hton + IPAddr.new(ifaddr).hton
    @udp.setsockopt(Socket::IPPROTO_IP, Socket::IP_ADD_MEMBERSHIP, mreq)
  end

  def recv
    return nil if IO.select([ @udp ], nil, nil, 10) == nil
    @udp.recv(4096)
  end
end

class SAPacket
  def initialize
    @packet = { :msgtype => :null }
  end

  def serialize
    @packet.to_msgpack
  end

  def msgtype=(cmd)
    @packet[:msgtype] = cmd
  end

  def add_data(hash)
    @packet.merge!(hash)
  end
end

class SAAdvertise < SAPacket
  def initialize
    super
    self.msgtype = :advertise
  end

  def add_payload(key, value)
    self.add_data(key => value)  
  end

  def add_payload_by_module(info_module)
    info_methods = info_module.methods
    Module.methods.each {|m| info_methods.delete(m) }

    info_methods.each do |m|
      add_payload(m, info_module.send(m.to_sym))
    end
  end
end

module HostInfo
  module_function

  def alert
    false
  end

  def hostname
    `hostname`.chomp
  end

  def ipaddr_private
    result = []
    ipaddrs = `ip addr list | grep inet`
    ipaddrs.each do |line|
      if line =~ /(10|172|192)\.(\d+\.\d+\.\d+)/
        result << "#{$1}.#{$2}"
      end
    end
    result
  end

  def ipaddr_global
    result = []
    ipaddrs = `ip addr list | grep inet`
    ipaddrs.each do |line|
      if line =~ /(\d+)\.(\d+\.\d+\.\d+)/
        a = $1.to_i
        next if a == 127 || a == 10 || a == 172 || a == 192
        result << "#{$1}.#{$2}"
      end
    end
    result
  end

  def server_role
    role = ''
    begin
      open('/etc/gree/server_role') do |io|
        role = io.readline
      end
    rescue
    end
    role.chomp
  end
end

def sender_thread(private_addr)
  Thread.new do
    delay = rand OPTIONS[:send_delay]
    puts "delay #{delay} sec" if OPTIONS[:verbose]
    sleep delay

    puts "send advertise message" if OPTIONS[:verbose]
    adv = SAAdvertise.new
    adv.add_payload_by_module(HostInfo)

    sender = SASender.new(private_addr)
    sender.send(adv.serialize)
  end
end

def receiver_proc(private_addr, filename)
  tmpname = "#{filename}.tmp"
  tstart = Time.now

  open(tmpname, 'w') do |io|
    receiver = SAReceiver.new(private_addr)
    result = []

    loop do
      if (r = receiver.recv) != nil
        result << MessagePack.unpack(r)
      else
        elapsed = Time.now - tstart
        break if elapsed > OPTIONS[:send_delay].to_i
      end
    end
    io.write(result.to_yaml)
    File.rename(tmpname, filename)
  end
end

opt = OptionParser.new
opt.on('-r FILENAME',    'enable receiver')             {|v| OPTIONS[:receive] = v }
opt.on('-v',             'verbose')                     {|v| OPTIONS[:verbose] = v }
opt.on('-w SEND_DELAY',  'send delay time in seconds')  {|v| OPTIONS[:send_delay] = v }
opt.parse!(ARGV)

private_addr = HostInfo.ipaddr_private[0]
sender = sender_thread(private_addr)

if OPTIONS[:receive]
  receiver_proc(private_addr, OPTIONS[:receive])
end

sender.join