|
| 1 | +import logging |
| 2 | +import operator |
| 3 | +import socket |
| 4 | +import re |
| 5 | +import time |
| 6 | + |
| 7 | +from pymemcache import MemcacheUnknownCommandError |
| 8 | +from pymemcache.client import Client |
| 9 | +from pymemcache.client.base import normalize_server_spec |
| 10 | +from pymemcache.client.hash import HashClient |
| 11 | +from pymemcache.client.rendezvous import RendezvousHash |
| 12 | + |
| 13 | + |
| 14 | +logger = logging.getLogger(__name__) |
| 15 | + |
| 16 | +_RE_AWS_ENDPOINT = re.compile( |
| 17 | + r"^(?:(?:[\w\d-]{0,61}[\w\d]\.)+[\w]{1,6}|\[(?:[\d]{1,3}\.){3}[\d]{1,3}\])\:\d{1,5}$" |
| 18 | +) |
| 19 | + |
| 20 | + |
| 21 | +class AWSElastiCacheHashClient(HashClient): |
| 22 | + """ |
| 23 | + This class is a subclass of HashClient and represents a client for interacting with an AWS ElastiCache cluster |
| 24 | + using a hash-based algorithm for key distribution. |
| 25 | +
|
| 26 | + *Connection * |
| 27 | +
|
| 28 | + Supports version 1.4.14 or higher |
| 29 | +
|
| 30 | + Example: |
| 31 | + >>> client = AWSElastiCacheServerlessClient('cluster.abcxyz.cfg.use1.cache.amazonaws.com') |
| 32 | + """ |
| 33 | + |
| 34 | + def __init__( |
| 35 | + self, |
| 36 | + cfg_node: object, |
| 37 | + hasher: object = RendezvousHash, |
| 38 | + serde: object = None, |
| 39 | + serializer: object = None, |
| 40 | + deserializer: object = None, |
| 41 | + connect_timeout: object = None, |
| 42 | + timeout: object = None, |
| 43 | + no_delay: object = False, |
| 44 | + socket_module: object = socket, |
| 45 | + socket_keepalive: object = None, |
| 46 | + key_prefix: object = b"", |
| 47 | + max_pool_size: object = None, |
| 48 | + pool_idle_timeout: object = 0, |
| 49 | + lock_generator: object = None, |
| 50 | + retry_attempts: object = 2, |
| 51 | + retry_timeout: object = 1, |
| 52 | + dead_timeout: object = 60, |
| 53 | + use_pooling: object = False, |
| 54 | + ignore_exc: object = False, |
| 55 | + allow_unicode_keys: object = False, |
| 56 | + default_noreply: object = True, |
| 57 | + encoding: object = "ascii", |
| 58 | + tls_context: object = None, |
| 59 | + use_vpc: object = True, |
| 60 | + ) -> object: |
| 61 | + """ |
| 62 | + Constructor. |
| 63 | +
|
| 64 | + Args: |
| 65 | + cfg_node: formatted string containing endpoint and port of the |
| 66 | + ElastiCache cluster endpoint. Ex.: |
| 67 | + `test-cluster.2os1zk.cfg.use1.cache.amazonaws.com:11211` |
| 68 | + serde: optional serializer object, see notes in the class docs. |
| 69 | + serializer: deprecated serialization function |
| 70 | + deserializer: deprecated deserialization function |
| 71 | + connect_timeout: optional float, seconds to wait for a connection to |
| 72 | + the memcached server. Defaults to "forever" (uses the underlying |
| 73 | + default socket timeout, which can be very long). |
| 74 | + timeout: optional float, seconds to wait for send or recv calls on |
| 75 | + the socket connected to memcached. Defaults to "forever" (uses the |
| 76 | + underlying default socket timeout, which can be very long). |
| 77 | + no_delay: optional bool, set the TCP_NODELAY flag, which may help |
| 78 | + with performance in some cases. Defaults to False. |
| 79 | + ignore_exc: optional bool, True to cause the "get", "gets", |
| 80 | + "get_many" and "gets_many" calls to treat any errors as cache |
| 81 | + misses. Defaults to False. |
| 82 | + socket_module: socket module to use, e.g. gevent.socket. Defaults to |
| 83 | + the standard library's socket module. |
| 84 | + socket_keepalive: Activate the socket keepalive feature by passing |
| 85 | + a KeepaliveOpts structure in this parameter. Disabled by default |
| 86 | + (None). This feature is only supported on Linux platforms. |
| 87 | + key_prefix: Prefix of key. You can use this as namespace. Defaults |
| 88 | + to b''. |
| 89 | + default_noreply: bool, the default value for 'noreply' as passed to |
| 90 | + store commands (except from cas, incr, and decr, which default to |
| 91 | + False). |
| 92 | + allow_unicode_keys: bool, support unicode (utf8) keys |
| 93 | + encoding: optional str, controls data encoding (defaults to 'ascii'). |
| 94 | + use_vpc: optional bool, if set False (defaults to True), the client |
| 95 | + will use FQDN to connect to nodes instead of IP addresses. See |
| 96 | + AWS Docs for extra info |
| 97 | + https://docs.aws.amazon.com/AmazonElastiCache/latest/mem-ug/ClientConfig.DNS.html |
| 98 | +
|
| 99 | + Notes: |
| 100 | + The constructor does not make a connection to memcached. The first |
| 101 | + call to a method on the object will do that. |
| 102 | + """ |
| 103 | + if not (_RE_AWS_ENDPOINT.fullmatch(cfg_node) and isinstance(cfg_node, str)): |
| 104 | + raise ValueError("Invalid AWS ElastiCache endpoint value '%s'" % cfg_node) |
| 105 | + |
| 106 | + self._cfg_node = cfg_node |
| 107 | + self.clients = {} |
| 108 | + self.retry_attempts = retry_attempts |
| 109 | + self.retry_timeout = retry_timeout |
| 110 | + self.dead_timeout = dead_timeout |
| 111 | + self.use_pooling = use_pooling |
| 112 | + self.key_prefix = key_prefix |
| 113 | + self.ignore_exc = ignore_exc |
| 114 | + self.allow_unicode_keys = allow_unicode_keys |
| 115 | + self._failed_clients = {} |
| 116 | + self._dead_clients = {} |
| 117 | + self._last_dead_check_time = time.time() |
| 118 | + |
| 119 | + self.hasher = hasher() |
| 120 | + |
| 121 | + self.default_kwargs = { |
| 122 | + "connect_timeout": connect_timeout, |
| 123 | + "timeout": timeout, |
| 124 | + "no_delay": no_delay, |
| 125 | + "socket_module": socket_module, |
| 126 | + "socket_keepalive": socket_keepalive, |
| 127 | + "key_prefix": key_prefix, |
| 128 | + "serde": serde, |
| 129 | + "serializer": serializer, |
| 130 | + "deserializer": deserializer, |
| 131 | + "allow_unicode_keys": allow_unicode_keys, |
| 132 | + "default_noreply": default_noreply, |
| 133 | + "encoding": encoding, |
| 134 | + "tls_context": tls_context, |
| 135 | + } |
| 136 | + |
| 137 | + if use_pooling is True: |
| 138 | + self.default_kwargs.update( |
| 139 | + { |
| 140 | + "max_pool_size": max_pool_size, |
| 141 | + "pool_idle_timeout": pool_idle_timeout, |
| 142 | + "lock_generator": lock_generator, |
| 143 | + } |
| 144 | + ) |
| 145 | + |
| 146 | + # server config returns as `[fqdn, ip, port]` if it's VPC installation you need to use ip |
| 147 | + self._use_vpc = int(use_vpc) |
| 148 | + |
| 149 | + self.reconfigure_nodes() |
| 150 | + |
| 151 | + self.encoding = encoding |
| 152 | + self.tls_context = tls_context |
| 153 | + |
| 154 | + def reconfigure_nodes(self): |
| 155 | + """ |
| 156 | + Reconfigures the nodes in the server cluster based on the provided configuration node. |
| 157 | +
|
| 158 | + May useful on error handling during cluster scale down or scale up |
| 159 | + """ |
| 160 | + old_clients = self.clients.copy() |
| 161 | + self.clients.clear() |
| 162 | + |
| 163 | + for server in self._get_nodes_list(): |
| 164 | + self.add_server(normalize_server_spec(server)) |
| 165 | + |
| 166 | + for client in old_clients.values(): |
| 167 | + client.close() |
| 168 | + |
| 169 | + def _get_nodes_list(self) -> list[tuple[str, int]]: |
| 170 | + """ |
| 171 | + Get the list of nodes from the cluster configuration. |
| 172 | +
|
| 173 | + Returns: |
| 174 | + A list of tuples containing the address and port of each node in the cluster. |
| 175 | + Each tuple has the format (address: str, port: int). |
| 176 | + """ |
| 177 | + addr, port = self._cfg_node.rsplit(":", maxsplit=1) |
| 178 | + client = Client((addr, port), **self.default_kwargs) |
| 179 | + |
| 180 | + # https://docs.aws.amazon.com/AmazonElastiCache/latest/mem-ug/AutoDiscovery.AddingToYourClientLibrary.html |
| 181 | + try: |
| 182 | + *_, config_line = client.raw_command( |
| 183 | + b"config get cluster", |
| 184 | + end_tokens=b"\n\r\nEND\r\n", |
| 185 | + ).splitlines() |
| 186 | + except MemcacheUnknownCommandError: |
| 187 | + logger.exception( |
| 188 | + "Can't retrieve cluster configuration from '%s:%s' " |
| 189 | + "Seems like it is ElastiCache Serverless or even isn't ElastiCache at all.", |
| 190 | + client.server, |
| 191 | + ) |
| 192 | + finally: |
| 193 | + client.close() |
| 194 | + |
| 195 | + servers = [ |
| 196 | + (server[self._use_vpc], server[2]) |
| 197 | + for server in map( |
| 198 | + operator.methodcaller("split", "|"), |
| 199 | + config_line.decode().split(" "), |
| 200 | + ) |
| 201 | + ] |
| 202 | + |
| 203 | + logger.debug("Got the next nodes from cluster config: %s", servers) |
| 204 | + |
| 205 | + return servers |
0 commit comments