dr-rubbis 0.0.0

data/lib/rubbis/state.rb

@@ -0,0 +1,407 @@
+require 'set'
+require 'rubbis/zset'
+
+module Rubbis
+	Error = Struct.new(:message) do
+		def self.incorrect_args(cmd)
+			new "wrong number of arguments for '#{cmd}' command"
+		end
+
+		def self.unknown_cmd(cmd)
+			new "unknown command '#{cmd}'"
+		end
+
+		def self.type_error
+			new "wrong type for command"
+		end 
+	end
+
+	class State
+		def initialize(clock)
+			@clock = clock
+			@data = {}
+			@log = []
+			@expires = {}
+			@watches = {}
+			@list_watches = {}
+			@ready_keys = []
+			@channels = Hash.new { |hash, key| hash[key] = Set.new }
+			@subscribers = Hash.new { |hash, key| hash[key] = Set.new }
+			@pchannels = Hash.new { |hash, key| hash[key] = Set.new }
+			@psubscribers = Hash.new { |hash, key| hash[key] = Set.new }
+		end
+
+		def self.valid_command?(cmd)
+			@valid_commands ||= Set.new(
+				StateCommands.public_instance_methods.map(&:to_s)
+				)
+			@valid_commands.include?(cmd) || blocking_command?(cmd)
+		end
+
+		def self.blocking_command?(cmd)
+			@blocking_commands ||= Set.new(
+				BlockingCommands.public_instance_methods.map(&:to_s)
+				)
+			@blocking_commands.include?(cmd)
+		end
+
+		def serialize
+			Marshal.dump [@data, @expires]
+		end
+
+		def deserialize(bytes)
+			@data, @expires = *Marshal.load(bytes)
+		end
+
+		def channel_count(client)
+			channels[client].length + pchannels[client].length
+		end
+
+		def apply_command(client, cmd)
+			unless State.valid_command?(cmd[0])
+				return Error.unknown_cmd(cmd[0])
+			end
+
+			if State.blocking_command?(cmd[0])
+				cmd << client
+			end
+			
+			public_send *cmd
+		end
+
+		def subscribers_for(channel)
+			(
+				subscribers[channel].to_a + 
+				psubscribers.select {|pattern, _|
+					File.fnmatch(pattern, channel)
+				}.values.map(&:to_a).flatten
+			)
+		end
+
+		module StateCommands
+			def publish(channel, message)
+				subscribers_for(channel).each do |client|
+					client.respond! ["message", channel, message]
+				end.length
+			end
+
+			def ping
+				:pong
+			end
+
+			def echo(text)
+				text
+			end
+
+			def watch(key, &block)
+				watches[key] ||= []
+				watches[key] << block if block
+				:ok
+			end
+
+			def unsubscribe_all(client)
+				(channels.delete(client) || Set.new).each do |channel|
+					subscribers.delete channel
+				end
+			end
+
+			def expire(key, secs)
+				pexpire(key, secs.to_i * 1000)
+			end
+
+			def pexpire(key, ms)
+				pexpireat(key, clock.now * 1000.0 + ms.to_i)
+			end
+
+			def minimal_log
+				data.map do |key, value|
+					case value
+					when String
+						[['set', key, value]]
+					when Array
+						value.reverse.map do |v|
+							['lpush', key, v]
+						end
+					else raise 'unimplemented'
+					end
+				end.reduce(:+) + expires.map do |key, value|
+					['pexpireat', key, (value * 1000).to_i.to_s]
+				end
+			end
+
+			def pexpireat(key, value)
+				if get(key)
+					expires[key] =  value.to_i / 1000.0
+					log << ['pexpireat', key, value.to_i.round.to_s]
+					1
+				else
+					0
+				end
+			end
+
+			def expire_keys!(n: 100, threshold: 0.25, rng: Random.new)
+				begin
+					expired = expires.keys.sample(n, random: rng).count do |key|
+						get(key)
+					end
+				end while expired > n * threshold
+			end
+
+			def process_list_watches!
+				ready_keys.each do |key|
+					list = get(key)
+					watches = list_watches.fetch(key, [])
+
+					while list.any? && watches.any?
+						op, client = *watches.shift
+						client.respond!(op.call) if client.active?
+					end
+				end
+
+				ready_keys.clear
+			end
+
+			def set(*args)
+				key, value ,modifier = *args
+
+				return Error.incorrect_args("set") unless key && value
+
+				exists = data.has_key?(key)
+				nx = modifier == "NX"
+				xx = modifier == "XX"
+
+				if (!nx && !xx) || (nx && !exists) || (xx && exists)
+					touch! key
+					log << ['set', key, value]
+					data[key] = value
+					:ok
+				end
+			end
+
+			def get(key)
+				expiry = expires[key]
+				del(key) if expiry && expiry <= clock.now
+
+				data[key]
+			end
+
+			def del(key)
+				touch! key
+				if data.delete(key)
+					log << ['del', key]
+				end
+				expires.delete(key)
+			end
+
+			def hset(*args)
+				hash, key, value = *args
+
+				return Error.incorrect_args("hset") unless hash && key && value
+				touch! key
+				data[hash] ||= {}
+				data[hash][key] = value
+				:ok
+			end
+
+			def hget(*args)
+				hash, key = *args
+
+				return Error.incorrect_args("hget") unless hash && key
+
+				value = get(hash)
+
+				value[key] if value
+			end
+
+			def hmget(hash, *keys)
+				existing = get(hash) || {}
+				if existing.is_a?(Hash)
+					existing.values_at(*keys)
+				else
+					Error.type_error
+				end
+			end
+
+			def hincrby(hash, key, amount)
+				value = get(hash)
+				if value
+					existing = value[key]
+					value[key] = existing.to_i + amount.to_i
+				end
+			end
+
+			def exists(key)
+				if data[key]
+					1
+				else
+					0
+				end
+			end
+
+			def keys(pattern)
+				if pattern == '*'
+					result = data.keys
+
+					result
+				else
+					raise 'unimplemented'
+				end
+			end
+
+			def zadd(key, score, member)
+				score = score.to_f
+
+				value = get(key) || data[key] = ZSet.new
+
+				value.add(score, member)
+			end
+
+			def zrange(key, start, stop)
+				value = get(key)
+				if value
+					value.range(start.to_i, stop.to_i)
+				else
+					[]
+				end
+			end
+
+			def zrank(key, member)
+				value = get(key)
+				if value
+					value.rank(member)
+				else
+					-1
+				end
+			end
+
+			def zscore(key, member)
+				value = get(key)
+				if value
+					value.score(member)
+				else
+					-1
+				end
+			end
+
+
+			def lpush(key, value)
+				list = get(key)
+				list ||= data[key] = []
+
+				if list_watches.fetch(key, []).any?
+					ready_keys << key
+				end
+				touch! key
+				list.unshift value
+				list.length
+			end
+
+			def llen(key)
+				list = get(key)
+				list ||= data[key] = []
+
+				list.length
+			end
+
+			def lrange(key, start, stop)
+				list = get(key)
+
+				if list
+					list[start.to_i..stop.to_i]
+				else
+					[]
+				end
+			end
+
+			def rpop(key)
+				list = get(key)
+				list ||= data[key] = []
+
+				touch! key
+				list.pop
+			end
+
+
+			def rpoplpush(pop_key, push_key)
+				item = rpop(pop_key)
+				return unless item
+				lpush push_key, item
+				item
+			end
+		end
+
+
+		module BlockingCommands
+			def subscribe(channel, client)
+				subscribers[channel] << client
+				channels[client] << channel
+
+				['subscribe', channel, channel_count(client)]
+			end
+
+			def unsubscribe(channel, client)
+				subscribers[channel].delete client
+				channels[client].delete channel
+
+				['unsubscribe', channel, channel_count(client)]
+			end
+
+			def psubscribe(channel, client)
+				psubscribers[channel] << client
+				pchannels[client] << channel
+
+				['psubscribe', channel, channel_count(client)]
+			end
+
+			def punsubscribe(channel, client)
+				psubscribers[channel].delete client
+				pchannels[client].delete channel
+
+				['punsubscribe', channel, channel_count(client)]
+			end
+
+			def brpoplpush(pop_key, push_key, timeout, client)
+
+				action = ->{ rpoplpush(pop_key, push_key) }
+
+				if llen(pop_key) == 0
+					list_watches[pop_key] ||= []
+					list_watches[pop_key] << [action, client]
+					:block
+				else
+					action.call
+				end
+			end			
+
+			def brpop(key, timeout, client)
+				list = get(key)
+				list ||= data[key] = []
+
+				action = ->{ rpop(key) }
+
+				if llen(key) == 0
+					list_watches[key] ||= []
+					list_watches[key] << [action, client]
+					:block
+				else
+					action.call
+				end
+			end
+		end
+
+		include StateCommands
+		include BlockingCommands
+		
+		attr_reader :log
+
+		private
+
+		attr_reader :data, :expires, :clock, :watches, :list_watches, :ready_keys,
+					:channels, :subscribers, :pchannels, :psubscribers
+
+		def touch!(key)
+			ws = watches.delete(key) || []
+			ws.each(&:call)
+		end
+	end
+end

