bitmap.go

Documentation: k8s.io/kubernetes/pkg/registry/core/service/allocator

     1  /*
     2  Copyright 2015 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package allocator
    18  
    19  import (
    20  	"errors"
    21  	"fmt"
    22  	"math/big"
    23  	"math/rand"
    24  	"sync"
    25  	"time"
    26  )
    27  
    28  // AllocationBitmap is a contiguous block of resources that can be allocated atomically.
    29  //
    30  // Each resource has an offset.  The internal structure is a bitmap, with a bit for each offset.
    31  //
    32  // If a resource is taken, the bit at that offset is set to one.
    33  // r.count is always equal to the number of set bits and can be recalculated at any time
    34  // by counting the set bits in r.allocated.
    35  //
    36  // TODO: use RLE and compact the allocator to minimize space.
    37  type AllocationBitmap struct {
    38  	// strategy carries the details of how to choose the next available item out of the range
    39  	strategy bitAllocator
    40  	// max is the maximum size of the usable items in the range
    41  	max int
    42  	// rangeSpec is the range specifier, matching RangeAllocation.Range
    43  	rangeSpec string
    44  
    45  	// lock guards the following members
    46  	lock sync.Mutex
    47  	// count is the number of currently allocated elements in the range
    48  	count int
    49  	// allocated is a bit array of the allocated items in the range
    50  	allocated *big.Int
    51  }
    52  
    53  // AllocationBitmap implements Interface and Snapshottable
    54  var _ Interface = &AllocationBitmap{}
    55  var _ Snapshottable = &AllocationBitmap{}
    56  
    57  // bitAllocator represents a search strategy in the allocation map for a valid item.
    58  type bitAllocator interface {
    59  	AllocateBit(allocated *big.Int, max, count int) (int, bool)
    60  }
    61  
    62  // NewAllocationMap creates an allocation bitmap using the random scan strategy.
    63  func NewAllocationMap(max int, rangeSpec string) *AllocationBitmap {
    64  	return NewAllocationMapWithOffset(max, rangeSpec, 0)
    65  }
    66  
    67  // NewAllocationMapWithOffset creates an allocation bitmap using a random scan strategy that
    68  // allows to pass an offset that divides the allocation bitmap in two blocks.
    69  // The first block of values will not be used for random value assigned by the AllocateNext()
    70  // method until the second block of values has been exhausted.
    71  // The offset value must be always smaller than the bitmap size.
    72  func NewAllocationMapWithOffset(max int, rangeSpec string, offset int) *AllocationBitmap {
    73  	a := AllocationBitmap{
    74  		strategy: randomScanStrategyWithOffset{
    75  			rand:   rand.New(rand.NewSource(time.Now().UnixNano())),
    76  			offset: offset,
    77  		},
    78  		allocated: big.NewInt(0),
    79  		count:     0,
    80  		max:       max,
    81  		rangeSpec: rangeSpec,
    82  	}
    83  
    84  	return &a
    85  }
    86  
    87  // Allocate attempts to reserve the provided item.
    88  // Returns true if it was allocated, false if it was already in use
    89  func (r *AllocationBitmap) Allocate(offset int) (bool, error) {
    90  	r.lock.Lock()
    91  	defer r.lock.Unlock()
    92  
    93  	// max is the maximum size of the usable items in the range
    94  	if offset < 0 || offset >= r.max {
    95  		return false, fmt.Errorf("offset %d out of range [0,%d]", offset, r.max)
    96  	}
    97  	if r.allocated.Bit(offset) == 1 {
    98  		return false, nil
    99  	}
   100  	r.allocated = r.allocated.SetBit(r.allocated, offset, 1)
   101  	r.count++
   102  	return true, nil
   103  }
   104  
   105  // AllocateNext reserves one of the items from the pool.
   106  // (0, false, nil) may be returned if there are no items left.
   107  func (r *AllocationBitmap) AllocateNext() (int, bool, error) {
   108  	r.lock.Lock()
   109  	defer r.lock.Unlock()
   110  
   111  	next, ok := r.strategy.AllocateBit(r.allocated, r.max, r.count)
   112  	if !ok {
   113  		return 0, false, nil
   114  	}
   115  	r.count++
   116  	r.allocated = r.allocated.SetBit(r.allocated, next, 1)
   117  	return next, true, nil
   118  }
   119  
   120  // Release releases the item back to the pool. Releasing an
   121  // unallocated item or an item out of the range is a no-op and
   122  // returns no error.
   123  func (r *AllocationBitmap) Release(offset int) error {
   124  	r.lock.Lock()
   125  	defer r.lock.Unlock()
   126  
   127  	if r.allocated.Bit(offset) == 0 {
   128  		return nil
   129  	}
   130  
   131  	r.allocated = r.allocated.SetBit(r.allocated, offset, 0)
   132  	r.count--
   133  	return nil
   134  }
   135  
   136  const (
   137  	// Find the size of a big.Word in bytes.
   138  	notZero   = uint64(^big.Word(0))
   139  	wordPower = (notZero>>8)&1 + (notZero>>16)&1 + (notZero>>32)&1
   140  	wordSize  = 1 << wordPower
   141  )
   142  
   143  // ForEach calls the provided function for each allocated bit.  The
   144  // AllocationBitmap may not be modified while this loop is running.
   145  func (r *AllocationBitmap) ForEach(fn func(int)) {
   146  	r.lock.Lock()
   147  	defer r.lock.Unlock()
   148  
   149  	words := r.allocated.Bits()
   150  	for wordIdx, word := range words {
   151  		bit := 0
   152  		for word > 0 {
   153  			if (word & 1) != 0 {
   154  				fn((wordIdx * wordSize * 8) + bit)
   155  				word = word &^ 1
   156  			}
   157  			bit++
   158  			word = word >> 1
   159  		}
   160  	}
   161  }
   162  
   163  // Has returns true if the provided item is already allocated and a call
   164  // to Allocate(offset) would fail.
   165  func (r *AllocationBitmap) Has(offset int) bool {
   166  	r.lock.Lock()
   167  	defer r.lock.Unlock()
   168  
   169  	return r.allocated.Bit(offset) == 1
   170  }
   171  
   172  // Free returns the count of items left in the range.
   173  func (r *AllocationBitmap) Free() int {
   174  	r.lock.Lock()
   175  	defer r.lock.Unlock()
   176  	return r.max - r.count
   177  }
   178  
   179  // Snapshot saves the current state of the pool.
   180  func (r *AllocationBitmap) Snapshot() (string, []byte) {
   181  	r.lock.Lock()
   182  	defer r.lock.Unlock()
   183  
   184  	return r.rangeSpec, r.allocated.Bytes()
   185  }
   186  
   187  // Restore restores the pool to the previously captured state.
   188  func (r *AllocationBitmap) Restore(rangeSpec string, data []byte) error {
   189  	r.lock.Lock()
   190  	defer r.lock.Unlock()
   191  
   192  	if r.rangeSpec != rangeSpec {
   193  		return errors.New("the provided range does not match the current range")
   194  	}
   195  
   196  	r.allocated = big.NewInt(0).SetBytes(data)
   197  	r.count = countBits(r.allocated)
   198  
   199  	return nil
   200  }
   201  
   202  // Destroy cleans up everything on shutdown.
   203  func (r *AllocationBitmap) Destroy() {
   204  }
   205  
   206  // randomScanStrategy chooses a random address from the provided big.Int, and then
   207  // scans forward looking for the next available address (it will wrap the range if
   208  // necessary).
   209  type randomScanStrategy struct {
   210  	rand *rand.Rand
   211  }
   212  
   213  func (rss randomScanStrategy) AllocateBit(allocated *big.Int, max, count int) (int, bool) {
   214  	if count >= max {
   215  		return 0, false
   216  	}
   217  	offset := rss.rand.Intn(max)
   218  	for i := 0; i < max; i++ {
   219  		at := (offset + i) % max
   220  		if allocated.Bit(at) == 0 {
   221  			return at, true
   222  		}
   223  	}
   224  	return 0, false
   225  }
   226  
   227  var _ bitAllocator = randomScanStrategy{}
   228  
   229  // randomScanStrategyWithOffset choose a random address from the provided big.Int and then scans
   230  // forward looking for the next available address. The big.Int range is subdivided so it will try
   231  // to allocate first from the reserved upper range of addresses (it will wrap the upper subrange if necessary).
   232  // If there is no free address it will try to allocate one from the lower range too.
   233  type randomScanStrategyWithOffset struct {
   234  	rand   *rand.Rand
   235  	offset int
   236  }
   237  
   238  func (rss randomScanStrategyWithOffset) AllocateBit(allocated *big.Int, max, count int) (int, bool) {
   239  	if count >= max {
   240  		return 0, false
   241  	}
   242  	// size of the upper subrange, prioritized for random allocation
   243  	subrangeMax := max - rss.offset
   244  	// try to get a value from the upper range [rss.reserved, max]
   245  	start := rss.rand.Intn(subrangeMax)
   246  	for i := 0; i < subrangeMax; i++ {
   247  		at := rss.offset + ((start + i) % subrangeMax)
   248  		if allocated.Bit(at) == 0 {
   249  			return at, true
   250  		}
   251  	}
   252  
   253  	start = rss.rand.Intn(rss.offset)
   254  	// subrange full, try to get the value from the first block before giving up.
   255  	for i := 0; i < rss.offset; i++ {
   256  		at := (start + i) % rss.offset
   257  		if allocated.Bit(at) == 0 {
   258  			return at, true
   259  		}
   260  	}
   261  	return 0, false
   262  }
   263  
   264  var _ bitAllocator = randomScanStrategyWithOffset{}
   265
View as plain text