Skip to content

Commit

Permalink
Memcache auto discovery
Browse files Browse the repository at this point in the history
  • Loading branch information
Vedanta Jha committed Oct 4, 2023
1 parent 8cd0d87 commit ea2e6da
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 3 deletions.
9 changes: 9 additions & 0 deletions pymemcache/client/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -1030,6 +1030,15 @@ def flush_all(self, delay: int = 0, noreply: Optional[bool] = None) -> bool:
return True
return results[0] == b"OK"

def auto_discover(self):
cmd = b"config get cluster"
data = self._misc_cmd([cmd], b"config get cluster", noreply=False)
lines = data.split(b'\n')
configs = [conf.split(b'|') for conf in lines[2].split(b' ')]
self.quit()
nodes = [(ip, int(port)) for host, ip, port in configs]
return nodes

def quit(self) -> None:
"""
The memcached "quit" command.
Expand Down
21 changes: 18 additions & 3 deletions pymemcache/client/hash.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,13 +48,15 @@ def __init__(
default_noreply=True,
encoding="ascii",
tls_context=None,
enable_auto_discovery=False
):
"""
Constructor.
Args:
servers: list() of tuple(hostname, port) or string containing a UNIX
socket path.
If enable_auto_discovery is set, just a configuration endpoint would suffice
hasher: optional class three functions ``get_node``, ``add_node``,
and ``remove_node``
defaults to Rendezvous (HRW) hash.
Expand All @@ -70,12 +72,14 @@ def __init__(
dead_timeout (float): Time in seconds before attempting to add a node
back in the pool.
encoding: optional str, controls data encoding (defaults to 'ascii').
enable_auto_discovery (bool): If enabled, nodes would be discovered from the configuration endpoint.
Further arguments are interpreted as for :py:class:`.Client`
constructor.
"""
self.clients = {}
self.retry_attempts = retry_attempts
self.connect_timeout = connect_timeout
self.retry_timeout = retry_timeout
self.dead_timeout = dead_timeout
self.use_pooling = use_pooling
Expand Down Expand Up @@ -112,11 +116,14 @@ def __init__(
"lock_generator": lock_generator,
}
)

for server in servers:
self.add_server(normalize_server_spec(server))
self.encoding = encoding
self.tls_context = tls_context
if not isinstance(servers, list):
if enable_auto_discovery:
servers = self._auto_discover(server)

for server in servers:
self.add_server(normalize_server_spec(server))

def _make_client_key(self, server):
if isinstance(server, (list, tuple)) and len(server) == 2:
Expand Down Expand Up @@ -340,6 +347,14 @@ def _set_many(self, client, values, *args, **kwargs):
succeeded = [key for key in values if key not in failed]
return succeeded, failed, None

def _auto_discover(self, configuration_endpoint, tls_context):
host, port = configuration_endpoint.split(':')
server = (host, port)
_class = PooledClient if self.use_pooling else self.client_class
client = _class(server)
nodes = client.auto_discover()
return nodes

def close(self):
for client in self.clients.values():
self._safely_run_func(client, client.close, False)
Expand Down

0 comments on commit ea2e6da

Please sign in to comment.