[r] Fix: Optimistic lock contention on HCA replicas (#6648)

src/azul/indexer/aggregate.py

@@ -33,6 +33,12 @@
 Entities = JSONs
 
 
+class TooManyEntities(Exception):
+
+    def __init__(self, max_size: int):
+        super().__init__(max_size)
+
+
 class Accumulator(metaclass=ABCMeta):
     """
     Accumulates multiple values into a single value, not necessarily of the same
@@ -92,7 +98,12 @@ class SetAccumulator(Accumulator):
     smallest values, it returns a sorted list of the first N distinct values.
     """
 
-    def __init__(self, max_size=None, key=None) -> None:
+    def __init__(self,
+                 max_size=None,
+                 key=None,
+                 *,
+                 raise_on_overflow: bool = False
+                 ) -> None:
         """
         :param max_size: the maximum number of elements to retain
 
@@ -101,15 +112,50 @@ def __init__(self, max_size=None, key=None) -> None:
                     be used. With that default key, if any None values were
                     placed in the accumulator, the first element, and only the
                     first element of the returned list will be None.
+
+        :param raise_on_overflow: If true, raise TooManyEntities if the size of
+                                  the accumulated set would exceed max_size.
         """
         super().__init__()
         self.value = set()
         self.max_size = max_size
         self.key = none_safe_key(none_last=True) if key is None else key
+        self.raise_on_overflow = raise_on_overflow
 
     def accumulate(self, value) -> bool:
         """
         :return: True, if the given value was incorporated into the set
+
+        >>> s = SetAccumulator(max_size=3)
+        >>> s.accumulate(1)
+        True
+
+        >>> s.accumulate(1)
+        False
+
+        >>> s.accumulate(2)
+        True
+
+        >>> s.accumulate([1, 2, 3])
+        True
+
+        >>> s.accumulate([2, 3])
+        False
+
+        >>> s.accumulate(4)
+        False
+
+        >>> s.get()
+        [1, 2, 3]
+
+        >>> s = SetAccumulator(max_size=3, raise_on_overflow=True)
+        >>> s.accumulate([1, 2, 3])
+        True
+
+        >>> s.accumulate(4)
+        Traceback (most recent call last):
+        ...
+        azul.indexer.aggregate.TooManyEntities: 3
         """
         if self.max_size is None or len(self.value) < self.max_size:
             before = len(self.value)
@@ -126,6 +172,8 @@ def accumulate(self, value) -> bool:
                 return False
             else:
                 assert False
+        elif self.raise_on_overflow and value not in self.value:
+            raise TooManyEntities(self.max_size)
         else:
             return False
 

src/azul/plugins/metadata/hca/indexer/aggregate.py

@@ -162,6 +162,8 @@ def _accumulator(self, field) -> Accumulator | None:
                                                          none_safe_itemgetter('value', 'unit')))
         elif field == 'donor_count':
             return UniqueValueCountAccumulator()
+        elif field == 'document_id':
+            return SetAccumulator(max_size=100, raise_on_overflow=True)
         else:
             return super()._accumulator(field)
 
@@ -197,6 +199,8 @@ class ProtocolAggregator(SimpleAggregator):
     def _accumulator(self, field) -> Accumulator | None:
         if field == 'assay_type':
             return FrequencySetAccumulator(max_size=100)
+        elif field == 'document_id':
+            return SetAccumulator(max_size=100, raise_on_overflow=True)
         else:
             return super()._accumulator(field)
 

src/azul/plugins/metadata/hca/indexer/transform.py

@@ -1430,7 +1430,12 @@ def entity_type(cls) -> str:
     @classmethod
     def hot_entity_types(cls) -> dict[EntityType, EntityType]:
         return {
-            'project': 'projects'
+            'project': 'projects',
+            'donor_organism': 'donors',
+            'analysis_protocol': 'analysis_protocols',
+            'imaging_protocol': 'imaging_protocols',
+            'library_preparation_protocol': 'library_preparation_protocols',
+            'sequencing_protocol': 'sequencing_protocols'
         }
 
     def _entities(self) -> Iterable[api.File]: