diff --git a/riak/transports/connection.py b/riak/transports/connection.py index d26b85f4..cfa23f6e 100644 --- a/riak/transports/connection.py +++ b/riak/transports/connection.py @@ -152,6 +152,9 @@ def _new_connection(self): class Socket(object): + # Allow extending classes to affect how sockets are created through options + _options = {} + def __init__(self, host, port): self.host = host self.port = port @@ -161,6 +164,8 @@ def __init__(self, host, port): def maybe_connect(self): if self.sock is None: self.sock = s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + if "timeout" in self._options: + self.sock.settimeout(self._options["timeout"]) try: s.connect((self.host, self.port)) diff --git a/riak/transports/pbc.py b/riak/transports/pbc.py index 19afe85f..706f6963 100644 --- a/riak/transports/pbc.py +++ b/riak/transports/pbc.py @@ -22,6 +22,7 @@ import errno import socket import struct +import types try: import json @@ -160,6 +161,13 @@ class RiakPbcTransport(RiakTransport): # The ConnectionManager class that this transport prefers. default_cm = connection.cm_using(SocketWithId) + @staticmethod + def with_timeout(timeout): + socket_with_timeout = type("SocketWithIdAndTimeout", (SocketWithId,), {"_options":dict(SocketWithId._options)}) + socket_with_timeout._options["timeout"] = timeout + cm_with_timeout = connection.cm_using(socket_with_timeout) + return type("RiackPbcTransportWithTimeout", (RiakPbcTransport,), {"default_cm":cm_with_timeout}) + def __init__(self, cm, client_id=None, max_attempts=1, **unused_options): """ Construct a new RiakPbcTransport object.