Merge branch 'main' into SIENTIAPDE-988-criar-ingestor-opc

This commit is contained in:
vitor-aignosi
2025-06-09 16:45:45 -03:00
committed by GitHub
24 changed files with 1735 additions and 444 deletions

View File

@@ -6,11 +6,17 @@ from kafka.errors import NoBrokersAvailable
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
import traceback
import ingestor.metrics as metrics
import os
class DataManager():
def __init__(self, kafka_servers: str, logger: Logger,
notification_handler: NotificationHandler) -> None:
class DataManager:
def __init__(
self,
kafka_servers: str,
logger: Logger,
notification_handler: NotificationHandler,
) -> None:
"""
Initializes the DataManager instance with a Kafka producer.
This constructor attempts to establish a connection to the specified Kafka servers
@@ -23,47 +29,63 @@ class DataManager():
NoBrokersAvailable: If the connection to Kafka servers fails after 3 attempts.
"""
self.pod_id = os.getenv("HOSTNAME", "localhost")
self.kafka_producer = None
for i in range(0, 3):
logger.info(
f"Trying ({i}) to initializing DataManager with Kafka servers: {kafka_servers}")
f"Trying ({i}) to initializing DataManager with Kafka servers: {kafka_servers}"
)
try:
self.kafka_producer = KafkaProducer(
bootstrap_servers=kafka_servers,
value_serializer=lambda v: json.dumps(v).encode(
'utf-8'), # Serialize JSON messages
key_serializer=lambda k: str(
k).encode('utf-8') if k else None,
"utf-8"
), # Serialize JSON messages
key_serializer=lambda k: str(k).encode("utf-8") if k else None,
)
# Kafka connected
metrics.KAFKA_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(1)
break
except NoBrokersAvailable:
logger.error(
f"Kafka servers {kafka_servers} are not available. Retrying...")
f"Kafka servers {kafka_servers} are not available. Retrying..."
)
sleep(5)
else:
# Kafka not connected
metrics.KAFKA_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(0)
logger.error(
f"Failed to connect to Kafka servers {kafka_servers} after 3 attempts.")
f"Failed to connect to Kafka servers {kafka_servers} after 3 attempts."
)
raise NoBrokersAvailable(
f"Failed to connect to Kafka servers {kafka_servers} after 3 attempts.")
f"Failed to connect to Kafka servers {kafka_servers} after 3 attempts."
)
logger.info(
f"DataManager initialized with Kafka servers: {kafka_servers}")
logger.info(f"DataManager initialized with Kafka servers: {kafka_servers}")
self.logger = logger
self.notification_handler = notification_handler
def __del__(self):
"""Destructor to close the producer connection."""
print("Closing Kafka producer...")
def shutdown(self):
"""Closes the Kafka producer connection."""
if self.kafka_producer:
self.kafka_producer.flush(timeout=10)
self.kafka_producer.close()
try:
self.kafka_producer.flush(timeout=10)
self.kafka_producer.close()
# Mark as disconnected
metrics.KAFKA_CONNECTION_STATUS.labels(pod_id=self.pod_id).set(0)
except Exception as e:
self.logger.error(f"Error closing Kafka producer: {e}")
else:
print("Kafka producer is already closed or not initialized.")
self.logger.warning("Kafka producer is already closed or not initialized.")
def __del__(self):
self.shutdown()
def delivery_report(self, msg: str):
"""Callback for delivery reports from Kafka."""
self.logger.debug(
f"Record successfully produced to {msg.topic} [{msg.partition}] at offset {msg.offset}")
f"Record successfully produced to {msg.topic} [{msg.partition}] at offset {msg.offset}"
)
def delivery_error(self, err: str):
"""Callback for delivery reports from Kafka."""
@@ -86,22 +108,22 @@ class DataManager():
try:
self.logger.debug(
f"Publishing message to topic {topic}: {data}")
self.kafka_producer.send(
topic=topic, value=data).add_callback(
self.delivery_report).add_errback(
self.delivery_error)
self.logger.debug(f"Publishing message to topic {topic}: {data}")
self.kafka_producer.send(topic=topic, value=data).add_callback(
self.delivery_report
).add_errback(self.delivery_error)
self.kafka_producer.flush(timeout=10)
metrics.KAFKA_MESSAGES_SENT.labels(pod_id=self.pod_id, topic=topic).inc()
except Exception as e:
metrics.KAFKA_MESSAGES_ERRORS.labels(pod_id=self.pod_id, topic=topic).inc()
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f"KAFKA_PRODUCER_ERROR_{topic}",
message=f"Error publishing message to topic {topic}: {e}",
block="kafka_producer",
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.logger.error(trace)

