lease implementation

Author: gabi

.coveragerc

@@ -1,12 +1,12 @@
 [run]
 branch=True
 source=rsocket
-parallel = true
 
 [coverage:report]
 skip_empty = true
 exclude_lines =
     pragma: no cover
+    pass
     @(abc\.)?abstractmethod
 
 [html]

examples/client.py

@@ -30,19 +30,16 @@ def on_subscribe(self, subscription):
 
 
 async def download(reader, writer):
-    socket = RSocketClient(reader, writer)
+    async with RSocketClient(reader, writer) as client:
+        payload = Payload(b'The quick brown fox', b'meta')
 
-    payload = Payload(b'The quick brown fox', b'meta')
+        result = await client.request_response(payload)
+        logging.info('RR: {}'.format(result))
 
-    result = await socket.request_response(payload)
-    logging.info('RR: {}'.format(result))
+        completion_event = Event()
+        client.request_stream(payload).subscribe(StreamSubscriber(completion_event))
 
-    completion_event = Event()
-    socket.request_stream(payload).subscribe(StreamSubscriber(completion_event))
-
-    await completion_event.wait()
-
-    await socket.close()
+        await completion_event.wait()
 
 
 if __name__ == '__main__':

examples/run_against_example_java_server.py

@@ -34,22 +34,21 @@ def on_error(self, exception: Exception):
             completion_event.set()
 
     connection = await asyncio.open_connection('localhost', 6565)
-    client = RSocketClient(*connection,
-                           metadata_encoding=WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA.value.name,
-                           data_encoding=WellKnownMimeTypes.APPLICATION_JSON.value.name)
+    async with RSocketClient(*connection,
+                             metadata_encoding=WellKnownMimeTypes.MESSAGE_RSOCKET_COMPOSITE_METADATA.value.name,
+                             data_encoding=WellKnownMimeTypes.APPLICATION_JSON.value.name) as client:
+        metadata = CompositeMetadata()
+        metadata.append(RoutingMetadata(['investigation.getInvestigationByContext']))
 
-    metadata = CompositeMetadata()
-    metadata.append(RoutingMetadata(['investigation.getInvestigationByContext']))
+        body = bytes(bytearray(map(ord, json.dumps({'active': True}))))
 
-    body = bytes(bytearray(map(ord, json.dumps({'active': True}))))
+        request = Payload(body, metadata.serialize())
 
-    request = Payload(body, metadata.serialize())
+        subscriber = Subscriber()
+        client.request_stream(request).subscribe(subscriber)
+        await completion_event.wait()
 
-    subscriber = Subscriber()
-    client.request_stream(request).subscribe(subscriber)
-    await completion_event.wait()
-
-    assert len(subscriber.values) == 2
+        assert len(subscriber.values) == 2
 
 
 asyncio.run(example())

rsocket/exceptions.py

@@ -15,15 +15,36 @@ class RSocketAuthenticationError(RSocketError):
     pass
 
 
+class RSocketLeaseNotImplementedError(RSocketError):
+    pass
+
+
+class RSocketLeaseNotReceivedTimeoutError(RSocketError):
+    pass
+
+
 class RSocketValueErrorException(RSocketError):
     pass
 
 
+class RSocketConnectionRejected(RSocketError):
+    pass
+
+
 class RSocketProtocolException(RSocketError):
     def __init__(self, error_code: ErrorCode, data: Optional[str] = None):
         self.error_code = error_code
         self.data = data
 
+    def __str__(self) -> str:
+        return 'RSocket error %s(%s): "%s"' % (self.error_code.name, self.error_code.value, self.data or '')
+
+
+class RSocketRejected(RSocketProtocolException):
+    def __init__(self, stream_id: Optional[int] = None):
+        super().__init__(ErrorCode.REJECTED)
+        self.stream_id = stream_id
+
 
 class RSocketFrameFragmentDifferentType(RSocketError):
     pass

rsocket/frame.py

@@ -4,7 +4,7 @@
 from typing import Tuple
 
 from rsocket.error_codes import ErrorCode
-from rsocket.exceptions import RSocketProtocolException
+from rsocket.exceptions import RSocketProtocolException, RSocketRejected
 
 PROTOCOL_MAJOR_VERSION = 1
 PROTOCOL_MINOR_VERSION = 0
@@ -605,6 +605,9 @@ def exception_to_error_frame(stream_id: int, exception: Exception) -> ErrorFrame
 
 
 def error_frame_to_exception(frame: ErrorFrame) -> Exception:
+    if frame.error_code == ErrorCode.REJECTED:
+        return RSocketRejected(frame.stream_id)
+
     if frame.error_code != ErrorCode.APPLICATION_ERROR:
         return RSocketProtocolException(frame.error_code)
 

rsocket/lease.py

