#!/usr/bin/env ruby
#
# kumofs
#
# Copyright (C) 2009 FURUHASHI Sadayuki
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
#

begin
	require 'rubygems'
rescue LoadError
end
require 'msgpack'
require 'socket'


class KumoRPC
	def initialize(host, port)
		@sock = TCPSocket.open(host, port)
		@pk = MessagePack::Unpacker.new
		@buffer = ''
		@nread = 0
		@seqid = rand(1<<16)  # FIXME 1 << 32
		@callback = {}
	end

	def close
		@sock.close
	end


	private
	module MessageType
		Request     = 0
		Response    = 1
		ClusterInit = 2
	end

	def send_request(seq, cmd, param)
		@sock.write [MessageType::Request, seq, cmd, param].to_msgpack
		@sock.flush
	rescue
		@sock.close
		raise
	end

	def receive_message
		while true
			if @buffer.length > @nread
				@nread = @pk.execute(@buffer, @nread)
				if @pk.finished?
					msg = @pk.data
					@pk.reset
					@buffer.slice!(0, @nread)
					@nread = 0
					if msg[0] == MessageType::Request
						process_request(msg[1], msg[2], msg[3])
					elsif msg[0] == MessageType::Response
						process_response(msg[1], msg[3], msg[2])
					else  # MessageType::ClusterInit
						raise "unexpected cluster init message"
					end
					return msg[1]
				end
			end
			@buffer << @sock.sysread(1024)
		end
	end

	def process_request(seqid, cmd, param)
		#raise "request received, excpect response"
		# ignore it
	end

	def process_response(seqid, res, err)
		if cb = @callback.delete(seqid)
			cb.call(res, err)
			GC.start
		end
	end

	def synchronize_response(seqid)
		while receive_message != seqid; end
	end

	def send_request_async(cmd, param, &callback)
		seqid = @seqid
		# FIXME 1 << 32
		@seqid += 1; if @seqid >= 1<<16 then @seqid = 0 end
		@callback[seqid] = callback if callback
		send_request(seqid, cmd, param)
		seqid
	end

	def send_request_sync(cmd, param)
		res = nil
		err = nil
		seqid = send_request_async(cmd, param) {|rres, rerr|
			res = rres
			err = rerr
		}
		synchronize_response(seqid)
		return [res, err]
	end

	def send_request_sync_ex(cmd, param)
		res, err = send_request_sync(cmd, param)
		raise "error #{err}" if err
		res
	end

	class HSSeed
		def initialize(seed)
			@clocktime = seed[1]
			@date = Time.at(@clocktime >> 32)
			@clock = @clocktime & ((1<<32)-1)

			@nodes = seed[0].map {|raw|
				active = (raw.slice!(0) == "\1"[0])
				HSSeed.rpc_addr(raw) + [active]
			}
		end
		attr_reader :clocktime, :date, :clock, :nodes

		def inspect
			%[hash space timestamp:\n] +
				%[  #{@date} clock #{@clock}\n] +
				%[node:\n] +
				@nodes.map {|addr, port, active|
					"  #{addr}:#{port}  (#{active ? "active":"fault"})"
				}.join("\n")
		end

		def self.parse(raw)
			self.new(raw)
		end

		def self.rpc_addr(raw)
			if raw.length == 6
				addr = Socket.pack_sockaddr_in(0, '0.0.0.0')
				addr[2,6] = raw[0,6]
			else
				addr = Socket.pack_sockaddr_in(0, '::')
				addr[2,2]  = raw[0,2]
				addr[8,20] = raw[2,20]
			end
			Socket.unpack_sockaddr_in(addr).reverse
		end
	end

	public
	def GetStatus
		res = send_request_sync_ex(Protocol::GetNodesInfo, [])
		form = {}

		seed = HSSeed.new(res[0])
		newcomers = res[1].map {|raw| HSSeed.rpc_addr(raw) }

		return [seed.nodes, newcomers, seed.date, seed.clock]
	end

	def AttachNewServers(replace)
		send_request_sync_ex(Protocol::AttachNewServers, [replace])
	end

	def DetachFaultServers(replace)
		send_request_sync_ex(Protocol::DetachFaultServers, [replace])
	end

	def CreateBackup(suffix)
		send_request_sync_ex(Protocol::CreateBackup, [suffix])
	end

	def SetAutoReplace(enable)
		send_request_sync_ex(Protocol::SetAutoReplace, [enable])
	end

	def StartReplace()
		send_request_sync_ex(Protocol::StartReplace, [])
	end

	module Protocol
		GetNodesInfo        = 0 << 16 |  99
		AttachNewServers    = 0 << 16 | 100
		DetachFaultServers  = 0 << 16 | 101
		CreateBackup        = 0 << 16 | 102
		SetAutoReplace      = 0 << 16 | 103
		StartReplace        = 0 << 16 | 104
		GetStatus           = 0 << 16 |  97
		SetConfig           = 0 << 16 |  98
	end

	MANAGER_DEFAULT_PORT = 19700
	SERVER_DEFAULT_PORT  = 19800

	class HashSpace
		VIRTUAL_NODE_NUMBER = 128

		class Node
			def initialize(addr, port, is_active)
				@addr = addr
				@port = port.to_i
				@is_active = is_active
			end
			attr_reader :addr, :port, :is_active
			def active?
				@active
			end
			def dump_addr
				Socket.pack_sockaddr_in(@port, @addr)[2, 6]
			end
		end

		class VirtualNode
			def initialize(hash, real)
				@hash = hash
				@real = real
			end
			attr_reader :hash, :real
			def <=>(o)
				@hash <=> o.hash
			end
		end

		def self.hash(str)
			require 'digest/sha1'
			Digest::SHA1.digest(str)[0,8].unpack('Q')[0]
		end

		def initialize
			@nodes = []
			@space = []
		end
		attr_reader :nodes, :space

		def add_server(addr, port, is_active = true)
			real = Node.new(addr, port, is_active)
			@nodes << real
			add_virtual_nodes(real)
			@space.sort!
			real
		end

		def find_each(h, &block)
			h = self.class.hash(h) if h.is_a?(String)
			first = 0
			@space.each {|v|
				break unless v.hash < h
				first += 1
			}
			i = first
			last = first + @space.length
			while true
				break unless yield @space[i % @space.length].real
				i += 1
				break if i == last
			end
		end

		private
		def add_virtual_nodes(real)
			x = self.class.hash(real.dump_addr)
			@space << VirtualNode.new(x, real)
			(VIRTUAL_NODE_NUMBER-1).times {
				x = self.class.hash([x].pack('Q'))
				@space << VirtualNode.new(x, real)
			}
			nil
		end

		def rehash
			@space.clear
			@nodes.each {|real|
				add_virtual_nodes(real)
			}
			@space.sort!
		end
	end
end


class KumoManager < KumoRPC
	def initialize(host, port)
		super(host, port)
	end

	def AttachNewServers(replace)
		send_request_sync_ex(Protocol::AttachNewServers, [replace])
	end

	def DetachFaultServers(replace)
		send_request_sync_ex(Protocol::DetachFaultServers, [replace])
	end

	def CreateBackup(suffix)
		send_request_sync_ex(Protocol::CreateBackup, [suffix])
	end

	def SetAutoReplace(enable)
		send_request_sync_ex(Protocol::SetAutoReplace, [enable])
	end

	def StartReplace(full = false)
		if full
			send_request_sync_ex(Protocol::StartReplace, [true])
		else
			send_request_sync_ex(Protocol::StartReplace, [])
		end
	end
end

if $0 == __FILE__



$now = Time.now.strftime("%Y%m%d")

def usage
	puts "Usage: #{File.basename($0)} address[:port=#{KumoRPC::MANAGER_DEFAULT_PORT}] command [options]"
	puts "command:"
	puts "   status                     get status"
	puts "   attach                     attach all new servers and start replace"
	puts "   attach-noreplace           attach all new servers"
	puts "   detach                     detach all fault servers and start replace"
	puts "   detach-noreplace           detach all fault servers"
	puts "   replace                    start replace without attach/detach"
	puts "   full-replace               start full-replace (repair consistency)"
	puts "   backup  [suffix=#{$now }]  create backup with specified suffix"
	puts "   enable-auto-replace        enable auto replace"
	puts "   disable-auto-replace       disable auto replace"
	exit 1
end

if ARGV.length < 2
	usage
end

addr = ARGV.shift
host, port = addr.split(':', 2)
port ||= KumoRPC::MANAGER_DEFAULT_PORT

cmd = ARGV.shift

case cmd
when "stat", "status"
	usage if ARGV.length != 0
	attached, not_attached, date, clock =
			KumoManager.new(host, port).GetStatus
	puts "hash space timestamp:"
	puts "  #{date} clock #{clock}"
	puts "attached node:"
	attached.each {|addr, port, active|
		puts "  #{addr}:#{port}  (#{active ? "active":"fault"})"
	}
	puts "not attached node:"
	not_attached.each {|addr, port|
		puts "  #{addr}:#{port}"
	}

when "attach"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).AttachNewServers(true)

when "attach-noreplace"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).AttachNewServers(false)

when "detach"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).DetachFaultServers(true)

when "detach-noreplace"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).DetachFaultServers(false)

when "enable-auto-replace"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).SetAutoReplace(true)

when "disable-auto-replace"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).SetAutoReplace(false)

when "backup"
	if ARGV.length == 0
		suffix = $now
	elsif ARGV.length == 1
		suffix = ARGV.shift
	else
		usage
	end
	puts "suffix=#{suffix}"
	p KumoManager.new(host, port).CreateBackup(suffix)

when "replace"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).StartReplace()

when "full-replace"
	usage if ARGV.length != 0
	p KumoManager.new(host, port).StartReplace(true)

else
	puts "unknown command #{cmd}"
	puts ""
	usage
end


end   # if $0 == __FILE__
