redis-allocator 0.0.1__py3-none-any.whl → 0.3.1__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- redis_allocator/__init__.py +5 -1
- redis_allocator/_version.py +1 -0
- redis_allocator/allocator.py +816 -280
- redis_allocator/lock.py +66 -17
- redis_allocator/task_queue.py +81 -57
- redis_allocator-0.3.1.dist-info/METADATA +529 -0
- redis_allocator-0.3.1.dist-info/RECORD +15 -0
- {redis_allocator-0.0.1.dist-info → redis_allocator-0.3.1.dist-info}/licenses/LICENSE +21 -21
- tests/conftest.py +160 -46
- tests/test_allocator.py +461 -488
- tests/test_lock.py +675 -338
- tests/test_task_queue.py +136 -136
- redis_allocator-0.0.1.dist-info/METADATA +0 -229
- redis_allocator-0.0.1.dist-info/RECORD +0 -14
- {redis_allocator-0.0.1.dist-info → redis_allocator-0.3.1.dist-info}/WHEEL +0 -0
- {redis_allocator-0.0.1.dist-info → redis_allocator-0.3.1.dist-info}/top_level.txt +0 -0
redis_allocator/allocator.py
CHANGED
@@ -1,45 +1,40 @@
|
|
1
1
|
"""Redis-based distributed memory allocation system.
|
2
2
|
|
3
|
-
This module provides the core functionality of the RedisAllocator system,
|
3
|
+
This module provides the core functionality of the RedisAllocator system,
|
4
4
|
allowing for distributed memory allocation with support for garbage collection,
|
5
5
|
thread health checking, and priority-based allocation mechanisms.
|
6
|
+
|
7
|
+
Key features:
|
8
|
+
1. Shared vs non-shared allocation modes:
|
9
|
+
- In shared mode, allocating an item simply removes it from the free list and puts it back to the tail
|
10
|
+
- In non-shared mode, allocation locks the item to prevent others from accessing it
|
11
|
+
2. Garbage collection for stale/unhealthy items:
|
12
|
+
- Items that are locked (unhealthy) but in the free list are removed
|
13
|
+
- Items that are not in the free list but haven't been updated within their timeout are freed
|
14
|
+
3. Soft binding mechanism:
|
15
|
+
- Maps object names to allocated keys for consistent allocation
|
16
|
+
- Prioritizes previously allocated keys when the same named object requests allocation
|
17
|
+
4. Support for an updater to refresh the pool's keys periodically
|
18
|
+
5. Policy-based control of allocation behavior through RedisAllocatorPolicy
|
6
19
|
"""
|
20
|
+
import atexit
|
7
21
|
import logging
|
8
22
|
import weakref
|
23
|
+
import contextlib
|
9
24
|
from abc import ABC, abstractmethod
|
25
|
+
from typing import Any, Callable
|
10
26
|
from functools import cached_property
|
11
27
|
from threading import current_thread
|
28
|
+
from concurrent.futures import ThreadPoolExecutor
|
12
29
|
from typing import (Optional, TypeVar, Generic,
|
13
30
|
Sequence, Iterable)
|
14
31
|
from redis import StrictRedis as Redis
|
32
|
+
from cachetools import cached, TTLCache
|
15
33
|
from .lock import RedisLockPool, Timeout
|
16
34
|
|
17
35
|
logger = logging.getLogger(__name__)
|
18
36
|
|
19
37
|
|
20
|
-
class RedisAllocatableClass(ABC):
|
21
|
-
@abstractmethod
|
22
|
-
def set_config(self, key: str, params: dict):
|
23
|
-
"""Set the configuration for the object.
|
24
|
-
|
25
|
-
Args:
|
26
|
-
key: The key to set the configuration for.
|
27
|
-
params: The parameters to set the configuration for.
|
28
|
-
"""
|
29
|
-
pass
|
30
|
-
|
31
|
-
def close(self):
|
32
|
-
"""close the object."""
|
33
|
-
pass
|
34
|
-
|
35
|
-
def name(self) -> Optional[str]:
|
36
|
-
"""Get the cache name of the object."""
|
37
|
-
return None
|
38
|
-
|
39
|
-
|
40
|
-
U = TypeVar('U', bound=RedisAllocatableClass)
|
41
|
-
|
42
|
-
|
43
38
|
class RedisThreadHealthCheckPool(RedisLockPool):
|
44
39
|
"""A class that provides a simple interface for managing the health status of a thread.
|
45
40
|
|
@@ -84,14 +79,54 @@ class RedisThreadHealthCheckPool(RedisLockPool):
|
|
84
79
|
self.unlock(self.current_thread_id)
|
85
80
|
|
86
81
|
|
82
|
+
class RedisAllocatableClass(ABC):
|
83
|
+
"""A class that can be allocated through RedisAllocator.
|
84
|
+
|
85
|
+
You should inherit from this class and implement the set_config method.
|
86
|
+
"""
|
87
|
+
|
88
|
+
@abstractmethod
|
89
|
+
def set_config(self, key: str, params: dict):
|
90
|
+
"""Set the configuration for the object.
|
91
|
+
|
92
|
+
Args:
|
93
|
+
key: The key to set the configuration for.
|
94
|
+
params: The parameters to set the configuration for.
|
95
|
+
"""
|
96
|
+
pass
|
97
|
+
|
98
|
+
def open(self):
|
99
|
+
"""Open the object."""
|
100
|
+
pass
|
101
|
+
|
102
|
+
def close(self):
|
103
|
+
"""close the object."""
|
104
|
+
pass
|
105
|
+
|
106
|
+
def is_healthy(self):
|
107
|
+
return True
|
108
|
+
|
109
|
+
@property
|
110
|
+
def name(self) -> Optional[str]:
|
111
|
+
"""Get the cache name of the object, if is none no soft binding will be used."""
|
112
|
+
return None
|
113
|
+
|
114
|
+
@property
|
115
|
+
def unique_id(self) -> str:
|
116
|
+
"""Get the unique ID of the object."""
|
117
|
+
return ""
|
118
|
+
|
119
|
+
|
120
|
+
U = TypeVar('U', bound=RedisAllocatableClass)
|
121
|
+
|
122
|
+
|
87
123
|
class RedisAllocatorObject(Generic[U]):
|
88
124
|
"""Represents an object allocated through RedisAllocator.
|
89
125
|
|
90
126
|
This class provides an interface for working with allocated objects
|
91
127
|
including locking and unlocking mechanisms for thread-safe operations.
|
92
|
-
|
93
128
|
"""
|
94
|
-
|
129
|
+
allocator: 'RedisAllocator' # Reference to the allocator that created this object
|
95
130
|
key: str # Redis key for this allocated object
|
96
131
|
params: Optional[dict] # Parameters associated with this object
|
97
132
|
obj: Optional[U] # The actual object being allocated
|
@@ -105,7 +140,7 @@ class RedisAllocatorObject(Generic[U]):
|
|
105
140
|
obj: The actual object being allocated
|
106
141
|
params: Additional parameters passed by local program
|
107
142
|
"""
|
108
|
-
self.
|
143
|
+
self.allocator = allocator
|
109
144
|
self.key = key
|
110
145
|
self.obj = obj
|
111
146
|
self.params = params
|
@@ -119,47 +154,358 @@ class RedisAllocatorObject(Generic[U]):
|
|
119
154
|
timeout: How long the lock should be valid (in seconds)
|
120
155
|
"""
|
121
156
|
if timeout > 0:
|
122
|
-
self.
|
157
|
+
self.allocator.update(self.key, timeout=timeout)
|
123
158
|
else:
|
124
|
-
self.
|
159
|
+
self.allocator.free(self)
|
160
|
+
|
161
|
+
def open(self):
|
162
|
+
"""Open the object."""
|
163
|
+
if self.obj is not None:
|
164
|
+
self.obj.open()
|
125
165
|
|
126
166
|
def close(self):
|
127
167
|
"""Kill the object."""
|
128
168
|
if self.obj is not None:
|
129
169
|
self.obj.close()
|
130
170
|
|
171
|
+
def is_healthy(self) -> bool:
|
172
|
+
"""Check if the object is healthy."""
|
173
|
+
if self.obj is not None:
|
174
|
+
return self.obj.is_healthy()
|
175
|
+
return True
|
176
|
+
|
177
|
+
def set_unhealthy(self, duration: int = 3600):
|
178
|
+
"""Set the object as unhealthy."""
|
179
|
+
if self.obj is not None and self.obj.name is not None:
|
180
|
+
self.allocator.unbind_soft_bind(self.obj.name)
|
181
|
+
self.allocator.update(self.key, timeout=duration)
|
182
|
+
|
183
|
+
def refresh(self, timeout: Timeout = 120):
|
184
|
+
"""Refresh the object."""
|
185
|
+
self.close()
|
186
|
+
new_obj = self.allocator.policy.malloc(self.allocator, timeout=timeout,
|
187
|
+
obj=self.obj, params=self.params)
|
188
|
+
if new_obj is not None:
|
189
|
+
self.obj = new_obj.obj
|
190
|
+
self.key = new_obj.key
|
191
|
+
self.params = new_obj.params
|
192
|
+
|
193
|
+
@property
|
194
|
+
def unique_id(self) -> str:
|
195
|
+
"""Get the unique ID of the object."""
|
196
|
+
if self.obj is None:
|
197
|
+
return self.key
|
198
|
+
return f"{self.key}:{self.obj.unique_id}"
|
199
|
+
|
200
|
+
@property
|
201
|
+
def name(self) -> Optional[str]:
|
202
|
+
"""Get the name of the object."""
|
203
|
+
if self.obj is None:
|
204
|
+
return None
|
205
|
+
return self.obj.name
|
206
|
+
|
131
207
|
def __del__(self):
|
132
208
|
"""Delete the object."""
|
133
209
|
self.close()
|
134
210
|
|
135
211
|
|
212
|
+
class RedisAllocatorUpdater:
|
213
|
+
"""A class that updates the allocator keys."""
|
214
|
+
|
215
|
+
def __init__(self, params: Sequence[Any]):
|
216
|
+
"""Initialize the allocator updater."""
|
217
|
+
assert len(params) > 0, "params should not be empty"
|
218
|
+
self.params = params
|
219
|
+
self.index = 0
|
220
|
+
|
221
|
+
@abstractmethod
|
222
|
+
def fetch(self, param: Any) -> Sequence[str]:
|
223
|
+
"""Fetch the keys from params."""
|
224
|
+
pass
|
225
|
+
|
226
|
+
def __call__(self):
|
227
|
+
"""Update the allocator key."""
|
228
|
+
current_param = self.params[self.index]
|
229
|
+
self.index = (self.index + 1) % len(self.params)
|
230
|
+
keys = self.fetch(current_param)
|
231
|
+
return keys
|
232
|
+
|
233
|
+
def __len__(self):
|
234
|
+
"""Get the length of the allocator updater."""
|
235
|
+
return len(self.params)
|
236
|
+
|
237
|
+
|
238
|
+
class RedisAllocatorPolicy(ABC):
|
239
|
+
"""Abstract base class for Redis allocator policies.
|
240
|
+
|
241
|
+
This class defines the interface for allocation policies that can be used
|
242
|
+
with RedisAllocator to control allocation behavior.
|
243
|
+
"""
|
244
|
+
|
245
|
+
def initialize(self, allocator: 'RedisAllocator'):
|
246
|
+
"""Initialize the policy with an allocator instance.
|
247
|
+
|
248
|
+
Args:
|
249
|
+
allocator: The RedisAllocator instance to use with this policy
|
250
|
+
"""
|
251
|
+
pass
|
252
|
+
|
253
|
+
@abstractmethod
|
254
|
+
def malloc(self, allocator: 'RedisAllocator', timeout: Timeout = 120,
|
255
|
+
obj: Optional[Any] = None, params: Optional[dict] = None,
|
256
|
+
cache_timeout: Timeout = 3600) -> Optional[RedisAllocatorObject]:
|
257
|
+
"""Allocate a resource according to the policy.
|
258
|
+
|
259
|
+
Args:
|
260
|
+
allocator: The RedisAllocator instance
|
261
|
+
timeout: How long the allocation should be valid (in seconds)
|
262
|
+
obj: The object to associate with the allocation
|
263
|
+
params: Additional parameters for the allocation
|
264
|
+
cache_timeout: Timeout for the soft binding cache entry (seconds).
|
265
|
+
Defaults to 3600.
|
266
|
+
|
267
|
+
Returns:
|
268
|
+
RedisAllocatorObject if allocation was successful, None otherwise
|
269
|
+
"""
|
270
|
+
pass
|
271
|
+
|
272
|
+
@abstractmethod
|
273
|
+
def refresh_pool(self, allocator: 'RedisAllocator'):
|
274
|
+
"""Refresh the allocation pool.
|
275
|
+
|
276
|
+
This method is called periodically to update the pool with new resources.
|
277
|
+
|
278
|
+
Args:
|
279
|
+
allocator: The RedisAllocator instance
|
280
|
+
"""
|
281
|
+
pass
|
282
|
+
|
283
|
+
def check_health_once(self, r_obj: RedisAllocatorObject, duration: int = 3600) -> bool:
|
284
|
+
"""Check the health of the object."""
|
285
|
+
with contextlib.closing(r_obj):
|
286
|
+
try:
|
287
|
+
if r_obj.is_healthy():
|
288
|
+
return True
|
289
|
+
else:
|
290
|
+
r_obj.set_unhealthy(duration)
|
291
|
+
return False
|
292
|
+
except Exception as e:
|
293
|
+
logger.error(f"Error checking health of {r_obj.key}: {e}")
|
294
|
+
r_obj.set_unhealthy(duration)
|
295
|
+
raise
|
296
|
+
|
297
|
+
def check_health(self, allocator: 'RedisAllocator', lock_duration: Timeout = 3600, max_threads: int = 8,
|
298
|
+
obj_fn: Optional[Callable[[str], Any]] = None,
|
299
|
+
params_fn: Optional[Callable[[str], dict]] = None) -> tuple[int, int]:
|
300
|
+
"""Check the health of the allocator.
|
301
|
+
|
302
|
+
Args:
|
303
|
+
allocator: The RedisAllocator instance
|
304
|
+
lock_duration: The duration of the lock (in seconds)
|
305
|
+
max_threads: The maximum number of threads to use
|
306
|
+
|
307
|
+
Returns:
|
308
|
+
A tuple containing the number of healthy and unhealthy items in the allocator
|
309
|
+
"""
|
310
|
+
with ThreadPoolExecutor(max_workers=max_threads) as executor:
|
311
|
+
for key in allocator.keys():
|
312
|
+
if params_fn is not None:
|
313
|
+
params = params_fn(key)
|
314
|
+
else:
|
315
|
+
params = None
|
316
|
+
if obj_fn is not None:
|
317
|
+
obj = obj_fn(key)
|
318
|
+
else:
|
319
|
+
obj = None
|
320
|
+
alloc_obj = RedisAllocatorObject(allocator, key, obj, params)
|
321
|
+
executor.submit(self.check_health_once, alloc_obj, lock_duration)
|
322
|
+
executor.shutdown(wait=True)
|
323
|
+
|
324
|
+
|
325
|
+
class DefaultRedisAllocatorPolicy(RedisAllocatorPolicy):
|
326
|
+
"""Default implementation of RedisAllocatorPolicy.
|
327
|
+
|
328
|
+
This policy provides the following features:
|
329
|
+
1. Garbage collection before allocation: Automatically performs garbage collection
|
330
|
+
operations before allocating resources to ensure stale resources are reclaimed.
|
331
|
+
|
332
|
+
2. Soft binding prioritization: Prioritizes allocation of previously bound keys
|
333
|
+
for named objects, creating a consistent mapping between object names and keys.
|
334
|
+
If a soft binding exists but the bound key is no longer in the pool, the binding is
|
335
|
+
ignored and a new key is allocated.
|
336
|
+
|
337
|
+
3. Periodic pool updates: Uses an optional updater to refresh the pool's keys at
|
338
|
+
configurable intervals. Only one process/thread (the one that acquires the update lock)
|
339
|
+
will perform the update.
|
340
|
+
|
341
|
+
4. Configurable expiry times: Allows setting default expiry durations for pool items,
|
342
|
+
ensuring automatic cleanup of stale resources even without explicit garbage collection.
|
343
|
+
|
344
|
+
The policy controls when garbage collection happens, when the pool is refreshed with new keys,
|
345
|
+
and how allocation prioritizes resources.
|
346
|
+
"""
|
347
|
+
|
348
|
+
def __init__(self, gc_count: int = 5, update_interval: int = 300,
|
349
|
+
expiry_duration: int = -1, updater: Optional[RedisAllocatorUpdater] = None):
|
350
|
+
"""Initialize the default allocation policy.
|
351
|
+
|
352
|
+
Args:
|
353
|
+
gc_count: Number of GC operations to perform before allocation
|
354
|
+
update_interval: Interval in seconds between pool updates
|
355
|
+
expiry_duration: Default timeout for pool items (-1 means no timeout)
|
356
|
+
updater: Optional updater for refreshing the pool's keys
|
357
|
+
"""
|
358
|
+
self.gc_count = gc_count
|
359
|
+
self.update_interval: float = update_interval
|
360
|
+
self.expiry_duration: float = expiry_duration
|
361
|
+
self.updater = updater
|
362
|
+
self._allocator: Optional[weakref.ReferenceType['RedisAllocator']] = None
|
363
|
+
self._update_lock_key: Optional[str] = None
|
364
|
+
self.objects: weakref.WeakValueDictionary[str, RedisAllocatorObject] = weakref.WeakValueDictionary()
|
365
|
+
|
366
|
+
def initialize(self, allocator: 'RedisAllocator'):
|
367
|
+
"""Initialize the policy with an allocator instance.
|
368
|
+
|
369
|
+
Args:
|
370
|
+
allocator: The RedisAllocator instance to use with this policy
|
371
|
+
"""
|
372
|
+
self._allocator = weakref.ref(allocator)
|
373
|
+
self._update_lock_key = f"{allocator._pool_str()}|policy_update_lock"
|
374
|
+
atexit.register(lambda: self.finalize(self._allocator()))
|
375
|
+
|
376
|
+
def malloc(self, allocator: 'RedisAllocator', timeout: Timeout = 120,
|
377
|
+
obj: Optional[Any] = None, params: Optional[dict] = None,
|
378
|
+
cache_timeout: Timeout = 3600) -> Optional[RedisAllocatorObject]:
|
379
|
+
"""Allocate a resource according to the policy.
|
380
|
+
|
381
|
+
This implementation:
|
382
|
+
1. Performs GC operations before allocation
|
383
|
+
2. Checks for soft binding based on object name
|
384
|
+
3. Falls back to regular allocation if no soft binding exists
|
385
|
+
|
386
|
+
Args:
|
387
|
+
allocator: The RedisAllocator instance
|
388
|
+
timeout: How long the allocation should be valid (in seconds)
|
389
|
+
obj: The object to associate with the allocation
|
390
|
+
params: Additional parameters for the allocation
|
391
|
+
cache_timeout: Timeout for the soft binding cache entry (seconds).
|
392
|
+
Defaults to 3600.
|
393
|
+
|
394
|
+
Returns:
|
395
|
+
RedisAllocatorObject if allocation was successful, None otherwise
|
396
|
+
"""
|
397
|
+
# Try to refresh the pool if necessary
|
398
|
+
self._try_refresh_pool(allocator)
|
399
|
+
|
400
|
+
# Perform GC operations before allocation
|
401
|
+
allocator.gc(self.gc_count)
|
402
|
+
|
403
|
+
# Fall back to regular allocation
|
404
|
+
# Explicitly call obj.name if obj exists
|
405
|
+
obj_name = obj.name if obj and hasattr(obj, 'name') and callable(
|
406
|
+
obj.name) else (obj.name if obj and hasattr(obj, 'name') else None)
|
407
|
+
key = allocator.malloc_key(timeout, obj_name,
|
408
|
+
cache_timeout=cache_timeout)
|
409
|
+
alloc_obj = RedisAllocatorObject(allocator, key, obj, params)
|
410
|
+
old_obj = self.objects.get(alloc_obj.unique_id, None)
|
411
|
+
if old_obj is not None:
|
412
|
+
old_obj.close()
|
413
|
+
self.objects[alloc_obj.unique_id] = alloc_obj
|
414
|
+
return alloc_obj
|
415
|
+
|
416
|
+
@cached(TTLCache(maxsize=64, ttl=5))
|
417
|
+
def _try_refresh_pool(self, allocator: 'RedisAllocator'):
|
418
|
+
"""Try to refresh the pool if necessary and if we can acquire the lock.
|
419
|
+
|
420
|
+
Args:
|
421
|
+
allocator: The RedisAllocator instance
|
422
|
+
"""
|
423
|
+
if self.updater is None:
|
424
|
+
return
|
425
|
+
if allocator.lock(self._update_lock_key, timeout=self.update_interval):
|
426
|
+
# If we got here, we acquired the lock, so we can update the pool
|
427
|
+
self.refresh_pool(allocator)
|
428
|
+
|
429
|
+
def refresh_pool(self, allocator: 'RedisAllocator'):
|
430
|
+
"""Refresh the allocation pool using the updater.
|
431
|
+
|
432
|
+
Args:
|
433
|
+
allocator: The RedisAllocator instance
|
434
|
+
"""
|
435
|
+
if self.updater is None:
|
436
|
+
return
|
437
|
+
|
438
|
+
keys = self.updater()
|
439
|
+
|
440
|
+
if len(keys) == 0:
|
441
|
+
logger.warning("No keys to update to the pool")
|
442
|
+
return
|
443
|
+
|
444
|
+
# Update the pool based on the number of keys
|
445
|
+
if len(self.updater) == 1:
|
446
|
+
allocator.assign(keys, timeout=self.expiry_duration)
|
447
|
+
else:
|
448
|
+
allocator.extend(keys, timeout=self.expiry_duration)
|
449
|
+
|
450
|
+
def finalize(self, allocator: 'RedisAllocator'):
|
451
|
+
"""Finalize the policy."""
|
452
|
+
for obj in self.objects.values():
|
453
|
+
obj.close()
|
454
|
+
|
455
|
+
|
136
456
|
class RedisAllocator(RedisLockPool, Generic[U]):
|
137
|
-
"""A Redis-based distributed
|
138
|
-
|
139
|
-
|
140
|
-
|
141
|
-
|
142
|
-
|
143
|
-
|
144
|
-
|
145
|
-
|
457
|
+
"""A Redis-based distributed allocation system.
|
458
|
+
|
459
|
+
Manages a pool of resource identifiers (keys) using Redis, allowing distributed
|
460
|
+
clients to allocate, free, and manage these resources. It leverages Redis's
|
461
|
+
atomic operations via Lua scripts for safe concurrent access.
|
462
|
+
|
463
|
+
The allocator maintains a doubly-linked list in a Redis hash to track available
|
464
|
+
(free) resources. Allocated resources are tracked using standard Redis keys
|
465
|
+
that act as locks.
|
466
|
+
|
467
|
+
Key Concepts:
|
468
|
+
- Allocation Pool: A set of resource identifiers (keys) managed by the allocator.
|
469
|
+
Stored in a Redis hash (`<prefix>|<suffix>|pool`) representing a doubly-linked list.
|
470
|
+
Head/Tail pointers are stored in separate keys (`<prefix>|<suffix>|pool|head`,
|
471
|
+
`<prefix>|<suffix>|pool|tail`).
|
472
|
+
- Free List: The subset of keys within the pool that are currently available.
|
473
|
+
Represented by the linked list structure within the pool hash.
|
474
|
+
- Allocated State: A key is considered allocated if a corresponding lock key exists
|
475
|
+
(`<prefix>|<suffix>:<key>`).
|
476
|
+
- Shared Mode: If `shared=True`, allocating a key moves it to the tail of the
|
477
|
+
free list but does *not* create a lock key. This allows multiple clients to
|
478
|
+
"allocate" the same key concurrently, effectively using the list as a rotating pool.
|
479
|
+
If `shared=False` (default), allocation creates a lock key, granting exclusive access.
|
480
|
+
- Soft Binding: Allows associating a logical name with an allocated key. If an object
|
481
|
+
provides a `name`, the allocator tries to reuse the previously bound key for that name.
|
482
|
+
Stored in Redis keys like `<prefix>|<suffix>-cache:bind:<name>`.
|
483
|
+
- Garbage Collection (GC): Periodically scans the pool to reconcile the free list
|
484
|
+
with the lock states. Removes expired/locked items from the free list and returns
|
485
|
+
items whose locks have expired back to the free list.
|
486
|
+
- Policies: Uses `RedisAllocatorPolicy` (e.g., `DefaultRedisAllocatorPolicy`)
|
487
|
+
to customize allocation behavior, GC triggering, and pool updates.
|
488
|
+
|
489
|
+
Generic type U should implement `RedisAllocatableClass`.
|
146
490
|
"""
|
147
491
|
|
148
492
|
def __init__(self, redis: Redis, prefix: str, suffix='allocator', eps=1e-6,
|
149
|
-
shared=False):
|
150
|
-
"""
|
151
|
-
|
493
|
+
shared=False, policy: Optional[RedisAllocatorPolicy] = None):
|
494
|
+
"""Initializes the RedisAllocator.
|
495
|
+
|
152
496
|
Args:
|
153
|
-
redis:
|
154
|
-
prefix: Prefix for all Redis keys used by this allocator
|
155
|
-
suffix: Suffix
|
156
|
-
eps: Small
|
157
|
-
shared:
|
497
|
+
redis: StrictRedis client instance (must decode responses).
|
498
|
+
prefix: Prefix for all Redis keys used by this allocator instance.
|
499
|
+
suffix: Suffix to uniquely identify this allocator instance's keys.
|
500
|
+
eps: Small float tolerance for comparisons (used by underlying lock).
|
501
|
+
shared: If True, operates in shared mode (keys are rotated, not locked).
|
502
|
+
If False (default), keys are locked upon allocation.
|
503
|
+
policy: Optional allocation policy. Defaults to `DefaultRedisAllocatorPolicy`.
|
158
504
|
"""
|
159
505
|
super().__init__(redis, prefix, suffix=suffix, eps=eps)
|
160
506
|
self.shared = shared
|
161
|
-
self.
|
162
|
-
self.
|
507
|
+
self.policy = policy or DefaultRedisAllocatorPolicy()
|
508
|
+
self.policy.initialize(self)
|
163
509
|
|
164
510
|
def object_key(self, key: str, obj: U):
|
165
511
|
"""Get the key for an object."""
|
@@ -169,35 +515,38 @@ class RedisAllocator(RedisLockPool, Generic[U]):
|
|
169
515
|
|
170
516
|
def _pool_pointer_str(self, head: bool = True):
|
171
517
|
"""Get the Redis key for the head or tail pointer of the allocation pool.
|
172
|
-
|
518
|
+
|
173
519
|
Args:
|
174
520
|
head: If True, get the head pointer key; otherwise, get the tail pointer key
|
175
|
-
|
521
|
+
|
176
522
|
Returns:
|
177
523
|
String representation of the Redis key for the pointer
|
178
524
|
"""
|
179
525
|
pointer_type = 'head' if head else 'tail'
|
180
526
|
return f'{self._pool_str()}|{pointer_type}'
|
181
527
|
|
182
|
-
def _gc_cursor_str(self):
|
183
|
-
"""Get the Redis key for the garbage collection cursor.
|
184
|
-
|
185
|
-
Returns:
|
186
|
-
String representation of the Redis key for the GC cursor
|
187
|
-
"""
|
188
|
-
return f'{self._pool_str()}|gc_cursor'
|
189
|
-
|
190
528
|
@property
|
191
529
|
def _lua_required_string(self):
|
192
|
-
"""
|
193
|
-
|
194
|
-
|
195
|
-
|
196
|
-
-
|
197
|
-
-
|
198
|
-
-
|
199
|
-
|
200
|
-
-
|
530
|
+
"""Base Lua script providing common functions for pool manipulation.
|
531
|
+
|
532
|
+
Includes functions inherited from RedisLockPool and adds allocator-specific ones:
|
533
|
+
- pool_pointer_str(head: bool): Returns the Redis key for the head/tail pointer.
|
534
|
+
- cache_str(): Returns the Redis key for the allocator's general cache.
|
535
|
+
- soft_bind_name(name: str): Returns the Redis key for a specific soft binding.
|
536
|
+
- split_pool_value(value: str): Parses the 'prev||next||expiry' string stored
|
537
|
+
in the pool hash for a key.
|
538
|
+
- join_pool_value(prev: str, next: str, expiry: int): Creates the value string.
|
539
|
+
- timeout_to_expiry(timeout: int): Converts relative seconds to absolute Unix timestamp.
|
540
|
+
- is_expiry_invalid(expiry: int): Checks if an absolute expiry time is in the past.
|
541
|
+
- is_expired(value: str): Checks if a pool item's expiry is in the past.
|
542
|
+
- push_to_tail(itemName: str, expiry: int): Adds/updates an item at the tail of the free list.
|
543
|
+
- pop_from_head(): Removes and returns the item from the head of the free list,
|
544
|
+
skipping expired or locked items. Returns (nil, -1) if empty.
|
545
|
+
- set_item_allocated(itemName: str): Removes an item from the free list structure.
|
546
|
+
- check_item_health(itemName: str, value: str|nil): Core GC logic for a single item.
|
547
|
+
- If item is marked #ALLOCATED but has no lock -> push to tail (return to free list).
|
548
|
+
- If item is in free list but expired -> remove from pool hash.
|
549
|
+
- If item is in free list but locked -> mark as #ALLOCATED (remove from free list).
|
201
550
|
"""
|
202
551
|
return f'''
|
203
552
|
{super()._lua_required_string}
|
@@ -208,185 +557,265 @@ class RedisAllocator(RedisLockPool, Generic[U]):
|
|
208
557
|
end
|
209
558
|
return '{self._pool_str()}|' .. pointer_type
|
210
559
|
end
|
211
|
-
local function
|
212
|
-
return '{self.
|
560
|
+
local function cache_str()
|
561
|
+
return '{self._cache_str}'
|
562
|
+
end
|
563
|
+
local function soft_bind_name(name)
|
564
|
+
if name == "" or not name then
|
565
|
+
return ""
|
566
|
+
end
|
567
|
+
return cache_str() .. ':bind:' .. name
|
568
|
+
end
|
569
|
+
local function split_pool_value(value)
|
570
|
+
if not value or value == "" then
|
571
|
+
return "", "", -1
|
572
|
+
end
|
573
|
+
value = tostring(value)
|
574
|
+
local prev, next, expiry = string.match(value, "(.*)||(.*)||(.*)")
|
575
|
+
return prev, next, tonumber(expiry)
|
576
|
+
end
|
577
|
+
local function join_pool_value(prev, next, expiry)
|
578
|
+
if expiry == nil then
|
579
|
+
expiry = -1
|
580
|
+
end
|
581
|
+
return tostring(prev) .. "||" .. tostring(next) .. "||" .. tostring(expiry)
|
582
|
+
end
|
583
|
+
local function timeout_to_expiry(timeout)
|
584
|
+
if timeout == nil or timeout <= 0 then
|
585
|
+
return -1
|
586
|
+
end
|
587
|
+
return os.time() + timeout
|
588
|
+
end
|
589
|
+
local function is_expiry_invalid(expiry)
|
590
|
+
return expiry ~= nil and expiry > 0 and expiry <= os.time()
|
591
|
+
end
|
592
|
+
local function is_expired(value)
|
593
|
+
local _, _, expiry = split_pool_value(value)
|
594
|
+
return is_expiry_invalid(expiry)
|
213
595
|
end
|
214
596
|
local poolItemsKey = pool_str()
|
215
597
|
local headKey = pool_pointer_str(true)
|
216
598
|
local tailKey = pool_pointer_str(false)
|
217
|
-
local function push_to_tail(itemName)
|
218
|
-
local tail
|
599
|
+
local function push_to_tail(itemName, expiry) -- push the item to the free list
|
600
|
+
local tail = redis.call("GET", tailKey)
|
219
601
|
if not tail then
|
220
602
|
tail = ""
|
221
603
|
end
|
222
|
-
redis.call("HSET", poolItemsKey, itemName,
|
223
|
-
if tail == "" then
|
604
|
+
redis.call("HSET", poolItemsKey, itemName, join_pool_value(tail, "", expiry))
|
605
|
+
if tail == "" then -- the free list is empty chain
|
224
606
|
redis.call("SET", headKey, itemName)
|
225
607
|
else
|
226
608
|
local tailVal = redis.call("HGET", poolItemsKey, tail)
|
227
|
-
local prev,
|
228
|
-
|
609
|
+
local prev, next, expiry = split_pool_value(tailVal)
|
610
|
+
assert(next == "", "tail is not the last item in the free list")
|
611
|
+
redis.call("HSET", poolItemsKey, tail, join_pool_value(prev, itemName, expiry))
|
229
612
|
end
|
230
|
-
redis.call("SET", tailKey, itemName)
|
613
|
+
redis.call("SET", tailKey, itemName) -- set the tail point to the new item
|
231
614
|
end
|
232
|
-
local function pop_from_head()
|
615
|
+
local function pop_from_head() -- pop the item from the free list
|
233
616
|
local head = redis.call("GET", headKey)
|
234
|
-
if head
|
235
|
-
return nil
|
617
|
+
if not head or head == "" then -- the free list is empty
|
618
|
+
return nil, -1
|
236
619
|
end
|
237
620
|
local headVal = redis.call("HGET", poolItemsKey, head)
|
238
|
-
|
239
|
-
|
240
|
-
|
241
|
-
|
242
|
-
|
243
|
-
|
621
|
+
assert(headVal ~= nil, "head should not nil")
|
622
|
+
local headPrev, headNext, headExpiry = split_pool_value(headVal)
|
623
|
+
-- Check if the head item has expired or is locked
|
624
|
+
if is_expiry_invalid(headExpiry) then -- the item has expired
|
625
|
+
redis.call("HDEL", poolItemsKey, head)
|
626
|
+
redis.call("SET", headKey, headNext)
|
627
|
+
return pop_from_head()
|
244
628
|
end
|
245
|
-
|
246
|
-
|
629
|
+
if redis.call("EXISTS", key_str(head)) > 0 then -- the item is locked
|
630
|
+
redis.call("HSET", poolItemsKey, head, join_pool_value("#ALLOCATED", "#ALLOCATED", headExpiry))
|
631
|
+
redis.call("SET", headKey, headNext)
|
632
|
+
return pop_from_head()
|
633
|
+
end
|
634
|
+
local prev, next, expiry = split_pool_value(headVal)
|
635
|
+
if next == "" then -- the item is the last in the free list
|
247
636
|
redis.call("SET", headKey, "")
|
248
637
|
redis.call("SET", tailKey, "")
|
249
638
|
else
|
250
639
|
local nextVal = redis.call("HGET", poolItemsKey, next)
|
251
|
-
local
|
252
|
-
redis.call("HSET", poolItemsKey, next,
|
640
|
+
local nextPrev, nextNext, nextExpiry = split_pool_value(nextVal)
|
641
|
+
redis.call("HSET", poolItemsKey, next, join_pool_value("", nextNext, nextExpiry))
|
253
642
|
redis.call("SET", headKey, next)
|
254
643
|
end
|
255
|
-
|
644
|
+
redis.call("HSET", poolItemsKey, head, join_pool_value("#ALLOCATED", "#ALLOCATED", expiry))
|
645
|
+
return head, headExpiry
|
256
646
|
end
|
257
|
-
local function
|
258
|
-
|
259
|
-
|
260
|
-
|
261
|
-
|
262
|
-
|
263
|
-
|
264
|
-
|
265
|
-
redis.call("
|
647
|
+
local function set_item_allocated(itemName, val)
|
648
|
+
if not val then
|
649
|
+
val = redis.call("HGET", poolItemsKey, itemName)
|
650
|
+
end
|
651
|
+
if val then
|
652
|
+
local prev, next, expiry = split_pool_value(val)
|
653
|
+
if prev ~= "#ALLOCATED" then
|
654
|
+
if is_expiry_invalid(expiry) then
|
655
|
+
redis.call("HDEL", poolItemsKey, itemName)
|
266
656
|
end
|
267
|
-
|
268
|
-
|
269
|
-
|
270
|
-
|
271
|
-
|
272
|
-
|
273
|
-
|
274
|
-
redis.call("
|
657
|
+
if prev ~= "" then
|
658
|
+
local prevVal = redis.call("HGET", poolItemsKey, prev)
|
659
|
+
if prevVal then
|
660
|
+
local prevPrev, prevNext, prevExpiry = split_pool_value(prevVal)
|
661
|
+
redis.call("HSET", poolItemsKey, prev, join_pool_value(prevPrev, next, prevExpiry))
|
662
|
+
end
|
663
|
+
else
|
664
|
+
redis.call("SET", headKey, next or "")
|
275
665
|
end
|
276
|
-
|
277
|
-
|
666
|
+
if next ~= "" then
|
667
|
+
local nextVal = redis.call("HGET", poolItemsKey, next)
|
668
|
+
if nextVal then
|
669
|
+
local nextPrev, nextNext, nextExpiry = split_pool_value(nextVal)
|
670
|
+
redis.call("HSET", poolItemsKey, next, join_pool_value(prev, nextNext, nextExpiry))
|
671
|
+
end
|
672
|
+
else
|
673
|
+
redis.call("SET", tailKey, prev or "")
|
674
|
+
end
|
675
|
+
redis.call("HSET", poolItemsKey, itemName, join_pool_value("#ALLOCATED", "#ALLOCATED", expiry))
|
278
676
|
end
|
279
677
|
end
|
280
|
-
redis.call("HDEL", poolItemsKey, itemName)
|
281
678
|
end
|
282
|
-
|
283
|
-
|
284
|
-
|
285
|
-
|
286
|
-
|
287
|
-
|
288
|
-
|
289
|
-
|
290
|
-
|
291
|
-
|
292
|
-
|
293
|
-
if
|
294
|
-
|
679
|
+
local function check_item_health(itemName, value)
|
680
|
+
if not value then
|
681
|
+
value = redis.call("HGET", pool_str(), itemName)
|
682
|
+
end
|
683
|
+
assert(value, "value should not be nil")
|
684
|
+
local prev, next, expiry = split_pool_value(value)
|
685
|
+
if is_expiry_invalid(expiry) then -- Check if the item has expired
|
686
|
+
set_item_allocated(itemName)
|
687
|
+
redis.call("HDEL", poolItemsKey, itemName)
|
688
|
+
end
|
689
|
+
local locked = redis.call("EXISTS", key_str(itemName)) > 0
|
690
|
+
if prev == "#ALLOCATED" then
|
691
|
+
if not locked then
|
692
|
+
push_to_tail(itemName, expiry)
|
693
|
+
end
|
295
694
|
else
|
296
|
-
|
695
|
+
if locked then
|
696
|
+
set_item_allocated(itemName)
|
697
|
+
end
|
297
698
|
end
|
298
699
|
end
|
299
700
|
'''
|
300
701
|
|
301
702
|
@cached_property
|
302
703
|
def _extend_script(self):
|
303
|
-
"""Cached
|
304
|
-
|
704
|
+
"""Cached Lua script to add or update keys in the pool.
|
705
|
+
|
706
|
+
Iterates through provided keys (ARGV[2...]).
|
707
|
+
If a key doesn't exist in the pool hash, it's added to the tail of the free list
|
708
|
+
using push_to_tail() with the specified expiry (calculated from ARGV[1] timeout).
|
709
|
+
If a key *does* exist, its expiry time is updated in the pool hash.
|
710
|
+
"""
|
711
|
+
return self.redis.register_script(f'''
|
712
|
+
{self._lua_required_string}
|
713
|
+
local timeout = tonumber(ARGV[1] or -1)
|
714
|
+
local expiry = timeout_to_expiry(timeout)
|
715
|
+
for i=2, #ARGV do
|
716
|
+
local itemName = ARGV[i]
|
717
|
+
local val = redis.call("HGET", poolItemsKey, itemName)
|
718
|
+
if val then -- only refresh the expiry timeout
|
719
|
+
local prev, next, _ = split_pool_value(val)
|
720
|
+
val = join_pool_value(prev, next, expiry)
|
721
|
+
redis.call("HSET", poolItemsKey, itemName, val)
|
722
|
+
else -- refresh the expiry timeout
|
723
|
+
push_to_tail(itemName, expiry)
|
724
|
+
end
|
725
|
+
end''')
|
305
726
|
|
306
|
-
def extend(self, keys: Optional[Sequence[str]] = None):
|
727
|
+
def extend(self, keys: Optional[Sequence[str]] = None, timeout: int = -1):
|
307
728
|
"""Add new resources to the allocation pool.
|
308
|
-
|
729
|
+
|
309
730
|
Args:
|
310
731
|
keys: Sequence of resource identifiers to add to the pool
|
732
|
+
timeout: Optional timeout in seconds for the pool items (-1 means no timeout)
|
311
733
|
"""
|
312
734
|
if keys is not None and len(keys) > 0:
|
313
|
-
|
314
|
-
|
315
|
-
|
316
|
-
def _shrink_lua_string(self):
|
317
|
-
"""LUA script for removing resources from the allocation pool.
|
318
|
-
|
319
|
-
This script removes specified items from the pool by deleting them
|
320
|
-
from the linked list structure.
|
321
|
-
"""
|
322
|
-
return f'''{self._lua_required_string}
|
323
|
-
for _, itemName in ipairs(ARGV) do
|
324
|
-
delete_item(itemName)
|
325
|
-
end
|
326
|
-
'''
|
735
|
+
# Ensure timeout is integer for Lua script
|
736
|
+
int_timeout = timeout if timeout is not None else -1
|
737
|
+
self._extend_script(args=[int_timeout] + list(keys))
|
327
738
|
|
328
739
|
@cached_property
|
329
740
|
def _shrink_script(self):
|
330
|
-
"""Cached
|
331
|
-
|
741
|
+
"""Cached Lua script to remove keys from the pool.
|
742
|
+
|
743
|
+
Iterates through provided keys (ARGV[1...]).
|
744
|
+
For each key:
|
745
|
+
1. Calls set_item_allocated() to remove it from the free list structure.
|
746
|
+
2. Deletes the key entirely from the pool hash using HDEL.
|
747
|
+
"""
|
748
|
+
return self.redis.register_script(f'''{self._lua_required_string}
|
749
|
+
for i=1, #ARGV do
|
750
|
+
local itemName = ARGV[i]
|
751
|
+
set_item_allocated(itemName)
|
752
|
+
redis.call("HDEL", poolItemsKey, itemName)
|
753
|
+
end''')
|
332
754
|
|
333
755
|
def shrink(self, keys: Optional[Sequence[str]] = None):
|
334
756
|
"""Remove resources from the allocation pool.
|
335
|
-
|
757
|
+
|
336
758
|
Args:
|
337
759
|
keys: Sequence of resource identifiers to remove from the pool
|
338
760
|
"""
|
339
761
|
if keys is not None and len(keys) > 0:
|
340
762
|
self._shrink_script(args=keys)
|
341
763
|
|
342
|
-
@
|
343
|
-
def
|
344
|
-
"""
|
345
|
-
|
346
|
-
|
347
|
-
|
348
|
-
|
349
|
-
|
350
|
-
|
351
|
-
|
352
|
-
|
353
|
-
|
764
|
+
@cached_property
|
765
|
+
def _assign_script(self):
|
766
|
+
"""Cached Lua script to set the pool to exactly the given keys.
|
767
|
+
|
768
|
+
1. Builds a Lua set (`assignSet`) of the desired keys (ARGV[2...]).
|
769
|
+
2. Fetches all current keys from the pool hash (HKEYS).
|
770
|
+
3. Iterates through current keys:
|
771
|
+
- If a key is *not* in `assignSet`, it's removed from the pool
|
772
|
+
(set_item_allocated() then HDEL).
|
773
|
+
- If a key *is* in `assignSet`, it's marked as processed by setting
|
774
|
+
`assignSet[key] = nil`.
|
775
|
+
4. Iterates through the remaining keys in `assignSet` (those not already
|
776
|
+
in the pool). These are added to the tail of the free list using
|
777
|
+
push_to_tail() with the specified expiry (from ARGV[1] timeout).
|
778
|
+
"""
|
779
|
+
return self.redis.register_script(f'''{self._lua_required_string}
|
780
|
+
local timeout = tonumber(ARGV[1] or -1)
|
781
|
+
local expiry = timeout_to_expiry(timeout)
|
782
|
+
local assignSet = {{}}
|
783
|
+
for i=2, #ARGV do
|
784
|
+
local k = ARGV[i]
|
785
|
+
assignSet[k] = true
|
354
786
|
end
|
355
787
|
local allItems = redis.call("HKEYS", poolItemsKey)
|
356
788
|
for _, itemName in ipairs(allItems) do
|
357
|
-
if not
|
358
|
-
|
789
|
+
if not assignSet[itemName] then
|
790
|
+
set_item_allocated(itemName)
|
791
|
+
redis.call("HDEL", poolItemsKey, itemName)
|
359
792
|
else
|
360
|
-
|
793
|
+
assignSet[itemName] = nil
|
361
794
|
end
|
362
795
|
end
|
363
|
-
for k, v in pairs(
|
796
|
+
for k, v in pairs(assignSet) do
|
364
797
|
if v then
|
365
|
-
push_to_tail(k)
|
798
|
+
push_to_tail(k, expiry)
|
366
799
|
end
|
367
800
|
end
|
368
|
-
'''
|
801
|
+
''')
|
369
802
|
|
370
|
-
|
371
|
-
def _assign_lua_script(self):
|
372
|
-
"""Cached Redis script for assigning resources to the allocation pool."""
|
373
|
-
return self.redis.register_script(self._assign_lua_string)
|
374
|
-
|
375
|
-
def assign(self, keys: Optional[Sequence[str]] = None):
|
803
|
+
def assign(self, keys: Optional[Sequence[str]] = None, timeout: int = -1):
|
376
804
|
"""Completely replace the resources in the allocation pool.
|
377
|
-
|
805
|
+
|
378
806
|
Args:
|
379
807
|
keys: Sequence of resource identifiers to assign to the pool,
|
380
808
|
replacing any existing resources
|
809
|
+
timeout: Optional timeout in seconds for the pool items (-1 means no timeout)
|
381
810
|
"""
|
382
811
|
if keys is not None and len(keys) > 0:
|
383
|
-
self.
|
812
|
+
self._assign_script(args=[timeout] + list(keys))
|
384
813
|
else:
|
385
814
|
self.clear()
|
386
815
|
|
387
816
|
def keys(self) -> Iterable[str]:
|
388
817
|
"""Get all resource identifiers in the allocation pool.
|
389
|
-
|
818
|
+
|
390
819
|
Returns:
|
391
820
|
Iterable of resource identifiers in the pool
|
392
821
|
"""
|
@@ -394,10 +823,10 @@ class RedisAllocator(RedisLockPool, Generic[U]):
|
|
394
823
|
|
395
824
|
def __contains__(self, key):
|
396
825
|
"""Check if a resource identifier is in the allocation pool.
|
397
|
-
|
826
|
+
|
398
827
|
Args:
|
399
828
|
key: Resource identifier to check
|
400
|
-
|
829
|
+
|
401
830
|
Returns:
|
402
831
|
True if the resource is in the pool, False otherwise
|
403
832
|
"""
|
@@ -406,7 +835,7 @@ class RedisAllocator(RedisLockPool, Generic[U]):
|
|
406
835
|
@property
|
407
836
|
def _cache_str(self):
|
408
837
|
"""Get the Redis key for the allocator's cache.
|
409
|
-
|
838
|
+
|
410
839
|
Returns:
|
411
840
|
String representation of the Redis key for the cache
|
412
841
|
"""
|
@@ -419,180 +848,287 @@ class RedisAllocator(RedisLockPool, Generic[U]):
|
|
419
848
|
|
420
849
|
def _soft_bind_name(self, name: str) -> str:
|
421
850
|
"""Get the Redis key for a soft binding.
|
422
|
-
|
851
|
+
|
423
852
|
Args:
|
424
853
|
name: Name of the soft binding
|
425
|
-
|
854
|
+
|
426
855
|
Returns:
|
427
856
|
String representation of the Redis key for the soft binding
|
428
857
|
"""
|
429
858
|
return f"{self._cache_str}:bind:{name}"
|
430
859
|
|
431
|
-
def update_soft_bind(self, name: str, key: str):
|
860
|
+
def update_soft_bind(self, name: str, key: str, timeout: Timeout = 3600):
|
432
861
|
"""Update a soft binding between a name and a resource.
|
433
|
-
|
862
|
+
|
863
|
+
Soft bindings create a persistent mapping between named objects and allocated keys,
|
864
|
+
allowing the same key to be consistently allocated to the same named object.
|
865
|
+
This is useful for maintaining affinity between objects and their resources.
|
866
|
+
|
434
867
|
Args:
|
435
868
|
name: Name to bind
|
436
869
|
key: Resource identifier to bind to the name
|
437
870
|
"""
|
438
|
-
self.update(self._soft_bind_name(name), key, timeout=
|
871
|
+
self.update(self._soft_bind_name(name), key, timeout=timeout)
|
439
872
|
|
440
873
|
def unbind_soft_bind(self, name: str):
|
441
874
|
"""Remove a soft binding.
|
442
|
-
|
875
|
+
|
876
|
+
This removes the persistent mapping between a named object and its allocated key,
|
877
|
+
allowing the key to be freely allocated to any requestor.
|
878
|
+
|
443
879
|
Args:
|
444
880
|
name: Name of the soft binding to remove
|
445
881
|
"""
|
446
882
|
self.unlock(self._soft_bind_name(name))
|
447
883
|
|
448
|
-
|
449
|
-
|
450
|
-
|
451
|
-
|
452
|
-
|
453
|
-
of the linked list and marking it as allocated. If the allocator
|
454
|
-
is not shared, the script also sets a lock on the allocated resource.
|
884
|
+
def get_soft_bind(self, name: str) -> Optional[str]:
|
885
|
+
"""Get the resource identifier bound to a name.
|
886
|
+
|
887
|
+
Args:
|
888
|
+
name: Name of the soft binding
|
455
889
|
"""
|
456
|
-
return
|
890
|
+
return self.redis.get(self._soft_bind_name(name))
|
891
|
+
|
892
|
+
@cached_property
|
893
|
+
def _malloc_script(self):
|
894
|
+
"""Cached Lua script to allocate a key from the pool.
|
895
|
+
|
896
|
+
Input ARGS: timeout, name (for soft binding), soft_bind_timeout
|
897
|
+
|
898
|
+
1. Soft Binding Check (if name provided):
|
899
|
+
- Tries to GET the bound key from the soft bind cache key.
|
900
|
+
- If found and the key is *not* currently locked (checked via EXISTS key_str(cachedKey)),
|
901
|
+
it refreshes the soft bind expiry and returns the cached key.
|
902
|
+
- If found but the key *is* locked, it deletes the stale soft bind entry.
|
903
|
+
2. Pop from Head: Calls `pop_from_head()` to get the next available key
|
904
|
+
from the free list head. This function internally skips expired/locked items.
|
905
|
+
3. Lock/Update (if key found):
|
906
|
+
- If `shared=False`: Sets the lock key (`key_str(itemName)`) with the specified timeout.
|
907
|
+
- If `shared=True`: Calls `push_to_tail()` to put the item back onto the free list immediately.
|
908
|
+
4. Update Soft Bind Cache (if key found and name provided):
|
909
|
+
- Sets the soft bind cache key to the allocated `itemName` with its timeout.
|
910
|
+
5. Returns the allocated `itemName` or nil if the pool was empty.
|
911
|
+
"""
|
912
|
+
return self.redis.register_script(f'''
|
457
913
|
{self._lua_required_string}
|
458
914
|
local shared = {1 if self.shared else 0}
|
915
|
+
|
459
916
|
local timeout = tonumber(ARGV[1])
|
460
|
-
local
|
461
|
-
|
462
|
-
|
463
|
-
|
464
|
-
|
465
|
-
|
466
|
-
|
467
|
-
|
468
|
-
|
469
|
-
|
917
|
+
local name_arg = ARGV[2] -- Original name argument
|
918
|
+
local cacheName = soft_bind_name(name_arg) -- Key for soft binding cache
|
919
|
+
local cacheTimeout = tonumber(ARGV[3]) -- Timeout for the soft binding cache entry
|
920
|
+
local function refresh_cache(cacheKey)
|
921
|
+
-- Only refresh if a valid name and timeout were provided
|
922
|
+
if name_arg ~= "" then
|
923
|
+
if cacheTimeout ~= nil and cacheTimeout > 0 then
|
924
|
+
redis.call("SET", cacheName, cacheKey, "EX", cacheTimeout)
|
925
|
+
else -- If timeout is invalid/zero, set without expiry
|
926
|
+
redis.call("SET", cacheName, cacheKey)
|
927
|
+
end
|
928
|
+
end
|
929
|
+
end
|
930
|
+
-- Check soft binding only if a name was provided
|
931
|
+
if name_arg ~= "" then
|
932
|
+
local cachedKey = redis.call("GET", cacheName)
|
933
|
+
if cachedKey then
|
934
|
+
-- Check if the cached key exists and is currently locked (in non-shared mode)
|
935
|
+
if redis.call("HEXISTS", poolItemsKey, cachedKey) <= 0 or redis.call("EXISTS", key_str(cachedKey)) > 0 then
|
936
|
+
-- Cached key is locked, binding is stale, remove it
|
937
|
+
redis.call("DEL", cacheName)
|
938
|
+
else
|
939
|
+
-- Cached key is valid (either not locked or in shared mode)
|
940
|
+
refresh_cache(cachedKey) -- Refresh the cache expiry
|
941
|
+
if shared == 0 then
|
942
|
+
redis.call("SET", key_str(cachedKey), "1", "EX", timeout)
|
943
|
+
set_item_allocated(cachedKey)
|
470
944
|
end
|
945
|
+
return cachedKey
|
471
946
|
end
|
472
947
|
end
|
473
948
|
end
|
474
|
-
return itemName
|
475
|
-
'''
|
476
949
|
|
477
|
-
|
478
|
-
|
479
|
-
|
480
|
-
|
950
|
+
-- No valid soft bind found, proceed with normal allocation
|
951
|
+
local itemName, expiry = pop_from_head()
|
952
|
+
if itemName ~= nil then
|
953
|
+
if shared == 0 then
|
954
|
+
-- Non-shared mode: Acquire lock
|
955
|
+
if timeout ~= nil and timeout > 0 then
|
956
|
+
redis.call("SET", key_str(itemName), "1", "EX", timeout)
|
957
|
+
else
|
958
|
+
redis.call("SET", key_str(itemName), "1") -- Set without expiry if timeout <= 0
|
959
|
+
end
|
960
|
+
else
|
961
|
+
-- Shared mode: Just put it back to the tail
|
962
|
+
push_to_tail(itemName, expiry)
|
963
|
+
end
|
964
|
+
end
|
481
965
|
|
482
|
-
|
966
|
+
-- If allocation was successful and a name was provided, update the soft bind cache
|
967
|
+
if itemName then
|
968
|
+
refresh_cache(itemName)
|
969
|
+
end
|
970
|
+
return itemName
|
971
|
+
''')
|
972
|
+
|
973
|
+
def malloc_key(self, timeout: Timeout = 120, name: Optional[str] = None,
|
974
|
+
cache_timeout: Timeout = 3600) -> Optional[str]:
|
483
975
|
"""Allocate a resource key from the pool.
|
484
|
-
|
976
|
+
|
977
|
+
The behavior depends on the allocator's shared mode:
|
978
|
+
- In non-shared mode (default): Locks the allocated key for exclusive access
|
979
|
+
- In shared mode: Simply removes the key from the free list without locking it
|
980
|
+
|
485
981
|
Args:
|
486
|
-
timeout: How long the allocation should be valid (in seconds)
|
487
|
-
|
982
|
+
timeout: How long the allocation lock should be valid (in seconds).
|
983
|
+
name: Optional name to use for soft binding.
|
984
|
+
cache_timeout: Timeout for the soft binding cache entry (seconds).
|
985
|
+
Defaults to 3600. If <= 0, cache entry persists indefinitely.
|
986
|
+
|
488
987
|
Returns:
|
489
988
|
Resource identifier if allocation was successful, None otherwise
|
490
989
|
"""
|
491
|
-
|
990
|
+
if name is None:
|
991
|
+
name = ""
|
992
|
+
# Convert timeout values to integers for Lua
|
993
|
+
lock_timeout_sec = int(self._to_seconds(timeout))
|
994
|
+
cache_timeout_sec = int(self._to_seconds(cache_timeout))
|
995
|
+
# Convert integers to strings for Lua script arguments
|
996
|
+
return self._malloc_script(args=[
|
997
|
+
lock_timeout_sec,
|
998
|
+
name,
|
999
|
+
cache_timeout_sec,
|
1000
|
+
])
|
492
1001
|
|
493
|
-
def malloc(self, timeout: Timeout = 120, obj: Optional[U] = None, params: Optional[dict] = None
|
1002
|
+
def malloc(self, timeout: Timeout = 120, obj: Optional[U] = None, params: Optional[dict] = None,
|
1003
|
+
cache_timeout: Timeout = 3600) -> Optional[RedisAllocatorObject[U]]:
|
494
1004
|
"""Allocate a resource from the pool and wrap it in a RedisAllocatorObject.
|
495
|
-
|
1005
|
+
|
1006
|
+
If a policy is configured, it will be used to control the allocation behavior.
|
1007
|
+
Otherwise, the basic allocation mechanism will be used.
|
1008
|
+
|
496
1009
|
Args:
|
497
|
-
timeout: How long the allocation should be valid (in seconds)
|
498
|
-
obj: The object to wrap in the RedisAllocatorObject
|
499
|
-
|
500
|
-
|
1010
|
+
timeout: How long the allocation lock should be valid (in seconds)
|
1011
|
+
obj: The object to wrap in the RedisAllocatorObject. If it has a `.name`,
|
1012
|
+
soft binding will be attempted.
|
1013
|
+
params: Additional parameters to associate with the allocated object.
|
1014
|
+
cache_timeout: Timeout for the soft binding cache entry (seconds).
|
1015
|
+
Defaults to 3600. Passed to the policy or `malloc_key`.
|
1016
|
+
|
501
1017
|
Returns:
|
502
1018
|
RedisAllocatorObject wrapping the allocated resource if successful, None otherwise
|
503
1019
|
"""
|
504
|
-
|
505
|
-
|
506
|
-
return
|
507
|
-
|
1020
|
+
if self.policy:
|
1021
|
+
# Pass cache_timeout to the policy's malloc method
|
1022
|
+
return self.policy.malloc(
|
1023
|
+
self, timeout, obj, params,
|
1024
|
+
cache_timeout=cache_timeout
|
1025
|
+
)
|
1026
|
+
# No policy, call malloc_key directly
|
1027
|
+
# Explicitly call obj.name if obj exists
|
1028
|
+
name = obj.name if obj and hasattr(obj, 'name') else None
|
1029
|
+
key = self.malloc_key(timeout, name, cache_timeout=cache_timeout)
|
1030
|
+
return RedisAllocatorObject(
|
1031
|
+
self, key, obj, params
|
1032
|
+
)
|
508
1033
|
|
509
|
-
@
|
510
|
-
def
|
511
|
-
"""
|
512
|
-
|
513
|
-
|
514
|
-
|
1034
|
+
@cached_property
|
1035
|
+
def _free_script(self):
|
1036
|
+
"""Cached Lua script to free allocated keys.
|
1037
|
+
|
1038
|
+
Iterates through provided keys (ARGV[2...]).
|
1039
|
+
For each key:
|
1040
|
+
1. Deletes the corresponding lock key (`key_str(k)`) using DEL.
|
1041
|
+
If the key existed (DEL returns 1), it proceeds.
|
1042
|
+
2. Adds the key back to the tail of the free list using `push_to_tail()`
|
1043
|
+
with the specified expiry (calculated from ARGV[1] timeout).
|
515
1044
|
"""
|
516
|
-
return f'''
|
1045
|
+
return self.redis.register_script(f'''
|
517
1046
|
{self._lua_required_string}
|
518
|
-
|
519
|
-
|
520
|
-
|
1047
|
+
local timeout = tonumber(ARGV[1] or -1)
|
1048
|
+
local expiry = timeout_to_expiry(timeout)
|
1049
|
+
for i=2, #ARGV do
|
1050
|
+
local k = ARGV[i]
|
1051
|
+
local deleted = redis.call('DEL', key_str(k))
|
1052
|
+
if deleted > 0 then -- Only push back to pool if it was actually locked/deleted
|
1053
|
+
push_to_tail(k, expiry)
|
1054
|
+
end
|
521
1055
|
end
|
522
|
-
'''
|
1056
|
+
''')
|
523
1057
|
|
524
|
-
|
525
|
-
def _free_script(self):
|
526
|
-
"""Cached Redis script for freeing allocated resources."""
|
527
|
-
return self.redis.register_script(self._free_lua_script)
|
528
|
-
|
529
|
-
def free_keys(self, *keys: str):
|
1058
|
+
def free_keys(self, *keys: str, timeout: int = -1):
|
530
1059
|
"""Free allocated resources.
|
531
|
-
|
1060
|
+
|
532
1061
|
Args:
|
533
1062
|
*keys: Resource identifiers to free
|
1063
|
+
timeout: Optional timeout in seconds for the pool items (-1 means no timeout)
|
534
1064
|
"""
|
535
|
-
|
1065
|
+
if keys:
|
1066
|
+
self._free_script(args=[timeout] + list(keys))
|
1067
|
+
|
1068
|
+
def free(self, obj: RedisAllocatorObject[U], timeout: int = -1):
|
1069
|
+
"""Free an allocated object.
|
536
1070
|
|
537
|
-
|
538
|
-
|
539
|
-
|
1071
|
+
Args:
|
1072
|
+
obj: The allocated object to free
|
1073
|
+
timeout: Optional timeout in seconds for the pool item (-1 means no timeout)
|
1074
|
+
"""
|
1075
|
+
self.free_keys(obj.key, timeout=timeout)
|
1076
|
+
|
1077
|
+
def _gc_cursor_str(self):
|
1078
|
+
"""Get the Redis key for the garbage collection cursor.
|
1079
|
+
|
1080
|
+
Returns:
|
1081
|
+
String representation of the Redis key for the GC cursor
|
1082
|
+
"""
|
1083
|
+
return f'{self._pool_str()}|gc_cursor'
|
540
1084
|
|
541
1085
|
@cached_property
|
542
1086
|
def _gc_script(self):
|
543
|
-
"""Cached
|
544
|
-
return self.redis.register_script(self._gc_lua_script)
|
1087
|
+
"""Cached Lua script for performing garbage collection.
|
545
1088
|
|
546
|
-
|
547
|
-
|
548
|
-
|
549
|
-
|
550
|
-
|
551
|
-
|
552
|
-
|
553
|
-
|
554
|
-
|
555
|
-
|
1089
|
+
Uses HSCAN to iterate through the pool hash incrementally.
|
1090
|
+
Input ARGS: count (max items to scan per call)
|
1091
|
+
|
1092
|
+
1. Gets the scan cursor from a dedicated key (`_gc_cursor_str()`).
|
1093
|
+
2. Calls HSCAN on the pool hash (`pool_str()`) starting from the cursor,
|
1094
|
+
requesting up to `count` items.
|
1095
|
+
3. Iterates through the key-value pairs returned by HSCAN.
|
1096
|
+
4. For each item, calls `check_item_health()` to reconcile its state
|
1097
|
+
(see `_lua_required_string` documentation).
|
1098
|
+
5. Saves the new cursor returned by HSCAN for the next GC call.
|
556
1099
|
"""
|
557
|
-
return f'''
|
1100
|
+
return self.redis.register_script(f'''
|
558
1101
|
{self._lua_required_string}
|
559
|
-
local
|
560
|
-
local
|
561
|
-
|
562
|
-
|
563
|
-
|
1102
|
+
local cursorKey = '{self._gc_cursor_str()}'
|
1103
|
+
local function get_cursor()
|
1104
|
+
local oldCursor = redis.call("GET", cursorKey)
|
1105
|
+
if not oldCursor or oldCursor == "" then
|
1106
|
+
return "0"
|
1107
|
+
else
|
1108
|
+
return oldCursor
|
1109
|
+
end
|
1110
|
+
end
|
1111
|
+
local function set_cursor(cursor)
|
1112
|
+
redis.call("SET", cursorKey, cursor)
|
564
1113
|
end
|
565
|
-
local
|
1114
|
+
local n = tonumber(ARGV[1])
|
1115
|
+
local scanResult = redis.call("HSCAN", pool_str(), get_cursor(), "COUNT", n)
|
566
1116
|
local newCursor = scanResult[1]
|
567
1117
|
local kvList = scanResult[2]
|
568
|
-
local tail = redis.call("GET", tailKey)
|
569
|
-
if not tail then
|
570
|
-
tail = ""
|
571
|
-
end
|
572
1118
|
for i = 1, #kvList, 2 do
|
573
1119
|
local itemName = kvList[i]
|
574
1120
|
local val = kvList[i + 1]
|
575
|
-
|
576
|
-
local locked = (redis.call("EXISTS", key_str(itemName)) == 1)
|
577
|
-
if not locked then
|
578
|
-
push_to_tail(itemName)
|
579
|
-
end
|
580
|
-
else
|
581
|
-
local locked = (redis.call("EXISTS", key_str(itemName)) == 1)
|
582
|
-
if locked then
|
583
|
-
delete_item(itemName)
|
584
|
-
redis.call("HSET", poolItemsKey, itemName, "#ALLOCATED")
|
585
|
-
end
|
586
|
-
end
|
1121
|
+
check_item_health(itemName, val)
|
587
1122
|
end
|
588
|
-
|
1123
|
+
set_cursor(newCursor)
|
1124
|
+
''')
|
589
1125
|
|
590
1126
|
def gc(self, count: int = 10):
|
591
1127
|
"""Perform garbage collection on the allocation pool.
|
592
|
-
|
1128
|
+
|
593
1129
|
This method scans through the pool and ensures consistency between
|
594
1130
|
the allocation metadata and the actual locks.
|
595
|
-
|
1131
|
+
|
596
1132
|
Args:
|
597
1133
|
count: Number of items to check in this garbage collection pass
|
598
1134
|
"""
|