data/lib/rubbis/transaction.rb

@@ -0,0 +1,34 @@
+module Rubbis
+	class Transaction
+		def initialize
+			@active = false
+			@buffer = []
+			@dirty = false
+		end
+
+		def start!
+			@active = true
+		end
+
+		def active?
+			@active
+		end
+
+		def dirty!
+			@dirty = true
+		end
+
+		def dirty?
+			@dirty
+		end
+
+		def queue(cmd)
+			raise unless @active
+			@buffer << cmd
+		end
+
+		def buffer
+			@buffer
+		end
+	end
+end

data/lib/rubbis/zset.rb

@@ -0,0 +1,27 @@
+class ZSet
+	# i've skipped refactoring this code to work in O(logn)
+	def initialize
+		@entries = []
+	end
+
+	def add(score, member)
+		entries << [score, member]
+		entries.sort!
+	end
+
+	def range(start, stop)
+		entries[start..stop].map {|x| x[1]}
+	end
+
+	def rank(member)
+		entries.index {|x| x[1] == member}
+	end
+
+	def score(member)
+		entries.detect {|x| x[1] == member}[0]
+	end
+
+	private
+
+	attr_reader :entries
+end

data/spec/acceptance/expiry_spec.rb

@@ -0,0 +1,27 @@
+require 'spec_helper'
+
+describe 'Rubbis', :acceptance do
+	it 'actively expires keys' do
+		with_server do
+			n = 10
+			n.times do |x|
+				kkey = "keep#{x}"
+				ekey = "expire#{x}"
+				client.set(kkey, "123")
+				client.set(ekey, "123")
+				client.pexpire(ekey, rand(600))
+			end
+			condition = ->{
+				client.keys("*").count {|x| x.start_with?("expire") } == 0
+			}
+
+			start_time = Time.now
+			while !condition.call && (Time.now < start_time + 2)
+				sleep 0.01
+			end
+
+			expect(condition.call).to eq(true)
+			expect(client.keys('*').size).to eq(n)
+		end
+	end
+end

