herd.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164
  1. import random
  2. import socket
  3. import time
  4. from collections import OrderedDict
  5. from django.conf import settings
  6. from redis.exceptions import ConnectionError, ResponseError, TimeoutError
  7. from ..exceptions import ConnectionInterrupted
  8. from .default import DEFAULT_TIMEOUT, DefaultClient
  9. _main_exceptions = (ConnectionError, ResponseError, TimeoutError, socket.timeout)
  10. class Marker:
  11. """
  12. Dummy class for use as
  13. marker for herded keys.
  14. """
  15. pass
  16. def _is_expired(x, herd_timeout: int) -> bool:
  17. if x >= herd_timeout:
  18. return True
  19. val = x + random.randint(1, herd_timeout)
  20. if val >= herd_timeout:
  21. return True
  22. return False
  23. class HerdClient(DefaultClient):
  24. def __init__(self, *args, **kwargs):
  25. self._marker = Marker()
  26. self._herd_timeout = getattr(settings, "CACHE_HERD_TIMEOUT", 60)
  27. super().__init__(*args, **kwargs)
  28. def _pack(self, value, timeout):
  29. herd_timeout = (timeout or self._backend.default_timeout) + int(time.time())
  30. return self._marker, value, herd_timeout
  31. def _unpack(self, value):
  32. try:
  33. marker, unpacked, herd_timeout = value
  34. except (ValueError, TypeError):
  35. return value, False
  36. if not isinstance(marker, Marker):
  37. return value, False
  38. now = int(time.time())
  39. if herd_timeout < now:
  40. x = now - herd_timeout
  41. return unpacked, _is_expired(x, self._herd_timeout)
  42. return unpacked, False
  43. def set(
  44. self,
  45. key,
  46. value,
  47. timeout=DEFAULT_TIMEOUT,
  48. version=None,
  49. client=None,
  50. nx=False,
  51. xx=False,
  52. ):
  53. if timeout is DEFAULT_TIMEOUT:
  54. timeout = self._backend.default_timeout
  55. if timeout is None or timeout <= 0:
  56. return super().set(
  57. key,
  58. value,
  59. timeout=timeout,
  60. version=version,
  61. client=client,
  62. nx=nx,
  63. xx=xx,
  64. )
  65. packed = self._pack(value, timeout)
  66. real_timeout = timeout + self._herd_timeout
  67. return super().set(
  68. key, packed, timeout=real_timeout, version=version, client=client, nx=nx
  69. )
  70. def get(self, key, default=None, version=None, client=None):
  71. packed = super().get(key, default=default, version=version, client=client)
  72. val, refresh = self._unpack(packed)
  73. if refresh:
  74. return default
  75. return val
  76. def get_many(self, keys, version=None, client=None):
  77. if client is None:
  78. client = self.get_client(write=False)
  79. if not keys:
  80. return {}
  81. recovered_data = OrderedDict()
  82. new_keys = [self.make_key(key, version=version) for key in keys]
  83. map_keys = dict(zip(new_keys, keys))
  84. try:
  85. results = client.mget(*new_keys)
  86. except _main_exceptions as e:
  87. raise ConnectionInterrupted(connection=client) from e
  88. for key, value in zip(new_keys, results):
  89. if value is None:
  90. continue
  91. val, refresh = self._unpack(self.decode(value))
  92. recovered_data[map_keys[key]] = None if refresh else val
  93. return recovered_data
  94. def set_many(
  95. self, data, timeout=DEFAULT_TIMEOUT, version=None, client=None, herd=True
  96. ):
  97. """
  98. Set a bunch of values in the cache at once from a dict of key/value
  99. pairs. This is much more efficient than calling set() multiple times.
  100. If timeout is given, that timeout will be used for the key; otherwise
  101. the default cache timeout will be used.
  102. """
  103. if client is None:
  104. client = self.get_client(write=True)
  105. set_function = self.set if herd else super().set
  106. try:
  107. pipeline = client.pipeline()
  108. for key, value in data.items():
  109. set_function(key, value, timeout, version=version, client=pipeline)
  110. pipeline.execute()
  111. except _main_exceptions as e:
  112. raise ConnectionInterrupted(connection=client) from e
  113. def incr(self, *args, **kwargs):
  114. raise NotImplementedError()
  115. def decr(self, *args, **kwargs):
  116. raise NotImplementedError()
  117. def touch(self, key, timeout=DEFAULT_TIMEOUT, version=None, client=None):
  118. if client is None:
  119. client = self.get_client(write=True)
  120. value = self.get(key, version=version, client=client)
  121. if value is None:
  122. return False
  123. self.set(key, value, timeout=timeout, version=version, client=client)
  124. return True