client-libraries/ruby/lib/redis.rb

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