client-libraries/ruby_2/rubyredis.rb

   1 # RubyRedis is an alternative implementatin of Ruby client library written
   2 # by Salvatore Sanfilippo.
   3 #
   4 # The aim of this library is to create an alternative client library that is
   5 # much simpler and does not implement every command explicitly but uses
   6 # method_missing instead.
   7
   8 require 'socket'
   9 require 'set'
  10
  11 begin
  12     if (RUBY_VERSION >= '1.9')
  13         require 'timeout'
  14         RedisTimer = Timeout
  15     else
  16         require 'system_timer'
  17         RedisTimer = SystemTimer
  18     end
  19 rescue LoadError
  20     RedisTimer = nil
  21 end
  22
  23 class RedisClient
  24     BulkCommands = {
  25         "set"=>true, "setnx"=>true, "rpush"=>true, "lpush"=>true, "lset"=>true,
  26         "lrem"=>true, "sadd"=>true, "srem"=>true, "sismember"=>true,
  27         "echo"=>true, "getset"=>true, "smove"=>true
  28     }
  29
  30     ConvertToBool = lambda{|r| r == 0 ? false : r}
  31     ConvertToSet = lambda{|r| Set.new(r)}
  32
  33     ReplyProcessor = {
  34         "exists" => ConvertToBool,
  35         "sismember"=> ConvertToBool,
  36         "sadd"=> ConvertToBool,
  37         "srem"=> ConvertToBool,
  38         "smove"=> ConvertToBool,
  39         "move"=> ConvertToBool,
  40         "setnx"=> ConvertToBool,
  41         "del"=> ConvertToBool,
  42         "renamenx"=> ConvertToBool,
  43         "expire"=> ConvertToBool,
  44         "smembers" => ConvertToSet,
  45         "sinter" => ConvertToSet,
  46         "sunion" => ConvertToSet,
  47         "sdiff" => ConvertToSet,
  48         "keys" => lambda{|r| r.split(" ")},
  49         "info" => lambda{|r|
  50             info = {}
  51             r.each_line {|kv|
  52                 k,v = kv.split(":",2).map{|x| x.chomp}
  53                 info[k.to_sym] = v
  54             }
  55             info
  56         }
  57     }
  58
  59     Aliases = {
  60         "flush_db" => "flushdb",
  61         "flush_all" => "flushall",
  62         "last_save" => "lastsave",
  63         "key?" => "exists",
  64         "delete" => "del",
  65         "randkey" => "randomkey",
  66         "list_length" => "llen",
  67         "push_tail" => "rpush",
  68         "push_head" => "lpush",
  69         "pop_tail" => "rpop",
  70         "pop_head" => "lpop",
  71         "list_set" => "lset",
  72         "list_range" => "lrange",
  73         "list_trim" => "ltrim",
  74         "list_index" => "lindex",
  75         "list_rm" => "lrem",
  76         "set_add" => "sadd",
  77         "set_delete" => "srem",
  78         "set_count" => "scard",
  79         "set_member?" => "sismember",
  80         "set_members" => "smembers",
  81         "set_intersect" => "sinter",
  82         "set_intersect_store" => "sinterstore",
  83         "set_inter_store" => "sinterstore",
  84         "set_union" => "sunion",
  85         "set_union_store" => "sunionstore",
  86         "set_diff" => "sdiff",
  87         "set_diff_store" => "sdiffstore",
  88         "set_move" => "smove",
  89         "set_unless_exists" => "setnx",
  90         "rename_unless_exists" => "renamenx"
  91     }
  92
  93     def initialize(opts={})
  94         @host = opts[:host] || '127.0.0.1'
  95         @port = opts[:port] || 6379
  96         @db = opts[:db] || 0
  97         @timeout = opts[:timeout] || 0
  98         connect_to_server
  99     end
 100
 101     def to_s
 102         "Redis Client connected to #{@host}:#{@port} against DB #{@db}"
 103     end
 104
 105     def connect_to_server
 106         @sock = connect_to(@host,@port,@timeout == 0 ? nil : @timeout)
 107         call_command(["select",@db]) if @db != 0
 108     end
 109
 110     def connect_to(host, port, timeout=nil)
 111         # We support connect() timeout only if system_timer is availabe
 112         # or if we are running against Ruby >= 1.9
 113         # Timeout reading from the socket instead will be supported anyway.
 114         if @timeout != 0 and RedisTimer
 115             begin
 116                 sock = TCPSocket.new(host, port, 0)
 117             rescue Timeout::Error
 118                 @sock = nil
 119                 raise Timeout::Error, "Timeout connecting to the server"
 120             end
 121         else
 122             sock = TCPSocket.new(host, port, 0)
 123         end
 124
 125         # If the timeout is set we set the low level socket options in order
 126         # to make sure a blocking read will return after the specified number
 127         # of seconds. This hack is from memcached ruby client.
 128         if timeout
 129             secs = Integer(timeout)
 130             usecs = Integer((timeout - secs) * 1_000_000)
 131             optval = [secs, usecs].pack("l_2")
 132             sock.setsockopt Socket::SOL_SOCKET, Socket::SO_RCVTIMEO, optval
 133             sock.setsockopt Socket::SOL_SOCKET, Socket::SO_SNDTIMEO, optval
 134         end
 135         sock
 136     end
 137
 138     def method_missing(*argv)
 139         call_command(argv)
 140     end
 141
 142     def call_command(argv)
 143         # this wrapper to raw_call_command handle reconnection on socket
 144         # error. We try to reconnect just one time, otherwise let the error
 145         # araise.
 146         connect_to_server if !@sock
 147         begin
 148             raw_call_command(argv)
 149         rescue Errno::ECONNRESET
 150             @sock.close
 151             connect_to_server
 152             raw_call_command(argv)
 153         end
 154     end
 155
 156     def raw_call_command(argv)
 157         bulk = nil
 158         argv[0] = argv[0].to_s.downcase
 159         argv[0] = Aliases[argv[0]] if Aliases[argv[0]]
 160         if BulkCommands[argv[0]]
 161             bulk = argv[-1].to_s
 162             argv[-1] = bulk.length
 163         end
 164         @sock.write(argv.join(" ")+"\r\n")
 165         @sock.write(bulk+"\r\n") if bulk
 166
 167         # Post process the reply if needed
 168         processor = ReplyProcessor[argv[0]]
 169         processor ? processor.call(read_reply) : read_reply
 170     end
 171
 172     def select(*args)
 173         raise "SELECT not allowed, use the :db option when creating the object"
 174     end
 175
 176     def [](key)
 177         get(key)
 178     end
 179
 180     def []=(key,value)
 181         set(key,value)
 182     end
 183
 184     def sort(key, opts={})
 185         cmd = []
 186         cmd << "SORT #{key}"
 187         cmd << "BY #{opts[:by]}" if opts[:by]
 188         cmd << "GET #{[opts[:get]].flatten * ' GET '}" if opts[:get]
 189         cmd << "#{opts[:order]}" if opts[:order]
 190         cmd << "LIMIT #{opts[:limit].join(' ')}" if opts[:limit]
 191         call_command(cmd)
 192     end
 193
 194     def incr(key,increment=nil)
 195         call_command(increment ? ["incrby",key,increment] :  ["incr",key])
 196     end
 197
 198     def decr(key,decrement=nil)
 199         call_command(decrement ? ["decrby",key,decrement] :  ["decr",key])
 200     end
 201
 202     def read_reply
 203         # We read the first byte using read() mainly because gets() is
 204         # immune to raw socket timeouts.
 205         begin
 206             rtype = @sock.read(1)
 207         rescue Errno::EAGAIN
 208             # We want to make sure it reconnects on the next command after the
 209             # timeout. Otherwise the server may reply in the meantime leaving
 210             # the protocol in a desync status.
 211             @sock = nil
 212             raise Errno::EAGAIN, "Timeout reading from the socket"
 213         end
 214
 215         raise Errno::ECONNRESET,"Connection lost" if !rtype
 216         line = @sock.gets
 217         case rtype
 218         when "-"
 219             raise "-"+line.strip
 220         when "+"
 221             line.strip
 222         when ":"
 223             line.to_i
 224         when "$"
 225             bulklen = line.to_i
 226             return nil if bulklen == -1
 227             data = @sock.read(bulklen)
 228             @sock.read(2) # CRLF
 229             data
 230         when "*"
 231             objects = line.to_i
 232             return nil if bulklen == -1
 233             res = []
 234             objects.times {
 235                 res << read_reply
 236             }
 237             res
 238         else
 239             raise "Protocol error, got '#{rtype}' as initial reply byte"
 240         end
 241     end
 242 end