View File

@@ -2,11 +2,12 @@ from logging import Logger
import traceback
from typing import Dict, List
from copy import deepcopy
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
from ingestor.managers.data_manager import DataManager
from ingestor.managers.opc_manager import OpcManager
from ingestor.managers.resource_manager import ResourceManager
from sientia_do.notifications.models import Notification, NotificationLevel
from sientia_do.notifications.handlers import NotificationHandler
import ingestor.metrics as metrics
class IngestorManager():
@@ -29,8 +30,10 @@ class IngestorManager():
self.opc_servers = {}
self.notification_handler = notification_handler
self.pod_id = pod_id
def initialize_opc_from_config(self, server_config: dict, data_manager: DataManager, logger: Logger) -> OpcManager | None:
def initialize_opc_from_config(self, server_config: dict,
data_manager: DataManager, logger: Logger) -> OpcManager | None:
"""
Initializes an OPC Manager instance using the provided server configuration.
Args:
@@ -53,9 +56,11 @@ class IngestorManager():
self.logger.info(
f"Initializing OpcManager at {server_config['url']}")
manager = OpcManager(
server_config['name'], server_config['url'], data_manager, logger, server_config['server_uri'],
self.notification_handler, server_config.get('cert_path'), server_config.get(
'private_key_path'), server_config.get('server_cert_path')
server_config['name'], server_config['url'],
data_manager, logger, server_config['server_uri'],
self.notification_handler, self.pod_id, server_config.get('cert_path'),
server_config.get('private_key_path'),
server_config.get('server_cert_path')
)
manager.config = server_config
@@ -76,6 +81,14 @@ class IngestorManager():
return manager
def shutdown(self):
for _server_name, server in self.opc_managers.items():
server.disconnect()
self.data_manager.shutdown()
def __del__(self):
self.shutdown()
def update_opc_servers(self):
"""
Updates the OPC (OLE for Process Control) server connections managed by the ingestor.
@@ -91,7 +104,8 @@ class IngestorManager():
4. Disconnects and removes OPC servers that are no longer registered.
Attributes:
self.managed_tags (dict): A nested dictionary containing slot and server configurations.
self.opc_managers (dict): A dictionary mapping server names to their OPC manager instances.
self.opc_managers (dict): A dictionary mapping server names to their
OPC manager instances.
self.data_manager: An object responsible for managing data operations.
self.logger: A logging object for recording warnings and other messages.
Raises:
@@ -101,27 +115,44 @@ class IngestorManager():
"""
registered_servers = []
for slot, slot_config in self.managed_tags.items():
current_managed_tags = deepcopy(self.managed_tags)
for _slot, slot_config in current_managed_tags.items():
for server, server_config in slot_config.items():
registered_servers.append(server)
server_config = server_config.copy()
server_config = deepcopy(server_config)
server_config.pop('tags', None)
server_instance = None
if server not in self.opc_managers:
server_instance = self.opc_managers.get(server, None)
if server_instance is None:
self.logger.debug(
f"Initializing OPC manager for server {server}"
)
server_instance = self.initialize_opc_from_config(
server_config, self.data_manager, self.logger
)
elif self.opc_managers[server].config != server_config:
self.opc_managers[server].disconnect()
elif server_instance.config != server_config:
self.logger.warning(
f"Reinitializing OPC manager for server {server}"
)
server_instance.disconnect()
del self.opc_managers[server]
server_instance = self.initialize_opc_from_config(
server_config, self.data_manager, self.logger
)
else:
self.logger.debug(
f"OPC manager for server {server} is already initialized and up to date"
)
if server_instance is not None:
self.opc_managers[server] = server_instance
else:
self.logger.warning(
f"Failed to initialize OPC manager for server {server}, "
f"removing server from managed tags."
)
_a = [self.managed_tags[slot].pop(server, None)
for slot, _value in current_managed_tags.items()]
for server in list(self.opc_managers.keys()):
if server not in registered_servers:
@@ -132,6 +163,8 @@ class IngestorManager():
self.opc_managers[server].disconnect()
self.opc_managers.pop(server, None)
metrics.OPC_MANAGERS_ACTIVE.labels(pod_id=self.pod_id).set(len(self.opc_managers))
def check_opc_servers_integrity(self):
"""
Checks the integrity of the OPC servers and updates the OPC servers if necessary.
@@ -143,18 +176,13 @@ class IngestorManager():
if is_lost:
self.logger.warning(
f"OPC server {server} is lost. "
f"Desconnecting from server."
f"Server will be disconnected."
)
opc_manager.disconnect()
self.opc_managers[server] = self.initialize_opc_from_config(
opc_manager.config, self.data_manager, self.logger
)
self.update_opc_servers()
for slot, slot_config in self.managed_tags.items():
if server in slot_config:
self.manage_server(
slot, server, slot_config[server], slot_config[server]['tags']
)
for slot, _config in self.managed_tags.items():
self.managed_tags[slot].pop(server, None)
metrics.OPC_MANAGERS_ACTIVE.labels(pod_id=self.pod_id).set(len(self.opc_managers))
def declare_active(self):
"""
@@ -188,6 +216,7 @@ class IngestorManager():
leases = self.resource_manager.get_all_leases()
self.number_of_slots = len(leases) if leases else 0
metrics.LEASES_TOTAL.set(self.number_of_slots)
return self.number_of_slots
def get_number_of_slots(self) -> int:
@@ -201,6 +230,7 @@ class IngestorManager():
slots = self.resource_manager.get_all_slots()
self.number_of_slots = len(slots) if slots else 0
metrics.SLOTS_TOTAL.set(self.number_of_slots)
return self.number_of_slots
def get_slot_leases(self, max_slots: int = 1) -> Dict:
@@ -212,14 +242,16 @@ class IngestorManager():
Dict: A dictionary where the keys are the slot identifiers (as strings)
and the values are the leased slot details.
Behavior:
- Iterates through available slots and attempts to lease them using the resource manager.
- Iterates through available slots and attempts to lease
them using the resource manager.
- Logs the leasing of each slot.
- Updates the `managed_tags` attribute with the acquired slots.
- Stops leasing once the specified `max_slots` are acquired.
- If unable to acquire the requested number of slots, logs a warning and returns the slots that were leased.
- If unable to acquire the requested number of slots,
logs a warning and returns the slots that were leased.
Notes:
- If a slot is leased but its details cannot be retrieved (i.e., `get_tag_slot` returns None),
that slot is skipped.
- If a slot is leased but its details cannot be retrieved
(i.e., `get_tag_slot` returns None), that slot is skipped.
"""
acquired = {}
@@ -230,9 +262,11 @@ class IngestorManager():
if slots is None:
continue
acquired[str(i)] = slots
metrics.SLOTS_ACQUIRED.labels(pod_id=self.pod_id).inc()
if len(acquired) >= max_slots:
self.managed_tags.update(acquired)
metrics.SLOTS_MANAGED.labels(pod_id=self.pod_id).set(len(self.managed_tags))
return acquired
self.logger.warning(
@@ -240,6 +274,7 @@ class IngestorManager():
f"Only {acquired} slots were leased."
)
self.managed_tags.update(acquired)
metrics.SLOTS_MANAGED.labels(pod_id=self.pod_id).set(len(self.managed_tags))
return acquired
def unsubscribe_slot(self, slot: str):
@@ -280,33 +315,17 @@ class IngestorManager():
removed_slots = []
update = {}
for slot, slot_config in self.managed_tags.items():
for slot, _slot_config in self.managed_tags.items():
self.resource_manager.renew_tag_lease(slot)
update = self.resource_manager.get_tag_slot(slot)
if update is None:
self.logger.warning(
f"Slot {slot} configuration not found. "
f"Removing slot from managed tags."
)
self.unsubscribe_slot(slot)
removed_slots.append(slot)
continue
if update != slot_config:
self.logger.info(
f"Slot {slot} configuration updated. "
f"Old: {slot_config}, New: {update}"
)
self.managed_tags[slot] = update
self.unsubscribe_slot(slot)
self.update_opc_servers()
self.subscribe_to_tags({slot: update})
self.managed_tags[slot] = update
for slot in removed_slots:
del self.managed_tags[slot]
self.update_opc_servers()
self.managed_tags.pop(slot, None)
def drop_slot_leases(self, ids: List[str]) -> None:
"""
@@ -321,8 +340,8 @@ class IngestorManager():
None
"""
for _id in ids:
self.resource_manager.drop_tag_lease(_id)
self.resource_manager.drop_tag_lease(lease_id)
metrics.SLOTS_RELEASED.labels(pod_id=self.pod_id).inc()
def manage_server(self, slot: str, server: str, server_config: dict, tags: dict) -> int:
"""
@@ -381,6 +400,7 @@ class IngestorManager():
tags_to_sub
)
except Exception as e:
metrics.OPC_SUBSCRIPTION_ERRORS.labels(pod_id=self.pod_id, server=server, slot=slot).inc()
trace = traceback.format_exc()
self.notification_handler.build_and_send_notification(
notification_id=f'OPC_SUBSCRIPTION_ERROR_{slot}:{server}',

View File

@@ -7,11 +7,12 @@ from asyncua.sync import Client
from sientia_do.notifications.models import NotificationLevel
from sientia_do.notifications.handlers import NotificationHandler
from ingestor.managers.data_manager import DataManager
import ingestor.metrics as metrics
class OpcManager():
def __init__(self, name: str, url: str, data_manager: DataManager,
logger: Logger, server_uri: str, notification_handler: NotificationHandler,
logger: Logger, server_uri: str, notification_handler: NotificationHandler, pod_id: str,
cert_path: str = None, private_key_path: str = None, server_cert_path: str = None):
self.url = url
self.name = name
@@ -26,8 +27,10 @@ class OpcManager():
self.nodes = {}
self.subscriptions = {}
self.data_manager = data_manager
self.notification_handler = notification_handler
self.pod_id = pod_id
metrics.OPC_CONNECTION_STATUS.labels(pod_id=self.pod_id, server_name=self.name, server_url=self.url).set(0)
metrics.OPC_TAGS_SUBSCRIBED.labels(pod_id=self.pod_id, server_name=self.name).set(0)
def __str__(self):
return f"OpcManager(name={self.name}, url={self.url}, server_uri={self.server_uri})\n" \
@@ -82,11 +85,20 @@ class OpcManager():
Exception: If the connection to the OPC server fails.
"""
self.client = Client(self.url)
if self.cert_path:
self.set_security()
self.logger.info('Starting connection...')
self.client.connect()
metrics.OPC_CONNECTIONS_TOTAL.labels(pod_id=self.pod_id, server_name=self.name).inc()
try:
self.client = Client(self.url)
if self.cert_path:
self.set_security()
self.logger.info(f'Starting connection to {self.name}...')
self.client.connect()
metrics.OPC_CONNECTION_STATUS.labels(pod_id=self.pod_id, server_name=self.name, server_url=self.url).set(1)
self.logger.info(f'Connection to {self.name} successful.')
except Exception as e:
metrics.OPC_CONNECTION_STATUS.labels(pod_id=self.pod_id, server_name=self.name, server_url=self.url).set(0)
metrics.OPC_CONNECTIONS_FAILED.labels(pod_id=self.pod_id, server_name=self.name).inc()
self.logger.error(f"Failed to connect to {self.name}: {e}")
raise
def create_subscription(self, name: str, period: int = 500):
"""
@@ -105,11 +117,14 @@ class OpcManager():
if not self.client:
raise ValueError("Client not connected. Call connect first.")
p = period if period != None else 500
self.subscriptions[name] = self.client.create_subscription(
p, self)
self.logger.info('Subscription created.')
try:
p = period if period is not None else 500
self.subscriptions[name] = self.client.create_subscription(p, self)
self.logger.info(f'Subscription {name} created on {self.name}.')
metrics.OPC_SUBSCRIPTIONS_CREATED.labels(pod_id=self.pod_id, server_name=self.name, slot_name=name).inc()
except Exception as e:
self.logger.error(f"Failed to create subscription {name} on {self.name}: {e}")
raise
def subscribe(self, subscription: str, nodes: dict, collect_period: int):
"""
@@ -129,14 +144,13 @@ class OpcManager():
"""
if not self.subscriptions.get(subscription):
raise ValueError(
"Subscription not created. Call create_subscription first.")
raise ValueError("Subscription not created. Call create_subscription first.")
self.logger.info(f"Subscribing to {subscription}...")
self.logger.info(f"Subscribing to {subscription} on {self.name}...")
self.logger.info(f"Subscribing to nodes: {nodes}")
self.addr_nodes = [self.client.get_node(
n) for n in nodes if n not in self.nodes]
self.addr_nodes = [self.client.get_node(n) for n in nodes if n not in self.nodes]
self.nodes.update(nodes)
metrics.OPC_TAGS_SUBSCRIBED.labels(pod_id=self.pod_id, server_name=self.name).set(len(self.nodes))
self.collect_period = collect_period
for node, config in self.nodes.items():
@@ -188,14 +202,24 @@ class OpcManager():
self.logger.warning("Client already disconnected.")
return
try:
[self.subscriptions[sub].delete() for sub in self.subscriptions]
_a = [self.subscriptions[sub].delete()
for sub in self.subscriptions]
self.logger.warning("Deleted all subscriptions.")
del self.client
self.client = None
self.logger.warning("Disconnected from OPC UA server.")
except Exception as sub_error:
self.logger.error(f"Failed to clean up subscription: {sub_error}")
try:
self.client.disconnect()
except Exception as conn_error:
self.logger.error(
f"Failed to disconnect from OPC UA server: {conn_error}")
finally:
del self.client
self.client = None
metrics.OPC_CONNECTION_STATUS.labels(pod_id=self.pod_id, server_name=self.name, server_url=self.url).set(0)
metrics.OPC_TAGS_SUBSCRIBED.labels(pod_id=self.pod_id, server_name=self.name).set(0)
self.logger.warning("Disconnected from OPC UA server.")
def datachange_notification(self, node, _val, data):
"""
Handles data change notifications for monitored OPC UA nodes.
@@ -223,13 +247,12 @@ class OpcManager():
tag = str(node)
self.logger.debug(
f"Data change notification received for tag: {tag} after {self.nodes[tag]['cycle_rule']['cycle_count']} cycles")
self.logger.debug(
f"Resetting cycle count for tag: {tag} after {self.non_receive_count} OPC cycles")
f"Data change notification received for tag:"
f"{tag} after {self.nodes[tag]['cycle_rule']['cycle_count']} cycles")
self.nodes[tag]['cycle_rule']['cycle_count'] = 0
self.non_receive_count = 0
metrics.OPC_CYCLES_WITHOUT_DATA.labels(pod_id=self.pod_id, server_name=self.name).set(0)
data = {
'tag': tag,
@@ -238,12 +261,13 @@ class OpcManager():
'value': value
}
[self.data_manager.publish(e, data)
for e in self.nodes[tag]['topics']]
_a = [self.data_manager.publish(e, data)
for e in self.nodes[tag]['topics']]
def check_cycles(self):
"""
Checks the cycle counts for all monitored nodes and sends notifications if thresholds are exceeded.
Checks the cycle counts for all monitored nodes and sends
notifications if thresholds are exceeded.
This method iterates through all monitored nodes and updates their cycle counts based on
configured increments. If a node's cycle count exceeds a threshold (5 cycles), it triggers
@@ -251,7 +275,8 @@ class OpcManager():
"""
for node, config in self.nodes.items():
self.nodes[node]['cycle_rule']['cycle_count'] += self.nodes[node]['cycle_rule']['cycle_increment']
self.nodes[node]['cycle_rule']['cycle_count'] += config[
'cycle_rule']['cycle_increment']
if self.nodes[node]['cycle_rule']['cycle_count'] >= 5:
name = config['tag_name']
cycles = self.nodes[node]['cycle_rule']['cycle_count']
@@ -270,14 +295,17 @@ class OpcManager():
"""
self.non_receive_count += 1
metrics.OPC_CYCLES_WITHOUT_DATA.labels(pod_id=self.pod_id, server_name=self.name).set(self.non_receive_count)
if self.non_receive_count >= 5:
self.notification_handler.build_and_send_notification(
notification_id=f'OPC_LISTENNING_STOPPED__{self.name}',
message=f'{self.non_receive_count} cycles without receive from OPC {self.name}. Tags: {json.dumps(self.nodes)}',
message=f'{self.non_receive_count} cycles without '
f'receive from OPC {self.name}. Tags: {json.dumps(self.nodes)}',
block="opc_manager",
level=NotificationLevel.ERROR
)
if self.non_receive_count >= 15:
metrics.OPC_RECONNECTIONS_TOTAL.labels(pod_id=self.pod_id, server_name=self.name).inc()
self.notification_handler.build_and_send_notification(
notification_id=f'OPC_CONNECTION_RETRY__{self.name}',
message=f'Retrying to connect to server {self.name}',

View File

@@ -1,17 +1,59 @@
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, password: str = None) -> None:
self.redis = Redis(host=host, port=port, decode_responses=True,
username=username, password=password)
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
self.pod_id = pod_id
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:
"""
@@ -19,11 +61,11 @@ class ResourceManager:
Args:
key (str): The key to look up in Redis.
Returns:
dict: The value associated with the key, parsed as a dictionary,
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.redis.get(key)
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:
@@ -48,14 +90,19 @@ class ResourceManager:
None
"""
self.redis.set(
f"heartbeat:ingestor:{self.pod_id}", 1, ex=self.heartbeat_ttl)
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
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.
@@ -63,8 +110,14 @@ class ResourceManager:
bool: True if the lease was successfully acquired, False otherwise.
"""
return self.redis.set(
f"lease:opc_tags:{tag_id}", self.pod_id, nx=True, ex=self.lease_ttl)
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:
"""
@@ -78,12 +131,14 @@ class ResourceManager:
bool: True if the lease was successfully renewed, False otherwise.
"""
current = self.redis.get(
f"lease:opc_tags:{tag_id}")
current = self._execute_redis_op(
"get", self.redis.get, f"lease:opc_tags:{tag_id}"
)
if current == self.pod_id:
self.redis.expire(f"lease:opc_tags:{tag_id}", self.lease_ttl)
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:
@@ -96,7 +151,7 @@ class ResourceManager:
None
"""
self.redis.delete(f"lease:opc_tags:{tag_id}")
self._execute_redis_op("delete", self.redis.delete, f"lease:opc_tags:{tag_id}")
def get_all_ingestors(self) -> List[str]:
"""
@@ -107,7 +162,7 @@ class ResourceManager:
list: A list of active ingestors.
"""
return self.redis.keys("heartbeat:ingestor:*")
return self._execute_redis_op("keys", self.redis.keys, "heartbeat:ingestor:*")
def get_all_slots(self) -> List[str]:
"""
@@ -118,7 +173,7 @@ class ResourceManager:
int: The number of slots available.
"""
return self.redis.keys("slot:opc_tags:*")
return self._execute_redis_op("keys", self.redis.keys, "slot:opc_tags:*")
def get_all_leases(self) -> List[str]:
"""
@@ -129,4 +184,4 @@ class ResourceManager:
list: A list of active leases.
"""
return self.redis.keys("lease:opc_tags:*")
return self._execute_redis_op("keys", self.redis.keys, "lease:opc_tags:*")