Repository: mbr/tinyrpc

Issue:
Request ID generator may produce memory leaks
There is a private method `JSONRPCProtocol._get_unique_id()` that is responsible for generating unique IDs for RPC client requests: 
https://github.com/mbr/tinyrpc/blob/2b2857d1947e7019bb01e50badc801194684226b/tinyrpc/protocols/jsonrpc.py#L476

It generates ID's by addition which may lead to a memory leak in case if client stays alive for a long time (month? year?). I believe some kind of sliding-window algorithm should be used there instead.


Discussion:

Patch:
diff --git a/docs/protocols.rst b/docs/protocols.rst
index 3084121..804f6ab 100644
--- a/docs/protocols.rst
+++ b/docs/protocols.rst
@@ -88,6 +88,38 @@ API Reference
     :member-order: bysource
 
 
+ID Generators
+-------------
+
+By default, the :py:class:`~tinyrpc.protocols.jsonrpc.JSONRPCProtocol` 
+and :py:class:`~tinyrpc.protocols.msgpackrpc.MSGPACKRPCProtocol` classes
+generates ids as sequential integers starting at 1.
+If alternative id generation is needed, you may supply your own
+generator.
+
+Example
+-------
+
+The following example shows how to use alternative id generators in a protocol
+that supports them.
+
+.. code-block:: python
+
+    from tinyrpc.protocols.jsonrpc import JSONRPCProtocol
+    
+    def collatz_generator():
+        """A sample generator for demonstration purposes ONLY."""
+        n = 27
+        while True:
+            if n % 2 != 0:
+                n = 3*n + 1
+            else:
+                n = n / 2
+            yield n
+
+    rpc = JSONRPCProtocol(id_generator=collatz_generator())
+
+
 Supported protocols
 -------------------
 
diff --git a/tinyrpc/protocols/__init__.py b/tinyrpc/protocols/__init__.py
index 014acf7..cb796a0 100644
--- a/tinyrpc/protocols/__init__.py
+++ b/tinyrpc/protocols/__init__.py
@@ -5,7 +5,8 @@
 Defines the abstract base classes from which a protocol definition must be constructed.
 """
 from abc import ABC
-from typing import Any, List, Dict, Union, Optional
+from typing import Any, Generator, List, Dict, Union, Optional
+import itertools
 
 from tinyrpc import exc
 
@@ -324,3 +325,16 @@ class RPCBatchProtocol(RPCProtocol, ABC):
         :rtype: :py:class:`RPCBatchRequest`
         """
         raise NotImplementedError()
+
+
+def default_id_generator(start: int = 1) -> Generator[int, None, None]:
+    """Generates sequential integers from `start`.
+
+    e.g. 1, 2, 3, .. 9, 10, 11, ...
+
+    :param start: The first value to start with.`
+    :type start: int
+    :return: A generator that yields a sequence of integers.
+    :rtype: :py:class:`Generator[int, None, None]`
+    """
+    return itertools.count(start)
diff --git a/tinyrpc/protocols/jsonrpc.py b/tinyrpc/protocols/jsonrpc.py
index 5a4795b..8e68054 100644
--- a/tinyrpc/protocols/jsonrpc.py
+++ b/tinyrpc/protocols/jsonrpc.py
@@ -12,12 +12,13 @@ Tinyrpc will detect the presence of jsonext and use it automatically.
 
 import json
 import sys
-from typing import Dict, Any, Union, Optional, List, Tuple, Callable
+from typing import Dict, Any, Union, Optional, List, Tuple, Callable, Generator
 
+from . import default_id_generator
 from .. import (
     RPCBatchProtocol, RPCRequest, RPCResponse, RPCErrorResponse,
     InvalidRequestError, MethodNotFoundError, InvalidReplyError, RPCError,
-    RPCBatchRequest, RPCBatchResponse, InvalidParamsError
+    RPCBatchRequest, RPCBatchResponse, InvalidParamsError,
 )
 
 if 'jsonext' in sys.modules:
@@ -505,13 +506,17 @@ class JSONRPCProtocol(RPCBatchProtocol):
     _ALLOWED_REPLY_KEYS = sorted(['id', 'jsonrpc', 'error', 'result'])
     _ALLOWED_REQUEST_KEYS = sorted(['id', 'jsonrpc', 'method', 'params'])
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(
+            self,
+            id_generator: Generator[object, None, None] = default_id_generator(),
+            *args,
+            **kwargs
+    ) -> None:
         super(JSONRPCProtocol, self).__init__(*args, **kwargs)
-        self._id_counter = 0
+        self._id_generator = id_generator
 
-    def _get_unique_id(self) -> int:
-        self._id_counter += 1
-        return self._id_counter
+    def _get_unique_id(self) -> object:
+        return next(self._id_generator)
 
     def request_factory(self) -> 'JSONRPCRequest':
         """Factory for request objects.
diff --git a/tinyrpc/protocols/msgpackrpc.py b/tinyrpc/protocols/msgpackrpc.py
index 7be837a..29ef60d 100644
--- a/tinyrpc/protocols/msgpackrpc.py
+++ b/tinyrpc/protocols/msgpackrpc.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+from . import default_id_generator
 from .. import (
     RPCError,
     RPCErrorResponse,
@@ -15,7 +16,7 @@ from .. import (
 import msgpack
 import six
 
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, Generator
 
 
 class FixedErrorMessageMixin(object):
@@ -236,13 +237,17 @@ class MSGPACKRPCRequest(RPCRequest):
 class MSGPACKRPCProtocol(RPCProtocol):
     """MSGPACKRPC protocol implementation."""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+            self,
+            id_generator: Generator[object, None, None] = default_id_generator(),
+            *args,
+            **kwargs
+    ) -> None:
         super(MSGPACKRPCProtocol, self).__init__(*args, **kwargs)
-        self._id_counter = 0
+        self._id_generator = id_generator
 
     def _get_unique_id(self):
-        self._id_counter += 1
-        return self._id_counter
+        return next(self._id_generator)
 
     def request_factory(self) -> "MSGPACKRPCRequest":
         """Factory for request objects.


Test patch:
diff --git a/tests/test_jsonrpc.py b/tests/test_jsonrpc.py
index d786b16..2a3fb8b 100644
--- a/tests/test_jsonrpc.py
+++ b/tests/test_jsonrpc.py
@@ -230,7 +230,8 @@ def test_request_generation(prot):
 
 def test_jsonrpc_spec_v2_example1(prot):
     # reset id counter
-    prot._id_counter = 0
+    from tinyrpc.protocols import default_id_generator
+    prot._id_generator = default_id_generator(1)
 
     request = prot.create_request('subtract', [42, 23])
 
@@ -264,7 +265,8 @@ def test_jsonrpc_spec_v2_example1(prot):
 
 def test_jsonrpc_spec_v2_example2(prot):
     # reset id counter
-    prot._id_counter = 2
+    from tinyrpc.protocols import default_id_generator
+    prot._id_generator = default_id_generator(3)
 
     request = prot.create_request('subtract',
                                   kwargs={'subtrahend': 23, 'minuend': 42})
diff --git a/tests/test_msgpackrpc.py b/tests/test_msgpackrpc.py
index 6186bf0..71a0f24 100644
--- a/tests/test_msgpackrpc.py
+++ b/tests/test_msgpackrpc.py
@@ -210,7 +210,8 @@ def test_request_generation(prot):
 
 def test_jsonrpc_spec_v2_example1(prot):
     # reset id counter
-    prot._id_counter = 0
+    from tinyrpc.protocols import default_id_generator
+    prot._id_generator = default_id_generator(1)
 
     request = prot.create_request("subtract", [42, 23])
 
diff --git a/tests/test_protocols.py b/tests/test_protocols.py
index 1dedc01..f5343e5 100644
--- a/tests/test_protocols.py
+++ b/tests/test_protocols.py
@@ -58,3 +58,10 @@ def test_parses_error_response(protocol):
     parsed = protocol.parse_reply(err_rep.serialize())
 
     assert hasattr(parsed, 'error')
+
+def test_default_id_generator():
+    from tinyrpc.protocols import default_id_generator
+    g = default_id_generator(1)
+    assert next(g) == 1
+    assert next(g) == 2
+    assert next(g) == 3