@@ -0,0 +1,68 @@
+import abc
+from datetime import timedelta, datetime
+from typing import Optional
+
+from rsocket.exceptions import RSocketRejected
+from rsocket.frame import LeaseFrame
+from rsocket.helpers import to_milliseconds
+
+MAX_31_BIT = pow(2, 31) - 1
+
+
+class Lease(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def assert_request_allowed(self, stream_id: Optional[int] = None):
+        ...
+
+    @abc.abstractmethod
+    def to_frame(self) -> LeaseFrame:
+        ...
+
+
+class NullLease(Lease):
+    def assert_request_allowed(self, stream_id: Optional[int] = None):
+        pass
+
+    def to_frame(self) -> LeaseFrame:
+        frame = LeaseFrame()
+        frame.number_of_requests = MAX_31_BIT
+        frame.time_to_live = MAX_31_BIT
+        return frame
+
+
+class DefinedLease(Lease):
+    __slots__ = (
+        '_maximum_request_count',
+        '_request_counter',
+        '_maximum_lease_time',
+        '_lease_created_at'
+    )
+
+    def __init__(self,
+                 maximum_request_count: int = MAX_31_BIT,
+                 maximum_lease_time: timedelta = timedelta(milliseconds=MAX_31_BIT)):
+        self._maximum_request_count = maximum_request_count
+        self._maximum_lease_time = maximum_lease_time
+        self._lease_created_at = datetime.now()
+        self._request_counter = 0
+
+    def assert_request_allowed(self, stream_id: Optional[int] = None):
+        if not self._is_request_allowed():
+            raise RSocketRejected(stream_id)
+
+    def _is_request_allowed(self) -> bool:
+        if self._lease_created_at + self._maximum_lease_time <= datetime.now():
+            return False
+
+        self._request_counter += 1
+
+        if self._request_counter > self._maximum_request_count:
+            return False
+
+        return True
+
+    def to_frame(self) -> LeaseFrame:
+        frame = LeaseFrame()
+        frame.number_of_requests = self._maximum_request_count
+        frame.time_to_live = to_milliseconds(self._maximum_lease_time)
+        return frame

rsocket/request_handler.py

@@ -6,8 +6,10 @@
 from reactivestreams.publisher import Publisher
 from reactivestreams.subscriber import Subscriber
 from reactivestreams.subscription import DefaultSubscription
+from rsocket.error_codes import ErrorCode
+from rsocket.exceptions import RSocketProtocolException
 from rsocket.extensions.composite_metadata import CompositeMetadata
-from rsocket.frame import LeaseFrame
+from rsocket.lease import Lease
 from rsocket.payload import Payload
 
 
@@ -23,8 +25,9 @@ async def on_setup(self,
                        metadata_encoding: bytes):
         ...
 
-    async def supply_lease(self):
-        """Not implemented by default"""
+    @abstractmethod
+    async def supply_lease(self) -> Lease:
+        ...
 
     @abstractmethod
     async def on_metadata_push(self, metadata: Payload):
@@ -56,13 +59,6 @@ def _parse_composite_metadata(self, metadata: bytes) -> CompositeMetadata:
         composite_metadata.parse(metadata)
         return composite_metadata
 
-    def _send_lease(self, stream: int, time_to_live: int, number_of_requests: int):
-        lease = LeaseFrame()
-        lease.stream_id = stream
-        lease.time_to_live = time_to_live
-        lease.number_of_requests = number_of_requests
-        self.socket.send_frame(lease)
-
 
 class BaseRequestHandler(RequestHandler):
     class UnimplementedPublisher(Publisher, DefaultSubscription):
@@ -76,6 +72,9 @@ async def on_setup(self,
                        metadata_encoding: bytes):
         """Nothing to do on setup by default"""
 
+    async def supply_lease(self) -> Lease:
+        raise RSocketProtocolException(error_code=ErrorCode.UNSUPPORTED_SETUP)
+
     async def request_channel(self, payload: Payload) -> Tuple[Publisher, Subscriber]:
         raise RuntimeError('Not implemented')
 

rsocket/routing/routing_request_handler.py

@@ -10,8 +10,7 @@
 from rsocket.extensions.composite_metadata import CompositeMetadata
 from rsocket.extensions.mimetypes import WellKnownMimeTypes
 from rsocket.extensions.routing import RoutingMetadata
-from rsocket.frame import LeaseFrame
-from rsocket.helpers import to_milliseconds
+from rsocket.lease import DefinedLease
 from rsocket.logger import logger
 from rsocket.payload import Payload
 from rsocket.routing.request_router import RequestRouter
@@ -47,12 +46,11 @@ def __init__(self,
         self._lease_ttl = lease_ttl
         self._lease_max_requests = lease_max_requests
 
-    async def supply_lease(self):
-        if self._lease_ttl is not None and self._lease_max_requests is not None:
-            frame = LeaseFrame()
-            frame.number_of_requests = self._lease_max_requests
-            frame.time_to_live = to_milliseconds(self._lease_ttl)
-            self.socket.send_frame(frame)
+    async def supply_lease(self) -> DefinedLease:
+        return DefinedLease(
+            self._lease_max_requests,
+            self._lease_ttl
+        )
 
     # noinspection PyAttributeOutsideInit
     async def on_setup(self,