Add retroactive feature

Author: attdona

src/rembus/__init__.py

@@ -13,6 +13,8 @@
     RembusError,
     CBOR,
     JSON,
+    LastReceived,
+    Now,
     QOS0,
     QOS1,
     QOS2,

src/rembus/admin.py

@@ -1,5 +1,6 @@
 import logging
 import rembus.protocol as rp
+#import rembus.db as rdb
 
 logger = logging.getLogger(__name__)
 
@@ -29,12 +30,14 @@ def remove_exposer(router, twin, topic):
     )
 
 
-def add_subscriber(router, twin, topic):
+def add_subscriber(router, twin, topic, msgfrom):
     logger.debug(
         "[%s] adding [%s] subscriber for topic [%s]", router, twin, topic
     )
     if topic not in router.subscribers:
         router.subscribers[topic] = []
+
+    twin.msg_from[topic] = msgfrom
     upsert_twin(router.subscribers[topic], twin)
 
 def remove_subscriber(router, twin, topic):
@@ -45,6 +48,12 @@ def remove_subscriber(router, twin, topic):
         "[%s] removed [%s] subscriber for topic [%s]", router, twin, topic
     )
 
+async def reactive(router, twin, status: bool):
+    logger.debug("[%s] reactive: %s", twin, status)
+    twin.isreactive = status
+    if status:
+        await router.inbox.put(rp.SendDataAtRest(twin))
+
 async def admin_command(msg: rp.AdminMsg):
     """Handle admin commands"""
     twin = msg.twin
@@ -61,8 +70,9 @@ async def admin_command(msg: rp.AdminMsg):
     elif cmd == rp.REMOVE_IMPL:
         remove_exposer(router, twin, topic)
     elif cmd == rp.ADD_INTEREST:
-        add_subscriber(router, twin, topic)
+        add_subscriber(router, twin, topic, msg.data[rp.MSG_FROM])
     elif cmd == rp.REMOVE_INTEREST:
         remove_subscriber(router, twin, topic)
-
+    elif cmd == rp.REACTIVE_CMD:
+        await reactive(router, twin, msg.data[rp.STATUS])
     await twin.response(rp.STS_OK, msg)

src/rembus/core.py

