contrib/lib9p/pytest/numalloc.py

   1 #! /usr/bin/env python
   2
   3 """
   4 Integer number allocator.
   5
   6 Basically, these keep track of a set of allocatable values in
   7 some range (you provide min and max) and let you allocate out of
   8 the range and return values into the range.
   9
  10 You may pick a value using "next since last time", or "next
  11 available after provided value".  Note that next-after will
  12 wrap around as needed (modular arithmetic style).
  13
  14 The free lists are thread-locked so that this code can be used
  15 with threads.
  16
  17     >>> a = NumAlloc(5, 10) # note closed interval: 5..10 inclusive
  18     >>> a
  19     NumAlloc(5, 10)
  20     >>> a.avail
  21     [[5, 10]]
  22     >>> a.alloc()
  23     5
  24     >>> a.avail
  25     [[6, 10]]
  26     >>> a.alloc(8)
  27     8
  28     >>> a.avail
  29     [[6, 7], [9, 10]]
  30     >>> a.free(5)
  31     >>> a.avail
  32     [[5, 7], [9, 10]]
  33     >>> a.free(8)
  34     >>> a.avail
  35     [[5, 10]]
  36
  37 Attempting to free a value that is already free is an error:
  38
  39     >>> a.free(5)
  40     Traceback (most recent call last):
  41        ...
  42     ValueError: free: 5 already available
  43
  44 You can, however, free a value that is outside the min/max
  45 range.  You can also free multiple values at once:
  46
  47     >>> a.free_multi([0, 1, 2, 4])
  48     >>> a.avail
  49     [[0, 2], [4, 10]]
  50     >>> a.free_multi([3, 12])
  51     >>> a.avail
  52     [[0, 10], [12, 12]]
  53
  54 Note that this changes the min/max values:
  55
  56     >>> a
  57     NumAlloc(0, 12)
  58
  59 To prevent adding values outside the min/max range, create the
  60 NumArray with autoextend=False, or set .autoextend=False at any
  61 time:
  62
  63     >>> a.autoextend = False
  64     >>> a
  65     NumAlloc(0, 12, autoextend=False)
  66     >>> a.free(13)
  67     Traceback (most recent call last):
  68        ...
  69     ValueError: free: 13 is outside range limit
  70
  71 You can create an empty range, which is really only useful once
  72 you free values into it:
  73
  74     >>> r = NumAlloc(0, -1)
  75     >>> r
  76     NumAlloc(0, -1)
  77     >>> r.alloc() is None
  78     True
  79     >>> r.free_multi(range(50))
  80     >>> r
  81     NumAlloc(0, 49)
  82
  83 Note that r.alloc() starts from where you last left off, even if
  84 you've freed a value:
  85
  86     >>> r.alloc()
  87     0
  88     >>> r.free(0)
  89     >>> r.alloc()
  90     1
  91
  92 Of course, in multithreaded code you can't really depend on this
  93 since it will race other threads.  Still, it generally makes for
  94 efficient allocation.  To force allocation to start from the
  95 range's minimum, provide the minimum (e.g., r.min_val) as an
  96 argument to r.alloc():
  97
  98     >>> r.alloc()
  99     2
 100     >>> r.alloc(r.min_val)
 101     0
 102
 103 Providing a number to alloc() tries to allocate that number,
 104 but wraps around to the next one if needed:
 105
 106     >>> r.alloc(49)
 107     49
 108     >>> r.alloc(49)
 109     3
 110     >>> r.alloc(99999)
 111     4
 112     >>> r.avail
 113     [[5, 48]]
 114
 115 There is currently no way to find all allocated values, although
 116 the obvious method (going through r.avail) will work.  Any iterator
 117 would not be thread-safe.
 118 """
 119
 120 import threading
 121
 122 class NumAlloc(object):
 123     """
 124     Number allocator object.
 125     """
 126     def __init__(self, min_val, max_val, autoextend=True):
 127         self.min_val = min_val
 128         self.max_val = max_val
 129         if min_val <= max_val:
 130             self.avail = [[min_val, max_val]]
 131         else:
 132             self.avail = []
 133         self.autoextend = autoextend
 134         self.last = None
 135         self.lock = threading.Lock()
 136
 137     def __repr__(self):
 138         myname = self.__class__.__name__
 139         if self.autoextend:
 140             ae = ''
 141         else:
 142             ae = ', autoextend=False'
 143         return '{0}({1}, {2}{3})'.format(myname, self.min_val, self.max_val, ae)
 144
 145     def _find_block(self, val):
 146         """
 147         Find the block that contains val, or that should contain val.
 148         Remember that self.avail is a list of avaliable ranges of
 149         the form [[min1, max1], [min2, max2], ..., [minN, maxN]]
 150         where max1 < min2, max2 < min3, ..., < minN.
 151
 152         The input value either falls into one of the available
 153         blocks, or falls into a gap between two available blocks.
 154         We want to know which block it goes in, or if it goes
 155         between two, which block it comes before.
 156
 157         We can do a binary search to find this block.  When we
 158         find it, return its index and its values.
 159
 160         If we find that val is not in a block, return the position
 161         where the value should go, were it to be put into a new
 162         block by itself.  E.g., suppose val is 17, and there is a
 163         block [14,16] and a block [18,20]. We would make this
 164         [14,16],[17,17],[18,20] by inserting [17,17] between them.
 165         (Afterward, we will want to fuse all three blocks to make
 166         [14,18].  However, if we insert as block 0, e.g., if the
 167         list starts with [18,20] and we insert to get
 168         [17,17][18,20], we really end up just modifying block 0 to
 169         [17,20].  Or, if we insert as the new final block, we
 170         might end up modifying the last block.)
 171         """
 172         low = 0
 173         high = len(self.avail) - 1
 174         while low <= high:
 175             mid = low + ((high - low) // 2)
 176             pair = self.avail[mid]
 177             if val < pair[0]:
 178                 # must go before block mid
 179                 high = mid - 1
 180             elif val > pair[1]:
 181                 # must go after block mid
 182                 low = mid + 1
 183             else:
 184                 # val >= first and val <= last, so we found it
 185                 return mid, pair
 186         # Low > high: no block actually contains val, or
 187         # there are no blocks at all.  If there are no blocks,
 188         # return block #0 and None.  Otherwise return the
 189         return low, None
 190
 191     def alloc(self, val=None):
 192         """
 193         Get new available value.
 194
 195         If val is None, we start from the most recently
 196         allocated value, plus 1.
 197
 198         If val is a numeric value, we start from that value.
 199         Hence, since the range is min_val..max_val, you can
 200         provide min_val to take the first available value.
 201
 202         This may return None, if no values are still available.
 203         """
 204         with self.lock:
 205             if val is None:
 206                 val = self.last + 1 if self.last is not None else self.min_val
 207             if val is None or val > self.max_val or val < self.min_val:
 208                 val = self.min_val
 209             i, pair = self._find_block(val)
 210             if pair is None:
 211                 # Value is is not available.  The next
 212                 # available value that is greater than val
 213                 # is in the block right after block i.
 214                 # If there is no block after i, the next
 215                 # available value is in block 0.  If there
 216                 # is no block 0, there are no available
 217                 # values.
 218                 nblocks = len(self.avail)
 219                 i += 1
 220                 if i >= nblocks:
 221                     if nblocks == 0:
 222                         return None
 223                     i = 0
 224                 pair = self.avail[i]
 225                 val = pair[0]
 226             # Value val is available - take it.
 227             #
 228             # There are four special cases to handle.
 229             #
 230             # 1. pair[0] < val < pair[1]: split the pair.
 231             # 2. pair[0] == val < pair[1]: increase pair[0].
 232             # 3. pair[0] == val == pair[1]: delete the pair
 233             # 4. pair[0] < val == pair[1]: decrease pair[1].
 234             assert pair[0] <= val <= pair[1]
 235             if pair[0] == val:
 236                 # case 2 or 3: Take the left edge or delete the pair.
 237                 if val == pair[1]:
 238                     del self.avail[i]
 239                 else:
 240                     pair[0] = val + 1
 241             else:
 242                 # case 1 or 4: split the pair or take the right edge.
 243                 if val == pair[1]:
 244                     pair[1] = val - 1
 245                 else:
 246                     newpair = [val + 1, pair[1]]
 247                     pair[1] = val - 1
 248                     self.avail.insert(i + 1, newpair)
 249             self.last = val
 250             return val
 251
 252     def free(self, val):
 253         "Free one value"
 254         self._free_multi('free', [val])
 255
 256     def free_multi(self, values):
 257         "Free many values (provide any iterable)"
 258         values = list(values)
 259         values.sort()
 260         self._free_multi('free_multi', values)
 261
 262     def _free_multi(self, how, values):
 263         """
 264         Free a (sorted) list of values.
 265         """
 266         if len(values) == 0:
 267             return
 268         with self.lock:
 269             while values:
 270                 # Take highest value, and any contiguous lower values.
 271                 # Note that it can be significantly faster this way
 272                 # since coalesced ranges make for shorter copies.
 273                 highval = values.pop()
 274                 val = highval
 275                 while len(values) and values[-1] == val - 1:
 276                     val = values.pop()
 277                 self._free_range(how, val, highval)
 278
 279     def _maybe_increase_max(self, how, val):
 280         """
 281         If needed, widen our range to include new high val -- i.e.,
 282         possibly increase self.max_val.  Do nothing if this is not a
 283         new all time high; fail if we have autoextend disabled.
 284         """
 285         if val <= self.max_val:
 286             return
 287         if self.autoextend:
 288             self.max_val = val
 289             return
 290         raise ValueError('{0}: {1} is outside range limit'.format(how, val))
 291
 292     def _maybe_decrease_min(self, how, val):
 293         """
 294         If needed, widen our range to include new low val -- i.e.,
 295         possibly decrease self.min_val.  Do nothing if this is not a
 296         new all time low; fail if we have autoextend disabled.
 297         """
 298         if val >= self.min_val:
 299             return
 300         if self.autoextend:
 301             self.min_val = val
 302             return
 303         raise ValueError('{0}: {1} is outside range limit'.format(how, val))
 304
 305     def _free_range(self, how, val, highval):
 306         """
 307         Free the range [val..highval].  Note, val==highval it's just
 308         a one-element range.
 309
 310         The lock is already held.
 311         """
 312         # Find the place to store the lower value.
 313         # We should never find an actual pair here.
 314         i, pair = self._find_block(val)
 315         if pair:
 316             raise ValueError('{0}: {1} already available'.format(how, val))
 317         # If we're freeing a range, check that the high val
 318         # does not span into the *next* range, either.
 319         if highval > val and i < len(self.avail):
 320             if self.avail[i][0] <= highval:
 321                 raise ValueError('{0}: {2} (from {{1}..{2}) already '
 322                                  'available'.format(how, val, highval))
 323
 324         # We'll need to insert a block and perhaps fuse it
 325         # with blocks before and/or after.  First, check
 326         # whether there *is* a before and/or after, and find
 327         # their corresponding edges and whether we abut them.
 328         if i > 0:
 329             abuts_below = self.avail[i - 1][1] + 1 == val
 330         else:
 331             abuts_below = False
 332         if i < len(self.avail):
 333             abuts_above = self.avail[i][0] - 1 == highval
 334         else:
 335             abuts_above = False
 336         # Now there are these four cases:
 337         # 1. abuts below and above: fuse the two blocks.
 338         # 2. abuts below only: adjust previous (i-1'th) block
 339         # 3. abuts above only: adjust next (i'th) block
 340         # 4. doesn't abut: insert new block
 341         if abuts_below:
 342             if abuts_above:
 343                 # case 1
 344                 self.avail[i - 1][1] = self.avail[i][1]
 345                 del self.avail[i]
 346             else:
 347                 # case 2
 348                 self._maybe_increase_max(how, highval)
 349                 self.avail[i - 1][1] = highval
 350         else:
 351             if abuts_above:
 352                 # case 3
 353                 self._maybe_decrease_min(how, val)
 354                 self.avail[i][0] = val
 355             else:
 356                 # case 4
 357                 self._maybe_decrease_min(how, val)
 358                 self._maybe_increase_max(how, highval)
 359                 newblock = [val, highval]
 360                 self.avail.insert(i, newblock)
 361
 362 if __name__ == '__main__':
 363     import doctest
 364     import sys
 365
 366     doctest.testmod()
 367     if sys.version_info[0] >= 3:
 368         xrange = range
 369     # run some worst case tests
 370     # NB: coalesce is terribly slow when done bottom up
 371     r = NumAlloc(0, 2**16 - 1)
 372     for i in xrange(r.min_val, r.max_val, 2):
 373         r.alloc(i)
 374     print('worst case alloc: len(r.avail) = {0}'.format(len(r.avail)))
 375     for i in xrange(r.max_val - 1, r.min_val, -2):
 376         r.free(i)
 377     print('free again; len(r.avail) should be 1; is {0}'.format(len(r.avail)))
 378     if len(r.avail) != 1:
 379         sys.exit('failure')