188 lines
6.6 KiB
Python
188 lines
6.6 KiB
Python
import json
|
|
from typing import List
|
|
from redis import Redis
|
|
from time import time
|
|
import ingestor.metrics as metrics
|
|
|
|
|
|
class ResourceManager:
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
port: int,
|
|
lease_ttl: int,
|
|
heartbeat_ttl: int,
|
|
pod_id: str,
|
|
username: str | None = None,
|
|
password: str | None = None,
|
|
) -> None:
|
|
self.pod_id = pod_id
|
|
try:
|
|
self.redis = Redis(
|
|
host=host,
|
|
port=port,
|
|
decode_responses=True,
|
|
username=username,
|
|
password=password,
|
|
)
|
|
self.redis.ping()
|
|
metrics.REDIS_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(1)
|
|
except Exception as e:
|
|
print(f"Failed to connect to Redis: {e}")
|
|
metrics.REDIS_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(0)
|
|
raise
|
|
|
|
self.lease_ttl = lease_ttl
|
|
self.heartbeat_ttl = heartbeat_ttl
|
|
|
|
def _execute_redis_op(self, operation_name: str, func, *args, **kwargs):
|
|
"""Wrapper to execute Redis operations and record metrics."""
|
|
start_time = time()
|
|
try:
|
|
result = func(*args, **kwargs)
|
|
metrics.REDIS_OPERATIONS_TOTAL.labels(
|
|
pod_id=self.pod_id, operation=operation_name
|
|
).inc()
|
|
duration = time() - start_time
|
|
metrics.REDIS_OPERATIONS_DURATION.labels(
|
|
pod_id=self.pod_id, operation=operation_name
|
|
).observe(duration)
|
|
return result
|
|
except Exception as e:
|
|
metrics.REDIS_OPERATIONS_ERRORS.labels(
|
|
pod_id=self.pod_id, operation=operation_name
|
|
).inc()
|
|
print(f"Error in Redis operation '{operation_name}': {e}")
|
|
raise
|
|
|
|
def get(self, key: str) -> dict:
|
|
"""
|
|
Retrieve a value from Redis by its key and return it as a dictionary.
|
|
Args:
|
|
key (str): The key to look up in Redis.
|
|
Returns:
|
|
dict: The value associated with the key, parsed as a dictionary,
|
|
or None if the key does not exist or the value is empty.
|
|
"""
|
|
|
|
history = self._execute_redis_op("get", self.redis.get, key)
|
|
return json.loads(history) if history else None
|
|
|
|
def get_tag_slot(self, id: str) -> dict:
|
|
"""
|
|
Retrieve the tag slot information for a given ID.
|
|
Args:
|
|
id (str): The unique identifier of the tag slot to retrieve.
|
|
Returns:
|
|
dict: A dictionary containing the tag slot information associated with the given ID.
|
|
"""
|
|
|
|
return self.get(f"slot:opc_tags:{id}")
|
|
|
|
def ingestor_heartbeat(self) -> None:
|
|
"""
|
|
Sends a heartbeat signal to Redis to indicate that the ingestor is active.
|
|
This method sets a key in Redis with a specific format that includes the
|
|
ingestor's pod ID. The key is set with a value of 1 and an expiration
|
|
time defined by `self.heartbeat_ttl`. This allows monitoring systems to
|
|
track the activity and health of the ingestor.
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
self._execute_redis_op(
|
|
"set",
|
|
self.redis.set,
|
|
f"heartbeat:ingestor:{self.pod_id}",
|
|
1,
|
|
ex=self.heartbeat_ttl,
|
|
)
|
|
|
|
def lease_tag(self, tag_id: str) -> bool:
|
|
"""
|
|
Attempts to lease a tag by setting a key in Redis with a specified TTL (time-to-live).
|
|
This method uses the Redis `SET` command with the `NX` option to ensure that the key
|
|
is only set if it does not already exist. The key is set with an expiration time
|
|
defined by `lease_ttl`.
|
|
Args:
|
|
tag_id (str): The unique identifier of the tag to be leased.
|
|
Returns:
|
|
bool: True if the lease was successfully acquired, False otherwise.
|
|
"""
|
|
|
|
return self._execute_redis_op(
|
|
"set_nx",
|
|
self.redis.set,
|
|
f"lease:opc_tags:{tag_id}",
|
|
self.pod_id,
|
|
nx=True,
|
|
ex=self.lease_ttl,
|
|
)
|
|
|
|
def renew_tag_lease(self, tag_id: str) -> bool:
|
|
"""
|
|
Renews the lease for a specific OPC tag if the current pod holds the lease.
|
|
This method checks if the current pod (identified by `self.pod_id`) holds the lease
|
|
for the given OPC tag. If so, it extends the lease by resetting its expiration time
|
|
in Redis to the configured lease TTL (`self.lease_ttl`).
|
|
Args:
|
|
tag_id (str): The identifier of the OPC tag whose lease is to be renewed.
|
|
Returns:
|
|
bool: True if the lease was successfully renewed, False otherwise.
|
|
"""
|
|
|
|
current = self._execute_redis_op(
|
|
"get", self.redis.get, f"lease:opc_tags:{tag_id}"
|
|
)
|
|
if current == self.pod_id:
|
|
self._execute_redis_op(
|
|
"expire", self.redis.expire, f"lease:opc_tags:{tag_id}", self.lease_ttl
|
|
)
|
|
return True
|
|
return False
|
|
|
|
def drop_tag_lease(self, tag_id: str) -> None:
|
|
"""
|
|
Drops the lease for a specific OPC tag.
|
|
This method removes the lease for the given OPC tag by deleting the corresponding key in Redis.
|
|
Args:
|
|
tag_id (str): The identifier of the OPC tag whose lease is to be dropped.
|
|
Returns:
|
|
None
|
|
"""
|
|
|
|
self._execute_redis_op("delete", self.redis.delete, f"lease:opc_tags:{tag_id}")
|
|
|
|
def get_all_ingestors(self) -> List[str]:
|
|
"""
|
|
Retrieves all active ingestors from Redis.
|
|
This method fetches all keys in Redis that match the pattern for ingestor leases
|
|
and returns a list of active ingestors.
|
|
Returns:
|
|
list: A list of active ingestors.
|
|
"""
|
|
|
|
return self._execute_redis_op("keys", self.redis.keys, "heartbeat:ingestor:*")
|
|
|
|
def get_all_slots(self) -> List[str]:
|
|
"""
|
|
Retrieves the number of slots available in Redis.
|
|
This method counts the number of keys in Redis that match the pattern for OPC tag leases
|
|
and returns the count.
|
|
Returns:
|
|
int: The number of slots available.
|
|
"""
|
|
|
|
return self._execute_redis_op("keys", self.redis.keys, "slot:opc_tags:*")
|
|
|
|
def get_all_leases(self) -> List[str]:
|
|
"""
|
|
Retrieves all active leases from Redis.
|
|
This method fetches all keys in Redis that match the pattern for OPC tag leases
|
|
and returns a list of active leases.
|
|
Returns:
|
|
list: A list of active leases.
|
|
"""
|
|
|
|
return self._execute_redis_op("keys", self.redis.keys, "lease:opc_tags:*")
|