Add IDLE backpressure detection
This commit is contained in:
Stanislav Hejny
@@ -0,0 +1,231 @@
|
||||
"""
|
||||
methods to create, serialize and deserialize DataMessage objects.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import json
|
||||
import re
|
||||
from io import BytesIO
|
||||
from datetime import datetime, timezone
|
||||
from typing import Dict, List
|
||||
|
||||
from amqp.model.model import DataMessage
|
||||
from amqp.model.snowflake_id import SnowflakeId
|
||||
from amqp.router.serializer import map_with_list_as_string, map_as_string, parse_map_list_string, parse_map_string
|
||||
from amqp.router.utils import sanitize_as_rabbitmq_name
|
||||
from amqp.adapter.file_downloader import download_buffer
|
||||
|
||||
|
||||
class DataMessageFactory:
|
||||
SNOWFLAKE_ID_BITS = 80
|
||||
SEQUENCE_BITS = 12
|
||||
MACHINE_ID_BITS = 24
|
||||
TIMESTAMP_BITS = SNOWFLAKE_ID_BITS - SEQUENCE_BITS - MACHINE_ID_BITS
|
||||
|
||||
MACHINE_ID_MASK = (1 << MACHINE_ID_BITS) - 1
|
||||
SEQUENCE_MASK = (1 << SEQUENCE_BITS) - 1
|
||||
MAGIC_BYTE_HTTP_REQUEST = "A"
|
||||
|
||||
_instance = None
|
||||
snowflake_epoch = datetime(2024, 10, 1, 0, 0, tzinfo=timezone.utc).timestamp()
|
||||
|
||||
def __init__(self, machine_id: int):
|
||||
self.machine_id = (machine_id & self.MACHINE_ID_MASK) << self.SEQUENCE_BITS
|
||||
self.snowflake_sequence = 0
|
||||
self.snowflake_last_timestamp = 0
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls, machine_id: int):
|
||||
"""
|
||||
return reusable factory instance
|
||||
:param machine_id: a generator ID, should be unique for each component to prevent duplicate IDs
|
||||
:return: factory instance
|
||||
"""
|
||||
if cls._instance is None:
|
||||
cls._instance = DataMessageFactory(machine_id)
|
||||
return cls._instance
|
||||
|
||||
def generate_snowflake_id(self) -> SnowflakeId:
|
||||
"""
|
||||
Generate an unique ID following the SnowflakeID structure, see https://en.wikipedia.org/wiki/Snowflake_ID
|
||||
:return: 80-bit wide ID
|
||||
"""
|
||||
timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)
|
||||
|
||||
if timestamp < self.snowflake_last_timestamp:
|
||||
raise RuntimeError("Clock moved backwards. Refusing to generate id")
|
||||
|
||||
if timestamp == self.snowflake_last_timestamp:
|
||||
self.snowflake_sequence = (self.snowflake_sequence + 1) & self.SEQUENCE_MASK
|
||||
if self.snowflake_sequence == 0:
|
||||
while timestamp == self.snowflake_last_timestamp:
|
||||
timestamp = int(datetime.now(timezone.utc).timestamp() * 1000)
|
||||
self.snowflake_last_timestamp = timestamp
|
||||
else:
|
||||
self.snowflake_sequence = 0
|
||||
self.snowflake_last_timestamp = timestamp
|
||||
|
||||
time = (timestamp - 1000 * int(DataMessageFactory.snowflake_epoch))
|
||||
return SnowflakeId(
|
||||
time >> (64 - self.SEQUENCE_BITS - self.MACHINE_ID_BITS),
|
||||
((time << (self.SEQUENCE_BITS + self.MACHINE_ID_BITS)) | self.machine_id | self.snowflake_sequence) & 0xFFFFFFFFFFFFFFFF
|
||||
)
|
||||
|
||||
def create_request_message(
|
||||
self,
|
||||
method: str,
|
||||
domain: str,
|
||||
path: str,
|
||||
headers: Dict[str, List[str]],
|
||||
message,
|
||||
) -> DataMessage:
|
||||
"""
|
||||
Encapsulate / hide the system level inputs, create the AMQ message from business relevant input
|
||||
:param method: a method requested (e.g. HTTP method like GET, POST in case of REST request)
|
||||
:param domain: service host name
|
||||
:param path: path at the service
|
||||
:param headers: HTTP headers
|
||||
:param message: payload
|
||||
:return: DataMessage object
|
||||
"""
|
||||
serializable_headers = headers.copy() if headers else {}
|
||||
payload = base64.b64encode(message.encode('utf-8') if isinstance(message, str) else message)
|
||||
trace_info = {}
|
||||
return DataMessage(
|
||||
self.MAGIC_BYTE_HTTP_REQUEST,
|
||||
self.generate_snowflake_id(),
|
||||
method,
|
||||
domain,
|
||||
path,
|
||||
serializable_headers,
|
||||
trace_info,
|
||||
payload,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def amq_message_routing_key(message: DataMessage) -> str:
|
||||
"""
|
||||
Constructs the message specific routing key.
|
||||
:param message: DataMessage input
|
||||
:return: routing key (as string) compliant to RabbitMQ allowed character set.
|
||||
"""
|
||||
return sanitize_as_rabbitmq_name(f"{message.domain}.{message.path}")
|
||||
|
||||
@staticmethod
|
||||
def get_timestamp_from_id(_id: SnowflakeId) -> int:
|
||||
"""
|
||||
Helper method - retrieve the timestamp from the ID
|
||||
:param _id: SnowflakeID (includes encoded timestamp)
|
||||
:return: the message timestamp
|
||||
"""
|
||||
bits = _id.get_bits()
|
||||
low = bits[0] >> (DataMessageFactory.SEQUENCE_BITS + DataMessageFactory.MACHINE_ID_BITS)
|
||||
high = bits[1] << (64 - DataMessageFactory.SEQUENCE_BITS - DataMessageFactory.MACHINE_ID_BITS)
|
||||
return (high | low) + DataMessageFactory.snowflake_epoch * 1000
|
||||
|
||||
@staticmethod
|
||||
def to_string(message: DataMessage) -> str:
|
||||
"""
|
||||
Human-readable representation.
|
||||
:param message: DataMessage input
|
||||
:return: as string
|
||||
"""
|
||||
return (
|
||||
f"DataMessage{{"
|
||||
f"id={message.id}, "
|
||||
f"method='{message.method}', "
|
||||
f"domain='{message.domain}', "
|
||||
f"path='{message.path}', "
|
||||
f"headers={{{map_with_list_as_string(message.headers)}}}, "
|
||||
f"traceInfo={{{map_as_string(message.trace_info)}}}, "
|
||||
f"base64body={message.base64_body.decode()}"
|
||||
f"}}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def serialize(message: DataMessage) -> bytes:
|
||||
"""
|
||||
Converts DataMessage to byte array
|
||||
:param message: message to serialize
|
||||
:return: messages as bytes
|
||||
"""
|
||||
id_bytes = str(message.id).encode()
|
||||
prolog = (
|
||||
f"|m={message.method}|d={message.domain}|p={message.path}|h={{{map_with_list_as_string(message.headers)}}}|t={{{map_as_string(message.trace_info)}}}|b="
|
||||
).encode('utf-8')
|
||||
header_length = 1 + len(id_bytes) + len(prolog)
|
||||
prolog_len = header_length.to_bytes(2, byteorder='big')
|
||||
msg_length = 2 + header_length + len(message.base64_body)
|
||||
with BytesIO(bytearray(msg_length)) as baos:
|
||||
baos.write(prolog_len)
|
||||
baos.write(DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST.encode('utf-8'))
|
||||
baos.write(id_bytes)
|
||||
baos.write(prolog)
|
||||
baos.write(message.base64_body)
|
||||
return baos.getvalue()
|
||||
|
||||
@staticmethod
|
||||
def from_bytes(input_bytes: bytes) -> DataMessage:
|
||||
"""
|
||||
Builds the DataMessage from its serialized (byte array) form.
|
||||
:param input_bytes: serialized message
|
||||
:return: DataMessage instance
|
||||
"""
|
||||
prolog_len = (input_bytes[0] << 8) + input_bytes[1]
|
||||
metadata = input_bytes[2:prolog_len + 2].decode()
|
||||
|
||||
pattern = re.compile(
|
||||
f"{DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST}"
|
||||
r"([0-9A-Fa-f]+)\|"
|
||||
r"m=([^|]*)\|"
|
||||
r"d=([^|]*)\|"
|
||||
r"p=([^|]*)\|"
|
||||
r"h=(\{[^}]*\})\|"
|
||||
r"t=(\{[^}]*\})\|"
|
||||
r"b="
|
||||
)
|
||||
match = pattern.match(metadata)
|
||||
if not match:
|
||||
raise ValueError(
|
||||
f"DataMessage format does not match expected format. "
|
||||
f"Input string [{metadata}] does not match pattern: {pattern.pattern}"
|
||||
)
|
||||
|
||||
id_str = match.group(1)
|
||||
method = match.group(2)
|
||||
domain = match.group(3)
|
||||
path = match.group(4)
|
||||
headers_str = match.group(5)
|
||||
trace_info_str = match.group(6)
|
||||
body_b64 = input_bytes[prolog_len + 2:]
|
||||
|
||||
headers = parse_map_list_string(headers_str)
|
||||
trace_info = parse_map_string(trace_info_str)
|
||||
|
||||
return DataMessage(
|
||||
DataMessageFactory.MAGIC_BYTE_HTTP_REQUEST,
|
||||
SnowflakeId.from_hex(id_str),
|
||||
method,
|
||||
domain,
|
||||
path,
|
||||
headers,
|
||||
trace_info,
|
||||
body_b64,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def from_stream(stream: bytes, rabbitmq_url: str, loop) -> DataMessage | None:
|
||||
"""
|
||||
Builds the DataMessage from its serialized (byte array) form.
|
||||
:param stream: serialized message
|
||||
:param rabbitmq_url: RabbitMQ connection URL
|
||||
:return: DataMessage instance
|
||||
"""
|
||||
# async implementation
|
||||
message_data = json.loads(stream.decode())
|
||||
req_info = message_data.get('req-info', '')
|
||||
if req_info:
|
||||
(body, eof) = await download_buffer(loop=loop, rabbitmq_url=rabbitmq_url, queue_name=req_info)
|
||||
if eof == 1:
|
||||
return DataMessageFactory.from_bytes(body)
|
||||
return None
|
||||
Reference in New Issue
Block a user