data/spec/acceptance/hash_spec.rb

@@ -0,0 +1,11 @@
+require 'spec_helper'
+
+describe 'Rubbis', :acceptance do
+	it 'supports hashes' do
+		with_server do
+			client.hset('myhash', 'abc', '123')
+			client.hset('myhash', 'def', '456')
+			expect(client.hmget("myhash", "abc", "def")).to eq(["123", "456"])
+		end
+	end
+end

data/spec/acceptance/list_spec.rb

@@ -0,0 +1,51 @@
+require 'spec_helper'
+
+describe "Rubbis", :acceptance do
+	it 'supports blocking pop' do
+		with_server do
+			items = %w(a b)
+			t1 = Thread.new do
+				item = client.brpop("queue")
+				expect(item).to eq(items.shift)
+			end
+			t2 = Thread.new do
+				item = client.brpop("queue")
+				expect(item).to eq(items.shift)
+			end
+
+			items.dup.each do |item|
+				expect(client.lpush("queue", item)).to eq(1)
+			end
+			t1.value
+			t2.value
+		end
+	end
+
+	it 'supports brpoplpush' do
+		with_server do
+			c = client
+
+			t1 = Thread.new do
+				item = c.brpoplpush("q","p")
+				expect(item).to eq('a')
+				expect(c.lrange('p', "0", "-1")).to eq(%w(a))
+			end
+
+			c.lpush("q", "a")
+			
+			t1.value
+		end
+	end
+
+	it 'handles disconnecting clients' do
+		with_server do
+
+			c = TCPSocket.new('localhost', TEST_PORT)
+			c.write ("*3\r\n$5\r\nbrpop\r\n$1\r\nq\r\n$1\r\n0\r\n")
+			c.close
+
+			expect(client.lpush("q", "a")).to eq(1)
+			expect(client.lpush("q", "b")).to eq(2)
+		end
+	end
+end

