def alloc(self, value=None):
# Allocates a index from the allocation list
if self._in_use.all():
idx = self._in_use.length()
if idx > self._max_alloc:
raise ResourceExhaustionError()
self._in_use.append(1)
else:
idx = self._in_use.index(0)
self._in_use[idx] = 1
try:
idx = self._get_zk_index_from_bit(idx)
except ResourceExhaustionError as e:
if self._in_use.all():
self._in_use.pop(1)
else:
self._in_use[idx] = 0
raise ResourceExhaustionError(str(e))
try:
# Create a node at path and return its integer value
id_str = "%(#)010d" % {'#': idx}
self._zookeeper_client.create_node(self._path + id_str, value)
return idx
except ResourceExistsError:
return self.alloc(value)
def _alloc_from_pools(self, pools=None):
if not pools:
raise ResourceExhaustionError()
if self._reverse:
pools = list(reversed(pools))
for pool in pools:
last_idx = self._in_use.length() - 1
pool_start = pool['start']
pool_end = pool['end']
pool_size = pool_end - pool_start + 1
if self._reverse:
start_zk_idx = pool_end
end_zk_idx = pool_start
else:
start_zk_idx = pool_start
end_zk_idx = pool_end
start_bit_idx = self._get_bit_from_zk_index(start_zk_idx)
end_bit_idx = self._get_bit_from_zk_index(end_zk_idx)
# if bitarray is less then start_bit_index,
# extend bit array to start_bit_idx and use that idx
if last_idx < start_bit_idx:
temp = bitarray(start_bit_idx - last_idx)
temp.setall(0)
self._in_use.extend(temp)
self._in_use[start_bit_idx] = 1
return start_bit_idx
# if bitarray is in between start_bit_idx and end_bit_idx
if last_idx >= start_bit_idx and last_idx <= end_bit_idx:
# we need to slice part of bitarray from
# start of the pool and end of array
pool_bitarray = self._in_use[start_bit_idx:]
else:
pool_bitarray = self._in_use[
start_bit_idx:end_bit_idx+1]
if pool_bitarray.all():
if last_idx >= end_bit_idx:
continue
idx = self._in_use.length()
self._in_use.append(1)
else:
idx = pool_bitarray.index(0)
idx += start_bit_idx
self._in_use[idx] = 1
return idx
raise ResourceExhaustionError()
def _get_zk_index_from_bit(self, idx):
size = idx
if self._reverse:
for alloc in reversed(self._alloc_list):
size -= alloc['end'] - alloc['start'] + 1
if size < 0:
return alloc['start'] - size - 1
else:
for alloc in self._alloc_list:
size -= alloc['end'] - alloc['start'] + 1
if size < 0:
return alloc['end'] + size + 1
raise ResourceExhaustionError('Cannot get zk index from bit %s' %
(idx))
def alloc(self, value):
if self._in_use.all():
idx = self._in_use.length()
if idx > self._size:
raise ResourceExhaustionError()
self._in_use.append(1)
else:
idx = self._in_use.index(0)
self._in_use[idx] = 1
idx = self._get_zk_index_from_bit(idx)
try:
# Create a node at path and return its integer value
id_str = "%(#)010d" % {'#': idx}
self._disc_service.create_node(self._path + id_str, value)
return idx
except kazoo.exceptions.NodeExistsError:
return self.alloc(value)