@@ -340,7 +340,7 @@ def __init__(
         self.config = rs.Config(name)
         self.policy = Policy(policy)
         self.owners = rs.load_tenants(self)
-        self.start_ts = time.time()
+        self.start_ts = rp.timestamp()
         self.msg_cache: list[rp.PubSubMsg] = []
         self.msg_topic_cache: dict[str, List[rp.PubSubMsg]] = {}
         self.tables: dict[str, rdb.Table] = {}
@@ -361,9 +361,10 @@ def __repr__(self):
 
     def isconnected(self, rid: str) -> bool:
         """Check if a component with the given rid is connected."""
-        for tk in self.id_twin:
+        for tk, twin in self.id_twin.items():
             if tk.startswith(rid + "@"):
-                return True
+                if twin.isopen():
+                    return True
         return False
 
     def uptime(self) -> str:
@@ -419,14 +420,16 @@ async def _broadcast(self, msg):
                 if t != twin:
                     # Do not send back to publisher.
                     await t.send(msg)
+                    t.mark = msg.recvts
 
         except Exception as e:  # pylint: disable=broad-exception-caught
             logger.warning("[%s] error in method invocation: %s", self, e)
             traceback.print_exc()
 
     async def _pubsub_msg(self, msg: rp.PubSubMsg):
-        msg.recvts = int(time.time())
         twin = msg.twin
+        ts = rp.timestamp()
+        msg.recvts = ts
         qos = msg.flags & rp.QOS2
         if qos > rp.QOS0 and msg.id:
             if twin.socket:
@@ -438,7 +441,8 @@ async def _pubsub_msg(self, msg: rp.PubSubMsg):
                     return
                 else:
                     # Save the message id to guarantee exactly one delivery.
-                    twin.ackdf[msg.id] = int(time.time())
+                    ##twin.ackdf[msg.id] = int(time.time())
+                    twin.ackdf[msg.id] = rp.timestamp()
 
         if self.db is not None:
             # save the message into msg_cache
@@ -526,7 +530,7 @@ async def _handle_identity(self, msg: rp.IdentityMsg) -> None:
             await self._auth_identity(msg)
 
     async def _handle_attestation(self, msg: rp.AttestationMsg) -> None:
-        sts = self._verify_signature(msg)
+        sts = await self._verify_signature(msg)
         await msg.twin.response(sts, msg)
 
     async def _handle_admin(self, msg: rp.AdminMsg) -> None:
@@ -554,6 +558,8 @@ async def _task_impl(self) -> None:
                 case "save_messages":
                     # save messages to the database periodically
                     rdb.save_data_at_rest(self)
+                case rp.SendDataAtRest():
+                    await rdb.send_data_at_rest(msg)
                 case rp.PubSubMsg():
                     await self._pubsub_msg(msg)
                 case rp.RpcReqMsg():
@@ -587,11 +593,8 @@ async def evaluate(self, twin, topic: str, data: Any) -> Any:
     async def _client_receiver(self, ws):
         """Receive messages from the client component."""
         url = RbURL()
-        if url.twkey in self.id_twin:
-            twin = self.id_twin[url.twkey]
-        else:
-            twin = Twin(url, bottom_router(self), False)
-            self.id_twin[url.twkey] = twin
+        twin = Twin(url, bottom_router(self), False)
+        self.id_twin[url.twkey] = twin
 
         twin.socket = ws
         await twin.twin_receiver()
@@ -627,13 +630,15 @@ def _needs_auth(self, cid: str):
         except FileNotFoundError:
             return False
 
-    def _update_twin(self, twin, identity):
+    async def _update_twin(self, twin, identity):
         logger.debug("[%s] setting name: [%s]", twin, identity)
         self.id_twin.pop(twin.twkey, twin)
         twin.rid = identity
         self.id_twin[twin.twkey] = twin
+        if twin.db is not None:
+            load_mark(twin)
 
-    def _verify_signature(self, msg: rp.AttestationMsg):
+    async def _verify_signature(self, msg: rp.AttestationMsg):
         """Verify the signature of the attestation message."""
         twin = msg.twin
         cid = msg.cid
@@ -654,7 +659,7 @@ def _verify_signature(self, msg: rp.AttestationMsg):
             elif isinstance(pubkey, ec.EllipticCurvePublicKey):
                 pubkey.verify(signature, plain, ec.ECDSA(hashes.SHA256()))
 
-            self._update_twin(twin, msg.cid)
+            await self._update_twin(twin, msg.cid)
             return rp.STS_OK
         except Exception as e:  # pylint: disable=broad-exception-caught
             logger.error("verification failed: %s (%s)", e, type(e))
@@ -680,7 +685,7 @@ async def _auth_identity(self, msg: rp.IdentityMsg):
             # component is provisioned, send the challenge
             response = self._challenge(msg)
         else:
-            self._update_twin(twin, identity)
+            await self._update_twin(twin, identity)
             response = rp.ResMsg(id=msg.id, status=rp.STS_OK)
 
         await twin.send(response)
@@ -764,6 +769,9 @@ def __init__(
         self.reconnect_task: Optional[asyncio.Task[None]] = None
         self.ackdf: dict[int, int] = {}  # msgid => ts
         self.handler["phase"] = lambda: "CLOSED"
+        self.isreactive: bool = False
+        self.msg_from: dict[str, float] = {}
+        self.mark: int = 0
         self.start()
 
     def __str__(self):
@@ -790,6 +798,7 @@ def rid(self):
 
     @rid.setter
     def rid(self, rid: str):
+        self.uid.hasname = True
         self.uid.id = rid
 
     @property
@@ -854,7 +863,10 @@ async def _reconnect(self):
     async def _shutdown(self):
         """Cleanup logic when shutting down the twin."""
         logger.debug("[%s] twin shutdown", self)
-
+        
+        if self.db is not None:
+            save_mark(self)
+        
         if self.isclient or self.uid.isrepl():
             await self._shutdown_router()
 
@@ -1113,7 +1125,7 @@ async def publish(self, topic: str, *data: Any, **kwargs):
         """Publish a message to the specified topic."""
 
         slot = kwargs.get("slot", None)
-        qos = kwargs.get("qos", rp.QOS0)
+        qos = kwargs.get("qos", rp.QOS0) & rp.QOS2
         if qos == rp.QOS0:
             msg = rp.PubSubMsg(topic=topic, data=data, slot=slot)
             msg.twin = self
@@ -1261,7 +1273,7 @@ async def unreactive(self):
     async def subscribe(
         self,
         fn: Callable[..., Any],
-        retroactive: bool = False,
+        msgfrom: float = rp.Now,
         topic: Optional[str] = None,
     ):
         """
@@ -1270,17 +1282,14 @@ async def subscribe(
         if topic is None:
             topic = fn.__name__
 
-        await self.setting(topic, rp.ADD_INTEREST, {"retroactive": retroactive})
+        await self.setting(topic, rp.ADD_INTEREST, {"msg_from": msgfrom})
         self.router.handler[topic] = fn
         return self
 
     async def unsubscribe(self, fn: Callable[..., Any] | str):
         """
         Unsubscribe the function from the corresponding topic.
         """
-        # if topic is None:
-        #    topic = fn.__name__
-
         if isinstance(fn, str):
             topic = fn
         else:
@@ -1306,9 +1315,6 @@ async def unexpose(
         """
         Unexpose the function as a remote procedure call(RPC) handler.
         """
-        # if topic is None:
-        #    topic = fn.__name__
-
         if isinstance(fn, str):
             topic = fn
         else:
@@ -1404,3 +1410,32 @@ async def component(
             await router.init_twin(RbURL(netlink), enc, isserver)
 
     return handle
+
+def load_mark(twin):
+    result =twin.db.sql("""
+        SELECT mark
+        FROM mark
+        WHERE twin = ? AND name = ?
+    """, params=[twin.rid, twin.router.id]).fetchone()
+
+    if result:
+        twin.mark = result[0]
+        logger.debug("[%s] loaded mark %s", twin, twin.mark)
+    else:
+        logger.debug("[%s] no mark found in database", twin)
+
+def save_mark(twin):
+    db = twin.db
+    mark = twin.mark
+    tid = twin.rid
+    name = twin.router.id
+    db.sql("""
+            MERGE INTO mark AS m
+            USING (SELECT ? AS name, ? AS twin, ? AS mark) AS t
+            ON m.name = t.name AND m.twin = t.twin
+            WHEN MATCHED THEN
+                UPDATE SET mark = t.mark
+            WHEN NOT MATCHED THEN
+                INSERT (name, twin, mark) VALUES (t.name, t.twin, t.mark)
+        """, params=[name, tid, mark])
+    logger.debug("[%s] saved mark %s", twin, mark)

src/rembus/db.py

@@ -14,7 +14,7 @@
 import pyarrow as pa
 from pydantic import BaseModel, Field, model_validator
 from rembus.settings import broker_dir, rembus_dir
-from rembus.protocol import tag2df, df2tag, PubSubMsg
+from rembus.protocol import tag2df, df2tag, timestamp, PubSubMsg
 
 logger = logging.getLogger(__name__)
 
@@ -116,7 +116,7 @@ def reset_db(broker_name):
 
         elif dl_url.startswith("sqlite"):
             logger.debug("removing db %s", db)
-            os.remove(db)
+            Path(db).unlink(missing_ok=True)
     else:
         broker_ducklake = Path(rembus_dir()) / f"{broker_name}.ducklake"
         broker_ducklake.unlink(True)
@@ -151,7 +151,7 @@ def init_db(router, schema):
             qos UTINYINT,
             uid UBIGINT,
             topic TEXT NOT NULL,
-            data TEXT
+            data BLOB
         )""",
         """
         CREATE TABLE IF NOT EXISTS exposer (
@@ -608,9 +608,7 @@ def msg_table(router, msg: PubSubMsg):
 
 def save_data_at_rest(router):
     """Save cached messages to the database."""
-    logger.debug(
-        "[save_data_at_rest] saving %d messages", len(router.msg_cache)
-    )
+    #logger.debug("[%s] saving %d messages", router, len(router.msg_cache))
     msgs = router.msg_cache
     if not msgs:
         return
@@ -628,6 +626,27 @@ def save_data_at_rest(router):
     router.msg_cache.clear()
     router.msg_topic_cache.clear()
 
+async def send_messages(twin, df, ts):
+    r = twin.router
+    for (name, recv, slot, qos, uid, topic, data) in df.iter_rows():
+        if recv>=ts and topic in r.subscribers and twin in r.subscribers[topic]: 
+            payload = cbor2.loads(data)
+            if payload:
+                await twin.publish(topic, *payload, slot=slot, qos=qos)
+            else:
+                await twin.publish(topic, slot=slot, qos=qos)
+
+async def send_data_at_rest(msg, max_period=3600000000000):
+    twin = msg.twin
+    r = twin.router
+    db = twin.db
+
+    ts = timestamp() - max_period
+    if twin.uid.hasname:
+        logger.debug("[%s] sending data at rest", twin)
+        df = db.execute(
+            f"SELECT * FROM message WHERE name='{r.id}' AND recv>={ts}").pl()
+        await send_messages(twin, df, ts)
 
 def build_message_batch(broker_id: str, msgs: list):
     """Build a PyArrow Table from a list of message tuples."""

src/rembus/protocol.py

@@ -8,6 +8,7 @@
 import logging
 from typing import Any, List
 import json
+import time
 import cbor2
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import rsa, ec
@@ -30,6 +31,9 @@
 CBOR = 0
 JSON = 1
 
+Now = 0.0
+LastReceived = float("inf")
+
 UInt8 = conint(ge=0, le=255)
 
 QOS0 = UInt8(0x00)
@@ -83,11 +87,16 @@
 
 BROKER_CONFIG = "__config__"
 COMMAND = "cmd"
+MSG_FROM = "msg_from"
 ADD_INTEREST = "subscribe"
 REMOVE_INTEREST = "unsubscribe"
 ADD_IMPL = "expose"
 REMOVE_IMPL = "unexpose"
+REACTIVE_CMD = "reactive"
+STATUS = "status"
 
+def timestamp():
+    return time.time_ns()
 
 def msgid():
     """Return an array of MSGID_SZ random bytes."""
@@ -116,6 +125,9 @@ def bytes2id(byte_data: bytearray) -> int:
     # Convert bytes to integer (assuming little-endian, adjust if big-endian)
     return from_bytes(byte_data)
 
+class SendDataAtRest:
+    def __init__(self, twin):
+        self.twin = twin
 
 class RembusException(Exception):
     """Base class for all Rembus exceptions."""