Merge pull request #8 from labteral/develop

v2

Author: brunneis

docker/Dockerfile

@@ -28,7 +28,7 @@ RUN \
     pip install --upgrade pip \
     && pip install \
     "cython==0.29.20" \
-    "easyrocks==0.0.16a0" \
+    "easyrocks>=2.214.0" \
     "msgpack==1.0.2" \
     "python-snappy==0.5.4" \
     "pyyaml==5.3.1" \

stopover/broker.py

@@ -98,8 +98,8 @@ def put_message(self, params: dict) -> dict:
 
         partition_numbers = self._get_stream_partition_numbers(stream)
         if partition_number is None:
-            partition_number = utils.get_partition_number(
-                partition_numbers, key)
+            partition_number = \
+                utils.get_partition_number(partition_numbers, key)
         elif partition_number not in partition_numbers:
             raise ValueError('partition does not exist')
 
@@ -276,7 +276,6 @@ def _get_stream_partition_numbers(self, stream: str):
                         raise FileNotFoundError(
                             f'missing partitions among {partition_numbers}')
 
-                    # Initialize new partitions
                     Partition(stream=stream,
                               number=partition_number,
                               data_dir=self.config['global']['data_dir'],
@@ -313,8 +312,8 @@ def _rebalance(self):
                             self.last_seen_by_group[receiver_group][receiver]
                         ) / 1000
 
-                        if receiver_unseen_time < self.config['global'][
-                                'receiver_timeout']:
+                        if receiver_unseen_time \
+                        < self.config['global']['receiver_timeout']:
                             stream_receiver_group_receivers.append(receiver)
 
                         else:
@@ -353,9 +352,9 @@ def _rebalance(self):
                                     stream_partition_numbers[index])
 
             for stream, receiver_group, receiver in receivers_to_remove:
-                logging.info(
-                    f'receiver "{receiver}" kicked from the receiver_group \
-                        "{receiver_group}" for the stream "{stream}"')
+                logging.info(f'receiver "{receiver}" kicked from the ' \
+                    f'receiver_group "{receiver_group}" ' \
+                    f'for the stream "{stream}"')
                 del self.partitions_by_group[stream][receiver_group][receiver]
                 if receiver in self.last_seen_by_group[receiver_group]:
                     del self.last_seen_by_group[receiver_group][receiver]

stopover/partition.py

@@ -3,10 +3,14 @@
 import utils
 from os import makedirs
 from easyrocks import RocksDB, WriteBatch, CompressionType
+from easyrocks.utils import int_to_padded_bytes
 from threading import Lock
 import logging
 from typing import Dict
 
+UINT_BYTES = 8
+MAX_UINT = 2**(UINT_BYTES * 8) - 1
+
 
 class PartitionItem:
     def __init__(self,
@@ -39,6 +43,10 @@ def _load_from_dict(self, value: Dict):
 
 
 class Partition:
+    MESSAGE = b'\x00'
+    INDEX = b'\x01'
+    OFFSET = b'\x02'
+
     def __init__(self,
                  stream: str,
                  number: int,
@@ -114,7 +122,7 @@ def set_offset(self, receiver: str, offset: int):
             offset_key = self._get_offset_key(receiver)
             self._store.put(offset_key, offset)
 
-    def prune(self, ttl):
+    def prune(self, ttl: int):
         ttl *= 1000  # milliseconds
 
         current_timestamp = utils.get_timestamp_ms()
@@ -149,43 +157,40 @@ def _get_by_index(self, index: int) -> bytes:
         partition_item = PartitionItem(item_dict=value)
         return partition_item
 
-    def _get_index(self):
-        index_key = self._get_index_key()
+    def _get_index(self) -> int:
+        index_key = Partition.INDEX
         index = self._store.get(index_key)
         if index is None:
             index = -1
         return index
 
-    def _increase_index(self, write_batch):
-        index_key = self._get_index_key()
-        index = self._get_index()
-        index += 1
-        self._store.put(index_key, index, write_batch=write_batch)
-
-    def _get_offset(self, receiver: str):
+    def _get_offset(self, receiver: str) -> int:
         offset_key = self._get_offset_key(receiver)
         offset = self._store.get(offset_key)
         if offset is None:
             offset = -1
         return offset
 
+    def _increase_index(self, write_batch: WriteBatch):
+        next_index = self._get_index() + 1
+        if next_index > MAX_UINT:
+            raise ValueError(next_index)
+        index_key = Partition.INDEX
+        self._store.put(index_key, next_index, write_batch=write_batch)
+
     def _increase_offset(self, receiver: str):
+        next_offset = self._get_offset(receiver) + 1
+        if next_offset > MAX_UINT:
+            raise ValueError(next_offset)
         offset_key = self._get_offset_key(receiver)
-        offset = self._get_offset(receiver)
-        offset += 1
-        self._store.put(offset_key, offset)
-
-    @staticmethod
-    def _get_index_key():
-        index_key = utils.get_padded_string('', prefix='_index:')
-        return index_key
+        self._store.put(offset_key, next_offset)
 
     @staticmethod
-    def _get_offset_key(receiver):
-        offset_key = utils.get_padded_string(receiver, prefix='_offset:')
+    def _get_offset_key(receiver: str) -> bytes:
+        offset_key = Partition.OFFSET + bytes(receiver, 'utf-8')
         return offset_key
 
     @staticmethod
-    def _get_message_key(index):
-        message_key = utils.get_padded_string(str(index), prefix='message:')
+    def _get_message_key(index: int) -> bytes:
+        message_key = Partition.MESSAGE + int_to_padded_bytes(index, UINT_BYTES)
         return message_key

stopover/stopover.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-version = '1.214.0'
+version = '2.214.0'
 
 banner = f"""
   ███████████                             ███████████