Add support for specifying TCP keepalive configuration

redis · Aug 13, 2012 · 5885967 · 5885967
1 parent a109bcf
commit 5885967
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 0 deletions.
diff --git a/lib/redis/client.rb b/lib/redis/client.rb
@@ -1,4 +1,5 @@
 require "redis/errors"
+require "socket"
 
 class Redis
   class Client
@@ -14,6 +15,7 @@ class Client
       :db => 0,
       :driver => nil,
       :id => nil,
+      :tcp_keepalive => 0
     }
 
     def scheme
@@ -350,6 +352,26 @@ def _parse_options(options)
       options[:db] = options[:db].to_i
       options[:driver] = _parse_driver(options[:driver]) || Connection.drivers.last
 
+      case options[:tcp_keepalive]
+      when Hash
+        [:time, :intvl, :probes].each do |key|
+          unless options[:tcp_keepalive][key].is_a?(Fixnum)
+            raise "Expected the #{key.inspect} key in :tcp_keepalive to be a Fixnum"
+          end
+        end
+
+      when Fixnum
+        if options[:tcp_keepalive] >= 60
+          options[:tcp_keepalive] = {:time => options[:tcp_keepalive] - 20, :intvl => 10, :probes => 2}
+
+        elsif options[:tcp_keepalive] >= 30
+          options[:tcp_keepalive] = {:time => options[:tcp_keepalive] - 10, :intvl => 5, :probes => 2}
+
+        elsif options[:tcp_keepalive] >= 5
+          options[:tcp_keepalive] = {:time => options[:tcp_keepalive] - 2, :intvl => 2, :probes => 1}
+        end
+      end
+
       options
     end
 

diff --git a/lib/redis/connection/ruby.rb b/lib/redis/connection/ruby.rb
@@ -177,9 +177,26 @@ def self.connect(config)
 
         instance = new(sock)
         instance.timeout = config[:timeout]
+        instance.set_tcp_keepalive config[:tcp_keepalive]
         instance
       end
 
+      if [:SOL_SOCKET, :SO_KEEPALIVE, :SOL_TCP, :TCP_KEEPIDLE, :TCP_KEEPINTVL, :TCP_KEEPCNT].all?{|c| Socket.const_defined? c}
+        def set_tcp_keepalive(keepalive)
+          return unless keepalive.is_a?(Hash)
+
+          @sock.setsockopt(Socket::SOL_SOCKET, Socket::SO_KEEPALIVE, true)
+          @sock.setsockopt(Socket::SOL_TCP, Socket::TCP_KEEPIDLE, keepalive[:time])
+          @sock.setsockopt(Socket::SOL_TCP, Socket::TCP_KEEPINTVL, keepalive[:intvl])
+          @sock.setsockopt(Socket::SOL_TCP, Socket::TCP_KEEPCNT, keepalive[:probes])
+        end
+
+      else
+        def set_tcp_keepalive(keepalive)
+          # NO-OP
+        end
+      end
+
       def initialize(sock)
         @sock = sock
       end

diff --git a/test/internals_test.rb b/test/internals_test.rb
@@ -98,6 +98,22 @@ def test_timeout
     end
   end
 
+  driver(:ruby) do
+    def test_tcp_keepalive
+      # Can only be tested on systems that support per-socket keepalive parameters
+      return unless [:SOL_SOCKET, :SO_KEEPALIVE, :SOL_TCP, :TCP_KEEPIDLE, :TCP_KEEPINTVL, :TCP_KEEPCNT].all?{|c| Socket.const_defined? c}
+
+      keepalive = {:time => 20, :intvl => 10, :probes => 5}
+
+      redis = Redis.new(OPTIONS.merge(:tcp_keepalive => keepalive))
+      redis.ping
+
+      assert_equal redis.instance_variable_get(:"@client").connection.instance_variable_get(:"@sock").getsockopt(Socket::SOL_TCP, Socket::TCP_KEEPIDLE).int, keepalive[:time]
+      assert_equal redis.instance_variable_get(:"@client").connection.instance_variable_get(:"@sock").getsockopt(Socket::SOL_TCP, Socket::TCP_KEEPINTVL).int, keepalive[:intvl]
+      assert_equal redis.instance_variable_get(:"@client").connection.instance_variable_get(:"@sock").getsockopt(Socket::SOL_TCP, Socket::TCP_KEEPCNT).int, keepalive[:probes]
+    end
+  end
+
   def test_time
     return if version < "2.5.4"