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.
@@ -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
- _allocator: 'RedisAllocator' # Reference to the allocator that created this object
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._allocator = allocator
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._allocator.update(self.key, timeout=timeout)
157
+ self.allocator.update(self.key, timeout=timeout)
123
158
  else:
124
- self._allocator.free(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 memory allocation system.
138
-
139
- This class implements a memory allocation interface using Redis as the backend store.
140
- It manages a pool of resources that can be allocated, freed, and garbage collected.
141
- The implementation uses a doubly-linked list structure stored in Redis hashes to
142
- track available and allocated resources.
143
-
144
- Generic type U must implement the RedisContextableObject protocol, allowing
145
- allocated objects to be used as context managers.
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
- """Initialize a RedisAllocator instance.
151
-
493
+ shared=False, policy: Optional[RedisAllocatorPolicy] = None):
494
+ """Initializes the RedisAllocator.
495
+
152
496
  Args:
153
- redis: Redis client for communication with Redis server
154
- prefix: Prefix for all Redis keys used by this allocator
155
- suffix: Suffix for Redis keys to uniquely identify this allocator
156
- eps: Small value for floating point comparisons
157
- shared: Whether resources can be shared across multiple consumers
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.soft_bind_timeout = 3600 # Default timeout for soft bindings (1 hour)
162
- self.objects: weakref.WeakValueDictionary[str, RedisAllocatorObject] = weakref.WeakValueDictionary()
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
- """LUA script containing helper functions for Redis operations.
193
-
194
- This script defines various LUA functions for manipulating the doubly-linked list
195
- structure that represents the allocation pool:
196
- - pool_pointer_str: Get Redis keys for head/tail pointers
197
- - gc_cursor_str: Get Redis key for GC cursor
198
- - push_to_tail: Add an item to the tail of the linked list
199
- - pop_from_head: Remove and return the item at the head of the linked list
200
- - delete_item: Remove an item from anywhere in the linked list
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 gc_cursor_str()
212
- return '{self._pool_str()}|gc_cursor'
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 = redis.call("GET", tailKey)
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, tostring(tail) .. "||")
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, _ = string.match(tailVal, "(.*)||(.*)")
228
- redis.call("HSET", poolItemsKey, tail, tostring(prev) .. "||" .. tostring(itemName))
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 == nil or head == "" then
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
- if headVal == nil then
239
- redis.call("SET", headKey, "")
240
- redis.call("SET", tailKey, "")
241
- return nil
242
- else
243
- assert(headVal ~= "#ALLOCATED", "head is allocated")
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
- local _, next = string.match(headVal, "(.*)||(.*)")
246
- if next == "" then
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 prev, _ = string.match(nextVal, "(.*)||(.*)")
252
- redis.call("HSET", poolItemsKey, next, tostring("") .. "||" .. tostring(prev))
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
- return head
644
+ redis.call("HSET", poolItemsKey, head, join_pool_value("#ALLOCATED", "#ALLOCATED", expiry))
645
+ return head, headExpiry
256
646
  end
257
- local function delete_item(itemName)
258
- local val = redis.call("HGET", poolItemsKey, itemName)
259
- if val ~= '#ALLOCATED' then
260
- local prev, next = string.match(val, "(.*)||(.*)")
261
- if prev ~= "" then
262
- local prevVal = redis.call("HGET", poolItemsKey, prev)
263
- if prevVal then
264
- local p, _ = string.match(prevVal, "(.*)||(.*)")
265
- redis.call("HSET", poolItemsKey, prev, tostring(p) .. "||" .. tostring(next))
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
- else
268
- redis.call("SET", headKey, next or "")
269
- end
270
- if next ~= "" then
271
- local nxtVal = redis.call("HGET", poolItemsKey, next)
272
- if nxtVal then
273
- local _, n = string.match(nxtVal, "(.*)||(.*)")
274
- redis.call("HSET", poolItemsKey, next, tostring(prev) .. "||" .. tostring(n))
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
- else
277
- redis.call("SET", tailKey, prev or "")
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
- @property
285
- def _extend_lua_string(self):
286
- """LUA script for extending the allocation pool with new resources.
287
-
288
- This script adds new items to the pool if they don't already exist.
289
- New items are added to the tail of the linked list.
290
- """
291
- return f'''{self._lua_required_string}
292
- for _, itemName in ipairs(ARGV) do
293
- if redis.call("HEXISTS", poolItemsKey, itemName) == 1 then
294
- -- todo: gc check
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
- push_to_tail(itemName)
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 Redis script for extending the allocation pool."""
304
- return self.redis.register_script(self._extend_lua_string)
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
- self._extend_script(args=keys)
314
-
315
- @property
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 Redis script for shrinking the allocation pool."""
331
- return self.redis.register_script(self._shrink_lua_string)
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
- @property
343
- def _assign_lua_string(self):
344
- """LUA script for completely replacing the resources in the allocation pool.
345
-
346
- This script clears the existing pool and replaces it with a new set of resources.
347
- Items not in the new set are removed, and items in the new set but not in the
348
- existing pool are added to the tail of the linked list.
349
- """
350
- return f'''{self._lua_required_string}
351
- local wantSet = {{}}
352
- for _, k in ipairs(ARGV) do
353
- wantSet[k] = true
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 wantSet[itemName] then
358
- delete_item(itemName)
789
+ if not assignSet[itemName] then
790
+ set_item_allocated(itemName)
791
+ redis.call("HDEL", poolItemsKey, itemName)
359
792
  else
360
- wantSet[itemName] = nil
793
+ assignSet[itemName] = nil
361
794
  end
362
795
  end
363
- for k, v in pairs(wantSet) do
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
- @cached_property
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._assign_lua_script(args=keys)
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=self.soft_bind_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
- @property
449
- def _malloc_lua_script(self):
450
- """LUA script for allocating a resource from the pool.
451
-
452
- This script allocates a resource by popping an item from the head
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 f'''
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 itemName = pop_from_head()
461
- if itemName ~= nil then
462
- if redis.call("EXISTS", key_str(itemName)) == 1 then
463
- itemName = nil
464
- else
465
- if not shared then
466
- if timeout ~= nil then
467
- redis.call("SET", key_str(itemName), "1", "EX", timeout)
468
- else
469
- redis.call("SET", key_str(itemName), "1")
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
- @cached_property
478
- def _malloc_script(self):
479
- """Cached Redis script for allocating a resource."""
480
- return self.redis.register_script(self._malloc_lua_script)
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
- def malloc_key(self, timeout: Timeout = 120) -> Optional[str]:
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
- return self._malloc_script(keys=[self._to_seconds(timeout)])
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) -> Optional[RedisAllocatorObject[U]]:
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
- params: Additional parameters to associate with the allocated object
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
- key = self.malloc_key(timeout)
505
- if key is None:
506
- return None
507
- return RedisAllocatorObject(self, key, obj, params)
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
- @property
510
- def _free_lua_script(self):
511
- """LUA script for freeing allocated resources.
512
-
513
- This script frees allocated resources by removing their locks
514
- and pushing them back to the tail of the linked list.
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
- for _, k in ipairs(KEYS) do
519
- redis.call('DEL', key_str(k))
520
- push_to_tail(k)
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
- @cached_property
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
- self._free_script(keys, args=keys)
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
- def free(self, obj: RedisAllocatorObject[U]):
538
- """Free an allocated object."""
539
- self.free_keys(obj.key)
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 Redis script for garbage collection."""
544
- return self.redis.register_script(self._gc_lua_script)
1087
+ """Cached Lua script for performing garbage collection.
545
1088
 
546
- @property
547
- def _gc_lua_script(self):
548
- """LUA script for garbage collection of the allocation pool.
549
-
550
- This script scans through the pool and performs two types of cleanup:
551
- 1. Resources marked as allocated but not actually locked are pushed back
552
- to the available pool
553
- 2. Resources not marked as allocated but actually locked are marked as allocated
554
-
555
- This ensures consistency between the allocation metadata and the actual locks.
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 n = tonumber(ARGV[1])
560
- local cursorKey = gc_cursor_str()
561
- local oldCursor = redis.call("GET", cursorKey)
562
- if not oldCursor or oldCursor == "" then
563
- oldCursor = "0"
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 scanResult = redis.call("HSCAN", poolItemsKey, oldCursor, "COUNT", n)
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
- if val == "#ALLOCATED" then
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
  """