420 lines
16 KiB
Python
420 lines
16 KiB
Python
"""
|
|
|
|
AUTHOR:
|
|
|
|
Khushal P Soonderji
|
|
|
|
DATE:
|
|
|
|
Thursday, 12th Dec., 2024
|
|
|
|
OBJECTIVE:
|
|
|
|
To handle all messages from one place.
|
|
|
|
REFERENCES:
|
|
|
|
N/A
|
|
|
|
DOWNLOADS:
|
|
|
|
N/A
|
|
|
|
"""
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** IMPORT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# To make sibling directories accessible for imports:
|
|
import sys
|
|
sys.path.append(".")
|
|
sys.path.append("..")
|
|
|
|
# My async utils:
|
|
from utils_v2.string import json
|
|
from utils_v2.date_time import date_time
|
|
from utils_v2.database.async_mongo_v2 import AsyncMongo, AsyncMongoStorage
|
|
|
|
# Models:
|
|
from models.servers.core import CoreServerInfoModel, CoreServerCheckLogModel
|
|
|
|
# To make API calls:
|
|
import httpx
|
|
|
|
# To work with MongoDB:
|
|
from bson import ObjectId
|
|
from pymongo import InsertOne, UpdateOne, ReplaceOne
|
|
|
|
# To work with datatypes:
|
|
from typing import Literal, List, Dict, Any
|
|
|
|
# To make deep-copies:
|
|
import copy
|
|
|
|
# To work with base-64 encoding:
|
|
import base64
|
|
|
|
# To work with date and time:
|
|
import datetime
|
|
import time
|
|
|
|
# For asynchronous activities:
|
|
import asyncio
|
|
|
|
# For debugging:
|
|
from icecream import IceCreamDebugger
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MACROS / ONE-TIME INIT ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** VARIABLES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** FUNCTIONS ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
# --- Nothing Yet
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** CLASSES ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
class CoreServerController:
|
|
|
|
# ┏┓┓ ┓┏
|
|
# ┃ ┃┏┓┏┏ ┃┃┏┓┏┓┏
|
|
# ┗┛┗┗┻┛┛ ┗┛┗┻┛ ┛
|
|
|
|
# For MongoDB:
|
|
SERVERS_COLLECTION = "_servers"
|
|
SERVER_CHECK_LOGS_COLLECTION = "_serverCheckLogs"
|
|
|
|
# ┏┓
|
|
# ┃ ┏┓┏┓┏╋┏┓┓┏┏╋┏┓┏┓
|
|
# ┗┛┗┛┛┗┛┗┛ ┗┻┗┗┗┛┛
|
|
|
|
def __init__(
|
|
self,
|
|
batch_size: int = 50,
|
|
batch_timeout: int | float = 60,
|
|
http_client: httpx.AsyncClient = None,
|
|
debug = True,
|
|
debug_prefix = "Srvr. (C) | ",
|
|
debug_only_errors = True
|
|
):
|
|
|
|
"""
|
|
To initialize the instance of this Server controller.
|
|
:param batch_size: When monitoring the servers, how many records will one worker pick at once to poll.
|
|
:param batch_timeout: When a particular server has been picked in a batch, for how long must another worker not
|
|
touch it before making his own attempt.
|
|
:param http_client: An instance of 'httpx' library's AsyncClient. If not given, one will be instantiated
|
|
internally. It is recommended that, for multi-bot use cases, you provide a common HTTP client from outside.
|
|
:param debug: Whether, or not, you would like to show debugging messages on the terminal.
|
|
:param debug_prefix: The prefix string to identify the debugging messages.
|
|
:param debug_only_errors: Whether you would like to show all debugging messages or just error messages.
|
|
"""
|
|
|
|
# Prepare the debugging utility:
|
|
self._debug_prefix = debug_prefix
|
|
self._printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
|
|
if not debug: self._printer.disable()
|
|
self._debug_only_errors = debug_only_errors
|
|
|
|
# Accept the configuration:
|
|
self._batch_size = batch_size
|
|
self._batch_timeout = batch_timeout
|
|
|
|
# Accept/create an HTTP client to work with:
|
|
if http_client: self.__http_client = http_client
|
|
else: self.__http_client = httpx.AsyncClient(
|
|
limits = httpx.Limits(
|
|
max_connections = 100, # ............ Maximum number of connections allowed in the pool.
|
|
max_keepalive_connections = 50, # ... Maximum number of connections that can be kept alive.
|
|
),
|
|
timeout = httpx.Timeout(
|
|
pool = 120.0, # .... Time to wait for a free connection from the pool.
|
|
connect = 2.5, # ... Time to wait for establishing a connection to the server.
|
|
write = 5.0, # ..... Time to wait for sending data.
|
|
read = 5.0 # ....... Time to wait for receiving data.
|
|
)
|
|
)
|
|
|
|
# ┏┓┳┓┳┳┳┓ ┏┓
|
|
# ┃ ┣┫┃┃┃┃ ━━ ┃ ┏┓┏┓┏┓╋┏┓
|
|
# ┗┛┛┗┗┛┻┛ ┗┛┛ ┗ ┗┻┗┗
|
|
|
|
async def register(
|
|
self,
|
|
mongo_data_conn: AsyncMongo,
|
|
server: CoreServerInfoModel
|
|
) -> CoreServerInfoModel | None:
|
|
|
|
"""
|
|
Register one server in the database.
|
|
:param mongo_data_conn: The instance of the database connector to use for the operation.
|
|
:param server: The message to save into the database.
|
|
:return: The object id of the inserted document.
|
|
"""
|
|
|
|
# Note down the current time:
|
|
event_ts = date_time.get_current_utc_date_time(as_string = False)
|
|
|
|
# Register/update the server's info:
|
|
_server = await mongo_data_conn.find_one_and_update(
|
|
collection = self.SERVERS_COLLECTION,
|
|
filter = {
|
|
"hostname": server.hostname,
|
|
"project": server.project,
|
|
"service": server.service,
|
|
"healthCheckUrl": server.healthCheckUrl
|
|
},
|
|
update = {
|
|
"$setOnInsert": {
|
|
"hostname": server.hostname,
|
|
"project": server.project,
|
|
"service": server.service,
|
|
"healthCheckUrl": server.healthCheckUrl,
|
|
"firstRegTs": event_ts
|
|
},
|
|
"$set": {
|
|
"os": server.os,
|
|
"cpu": server.cpu,
|
|
"pid": server.pid,
|
|
"ppid": server.ppid,
|
|
"ipAddr": server.ipAddr,
|
|
"portNo": server.portNo,
|
|
"description": server.description,
|
|
"healthCheckInterval": server.healthCheckInterval,
|
|
"healthAlertUrl": server.healthAlertUrl,
|
|
"online": True,
|
|
"batchId": server.batchId,
|
|
"batchTs": server.batchTs,
|
|
"lastRegTs": server.lastRegTs,
|
|
"lastCheckTs": server.lastCheckTs,
|
|
"checkAfterTs": server.checkAfterTs,
|
|
}
|
|
},
|
|
upsert = True,
|
|
return_updated = True
|
|
)
|
|
|
|
# Done here:
|
|
return CoreServerInfoModel(**_server) if _server else None
|
|
|
|
# ┏┓┳┓┳┳┳┓ ┳┓ •
|
|
# ┃ ┣┫┃┃┃┃ ━━ ┣┫┏┓╋┏┓┓┏┓┓┏┏┓
|
|
# ┗┛┛┗┗┛┻┛ ┛┗┗ ┗┛ ┗┗ ┗┛┗
|
|
|
|
async def get_batch(
|
|
self,
|
|
mongo_data_conn: AsyncMongo,
|
|
) -> List[CoreServerInfoModel] | None:
|
|
|
|
# Get a batch id and note down the time:
|
|
batch_id = ObjectId()
|
|
batch_ts = date_time.get_current_utc_date_time(as_string = False)
|
|
|
|
# Find servers that need to be checked:
|
|
servers = await mongo_data_conn.find_many(
|
|
collection = self.SERVERS_COLLECTION,
|
|
filter = {
|
|
"checkAfterTs": {"$lte": batch_ts},
|
|
"$or": [
|
|
{
|
|
"batchId": {"$eq": None}
|
|
},
|
|
{
|
|
"batchTs": {"$lt": batch_ts - datetime.timedelta(seconds = self._batch_timeout)}
|
|
}
|
|
]
|
|
},
|
|
sort = {"checkAfterTs": 1},
|
|
limit = self._batch_size
|
|
)
|
|
|
|
# If we didn't find any servers:
|
|
if not servers: return servers
|
|
|
|
# Model the servers:
|
|
servers = [CoreServerInfoModel(**s) for s in servers]
|
|
|
|
# Mark these servers as picked for checking:
|
|
updated_count = await mongo_data_conn.update_many(
|
|
collection = self.SERVERS_COLLECTION,
|
|
filter = {"_id": {"$in": [s.serverId for s in servers]}},
|
|
update = {
|
|
"$set": {
|
|
"batchId": batch_id,
|
|
"batchTs": batch_ts
|
|
}
|
|
},
|
|
upsert = False
|
|
)
|
|
|
|
# Done here:
|
|
if updated_count < len(servers):
|
|
self._printer("Mismatched count", len(servers), updated_count)
|
|
return None
|
|
else: return servers
|
|
|
|
# ┏┓┳┓┳┳┳┓ ┳┳ ┓
|
|
# ┃ ┣┫┃┃┃┃ ━━ ┃┃┏┓┏┫┏┓╋┏┓
|
|
# ┗┛┛┗┗┛┻┛ ┗┛┣┛┗┻┗┻┗┗
|
|
# ┛
|
|
|
|
async def check_one_server(
|
|
self,
|
|
mongo_data_conn: AsyncMongo,
|
|
server: CoreServerInfoModel,
|
|
insert_check_log: bool = True,
|
|
release_from_batch: bool = True
|
|
) -> CoreServerCheckLogModel:
|
|
|
|
# Note down some assumptions and starting parameters:
|
|
server_check_log = CoreServerCheckLogModel(serverId = server.serverId)
|
|
start_time = time.time()
|
|
|
|
try:
|
|
|
|
# Make the API call:
|
|
server_response = await self.__http_client.get(url = server.healthCheckUrl)
|
|
server_response.raise_for_status()
|
|
server_check_log.online = True
|
|
server_check_log.message = "The server responded."
|
|
|
|
# Catch various exceptions:
|
|
except httpx.HTTPStatusError as e: server_check_log.message = f"HTTP Error: {e.response.status_code}"
|
|
except httpx.ConnectTimeout as e: server_check_log.message = f"Connect Timeout Error: {e}"
|
|
except httpx.ReadTimeout as e: server_check_log.message = f"Read Timeout Error: {e}"
|
|
except httpx.TimeoutException as e: server_check_log.message = f"Misc. Timeout Error: {e}"
|
|
except httpx.NetworkError as e: server_check_log.message = f"Network Error: {e}"
|
|
except httpx.RequestError as e: server_check_log.message = f"Request Error: {e}"
|
|
|
|
# Note down the latency:
|
|
server_check_log.latency = time.time() - start_time
|
|
|
|
# In case of failure, we try to hit the alert URL:
|
|
if not server_check_log.online:
|
|
try:
|
|
alert_url_response = await self.__http_client.get(url = server.healthAlertUrl)
|
|
server_check_log.alertRaised = True if alert_url_response.status_code == 200 else False
|
|
except:
|
|
server_check_log.alertRaised = False
|
|
|
|
# Store the log in the database:
|
|
if insert_check_log: await mongo_data_conn.insert_one(
|
|
collection = self.SERVER_CHECK_LOGS_COLLECTION,
|
|
document = server_check_log.model_dump()
|
|
)
|
|
|
|
# Release the picked server from the batch:
|
|
if release_from_batch: await mongo_data_conn.update_one(
|
|
collection = self.SERVERS_COLLECTION,
|
|
filter = {"_id": server.serverId},
|
|
update = {
|
|
"$set": {
|
|
"batchId": None,
|
|
"batchTs": None,
|
|
"online": server_check_log.online,
|
|
"lastCheckTs": server_check_log.ts,
|
|
"checkAfterTs": server_check_log.ts + datetime.timedelta(seconds = server.healthCheckInterval)
|
|
}
|
|
},
|
|
upsert = False
|
|
)
|
|
|
|
# Done here:
|
|
return server_check_log
|
|
|
|
async def check_server_batch(
|
|
self,
|
|
mongo_data_conn: AsyncMongo,
|
|
servers: List[CoreServerInfoModel]
|
|
) -> List[CoreServerCheckLogModel]:
|
|
|
|
# If the servers list is a blank array:
|
|
if not servers: return []
|
|
|
|
# Create tasks to check individual servers,
|
|
# tell the tasks not to update the database individually:
|
|
tasks = [
|
|
self.check_one_server(
|
|
mongo_data_conn = mongo_data_conn,
|
|
server = server,
|
|
insert_check_log = False,
|
|
release_from_batch = False
|
|
) for server in servers
|
|
]
|
|
|
|
# Fire the tasks and note down the results:
|
|
server_check_logs = await asyncio.gather(*tasks)
|
|
|
|
# Insert the logs:
|
|
await mongo_data_conn.insert_many(
|
|
collection = self.SERVER_CHECK_LOGS_COLLECTION,
|
|
documents = [log.model_dump() for log in server_check_logs]
|
|
)
|
|
|
|
# Release the batches:
|
|
await mongo_data_conn.bulk_write(
|
|
collection = self.SERVERS_COLLECTION,
|
|
requests = [
|
|
UpdateOne(
|
|
filter = {"_id": server.serverId},
|
|
update = {
|
|
"$set": {
|
|
"batchId": None,
|
|
"batchTs": None,
|
|
"online": log.online,
|
|
"lastCheckTs": log.ts,
|
|
"checkAfterTs": log.ts + datetime.timedelta(seconds = server.healthCheckInterval)
|
|
}
|
|
},
|
|
upsert = False
|
|
) for server, log in zip(servers, server_check_logs)
|
|
]
|
|
)
|
|
|
|
# Done here:
|
|
return server_check_logs
|
|
|
|
|
|
# *****************************************************************************************************************
|
|
# ***** ****
|
|
# *** MAIN PROGRAM ***
|
|
# ***** ****
|
|
# *****************************************************************************************************************
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
pass
|