data/spec/acceptance/persistence_spec.rb

@@ -0,0 +1,25 @@
+require 'spec_helper'
+
+describe 'Rubbis', :acceptance do
+	it 'can persist data' do 
+		with_server do 
+			client.set("a", "1")
+			client.bgsave
+		end
+
+		with_server do
+			expect(client.get("a")).to eq("1")
+		end
+	end
+
+	it 'can persist data with aof' do 
+		with_server(server_file: false, aof_file: true) do 
+			client.set("a", "1")
+			client.bgsave
+		end
+
+		with_server(server_file: false, aof_file: true) do
+			expect(client.get("a")).to eq("1")
+		end
+	end
+end

data/spec/acceptance/pubsub_spec.rb

@@ -0,0 +1,48 @@
+require 'spec_helper'
+
+describe 'Rubbis', :acceptance do
+	it 'support pubsub' do
+		with_server do 
+			received = nil
+			c = client
+
+			t1 = Thread.new do
+				c.subscribe("mychan") do |on|
+					on.message do |channel, msg|
+						received = msg
+						c.unsubscribe("mychan")
+					end
+				end
+			end
+
+			sleep 0.1
+
+			client.publish('mychan', 'hello')
+			t1.value
+			expect(received).to eq('hello')
+		end
+	end
+
+	it 'supports pubsub with patterns' do
+		with_server do
+			received = nil
+			c = client
+
+			t1 = Thread.new do
+				c.psubscribe("my.*") do |on|
+					on.message do |channel, msg|
+						received = msg
+						c.punsubscribe("my.*")
+					end
+				end
+			end
+
+			sleep 0.1
+
+			client.publish('your.chan', 'bogus')
+			client.publish('my.chan', 'hello')
+			t1.value
+			expect(received).to eq('hello')
+		end
+	end
+end

data/spec/acceptance/skeleton_spec.rb

@@ -0,0 +1,34 @@
+require "spec_helper"
+
+describe 'Rubbis', :acceptance do
+	it 'responds to ping' do 
+		with_server do
+			c = client
+			c.without_reconnect do
+				expect(c.ping).to eq("PONG")
+				expect(c.ping).to eq("PONG")
+			end
+		end
+	end
+	
+	it 'echoes msgs' do 
+		with_server do
+ 			expect(client.echo("hello\nthere")).to eq("hello\nthere")
+		end
+	end
+
+	it 'supports multiple clients simultaneously' do 
+		with_server do
+ 			expect(client.echo("hello\nthere")).to eq("hello\nthere")
+ 			expect(client.echo("hello\nthere")).to eq("hello\nthere")
+		end
+	end
+
+	it 'gets and sets values' do
+		with_server do
+			expect(client.get("abc")).to eq(nil)
+			expect(client.set("abc", "123")).to eq("OK")
+			expect(client.get("abc")).to eq("123")
+		end
+	end
+end

data/spec/acceptance/transactions_spec.rb

@@ -0,0 +1,49 @@
+require 'spec_helper'
+
+describe 'Rubbis', :acceptance do
+	it 'handles transactions' do
+		with_server do
+			c = client
+
+			result = c.multi do
+				c.set("abc", "123")
+				c.get("abc")
+			end
+
+			expect(result).to eq(%w(OK 123))
+
+			begin
+				c.multi do
+					c.set("abc", "456")
+					raise 
+				end
+			rescue
+			end
+
+			expect(c.get("abc")).to eq("123")
+		end
+	end
+
+	it 'supports WATCH' do
+		with_server do
+			c1 = client
+			c2 = client
+
+			c1.set("abc", 1)
+
+			c1.watch("abc")
+			c1.watch("def")
+
+			prev = c1.get("abc")
+
+			c2.set("abc", "10")
+
+			c1.multi do
+				c1.set("abc", prev.to_i + 1)
+			end
+
+			expect(c1.get("abc")).to eq("10")
+		end
+
+	end
+end