def __init__(self, rules):
t = RBTree()
## to avoid exceptions during LE lookup
t.insert(0, [])
nets = []
for rule in rules:
net = IPNetwork(rule.net)
t.insert(net.first, [])
t.insert(net.last+1, [])
for prio, rule in enumerate(rules):
net = IPNetwork(rule.net)
for k,v in t.iter_items(net.first, net.last+1):
v.append((prio, rule))
self.t = t
class CompletedKeys(object):
def __init__(self, max_index, min_index=0):
self._max_index = max_index
self._min_index = min_index
self.num_remaining = max_index - min_index
self._slabs = RBTree()
def _get_previous_or_none(self, index):
try:
return self._slabs.floor_item(index)
except KeyError:
return None
def is_available(self, index):
logger.debug("Testing index %s", index)
if index >= self._max_index or index < self._min_index:
logger.debug("Index out of range")
return False
try:
prev_start, prev_length = self._slabs.floor_item(index)
logger.debug("Prev range: %s-%s", prev_start,
prev_start + prev_length)
return (prev_start + prev_length) <= index
except KeyError:
return True
def mark_completed(self, start_index, past_last_index):
logger.debug("Marking the range completed: %s-%s", start_index,
past_last_index)
num_completed = min(past_last_index, self._max_index) - max(
start_index, self._min_index)
# Find the item directly before this and see if there is overlap
to_discard = set()
try:
prev_start, prev_length = self._slabs.floor_item(start_index)
max_prev_completed = prev_start + prev_length
if max_prev_completed >= start_index:
# we are going to merge with the range before us
logger.debug("Merging with the prev range: %s-%s", prev_start,
prev_start + prev_length)
to_discard.add(prev_start)
num_completed = max(
num_completed - (max_prev_completed - start_index), 0)
start_index = prev_start
past_last_index = max(past_last_index,
prev_start + prev_length)
except KeyError:
pass
# Find all keys between the start and last index and merge them into one block
for merge_start, merge_length in self._slabs.iter_items(
start_index, past_last_index + 1):
if merge_start in to_discard:
logger.debug("Already merged with block %s-%s", merge_start,
merge_start + merge_length)
continue
candidate_next_index = merge_start + merge_length
logger.debug("Merging with block %s-%s", merge_start,
candidate_next_index)
num_completed -= merge_length - max(
candidate_next_index - past_last_index, 0)
to_discard.add(merge_start)
past_last_index = max(past_last_index, candidate_next_index)
# write the new block which is fully merged
discard = False
if past_last_index >= self._max_index:
logger.debug("Discarding block and setting new max to: %s",
start_index)
self._max_index = start_index
discard = True
if start_index <= self._min_index:
logger.debug("Discarding block and setting new min to: %s",
past_last_index)
self._min_index = past_last_index
discard = True
if to_discard:
logger.debug("Discarding %s obsolete blocks", len(to_discard))
self._slabs.remove_items(to_discard)
if not discard:
logger.debug("Writing new block with range: %s-%s", start_index,
past_last_index)
self._slabs.insert(start_index, past_last_index - start_index)
# Update the number of remaining items with the adjustments we've made
assert num_completed >= 0
self.num_remaining -= num_completed
logger.debug("Total blocks: %s", len(self._slabs))
def get_block_start_index(self, block_size_estimate):
logger.debug("Total range: %s-%s", self._min_index, self._max_index)
if self._max_index <= self._min_index:
raise NoAvailableKeysError(
"All indexes have been marked completed")
num_holes = len(self._slabs) + 1
random_hole = random.randint(0, num_holes - 1)
logger.debug("Selected random hole %s with %s total holes",
random_hole, num_holes)
hole_start = self._min_index
past_hole_end = self._max_index
# Now that we have picked a hole, we need to define the bounds
if random_hole > 0:
# There will be a slab before this hole, find where it ends
bound_entries = self._slabs.nsmallest(random_hole + 1)[-2:]
left_index, left_len = bound_entries[0]
logger.debug("Left range %s-%s", left_index, left_index + left_len)
hole_start = left_index + left_len
if len(bound_entries) > 1:
right_index, right_len = bound_entries[1]
logger.debug("Right range %s-%s", right_index,
right_index + right_len)
past_hole_end, _ = bound_entries[1]
elif not self._slabs.is_empty():
right_index, right_len = self._slabs.nsmallest(1)[0]
logger.debug("Right range %s-%s", right_index,
right_index + right_len)
past_hole_end, _ = self._slabs.nsmallest(1)[0]
# Now that we have our hole bounds, select a random block from [0:len - block_size_estimate]
logger.debug("Selecting from hole range: %s-%s", hole_start,
past_hole_end)
rand_max_bound = max(hole_start, past_hole_end - block_size_estimate)
logger.debug("Rand max bound: %s", rand_max_bound)
return random.randint(hole_start, rand_max_bound)
