vitess.io/vitess@v0.16.2/go/pools/id_pool.go (about)

     1  /*
     2  Copyright 2019 The Vitess 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 pools
    18  
    19  import (
    20  	"fmt"
    21  	"sync"
    22  )
    23  
    24  // IDPool is used to ensure that the set of IDs in use concurrently never
    25  // contains any duplicates. The IDs start at 1 and increase without bound, but
    26  // will never be larger than the peak number of concurrent uses.
    27  //
    28  // IDPool's Get() and Put() methods can be used concurrently.
    29  type IDPool struct {
    30  	sync.Mutex
    31  
    32  	// used holds the set of values that have been returned to us with Put().
    33  	used map[uint32]bool
    34  	// maxUsed remembers the largest value we've given out.
    35  	maxUsed uint32
    36  }
    37  
    38  // NewIDPool creates and initializes an IDPool.
    39  func NewIDPool(initValue uint32) *IDPool {
    40  	return &IDPool{
    41  		used:    make(map[uint32]bool),
    42  		maxUsed: initValue,
    43  	}
    44  }
    45  
    46  // Get returns an ID that is unique among currently active users of this pool.
    47  func (pool *IDPool) Get() (id uint32) {
    48  	pool.Lock()
    49  	defer pool.Unlock()
    50  
    51  	// Pick a value that's been returned, if any.
    52  	for key := range pool.used {
    53  		delete(pool.used, key)
    54  		return key
    55  	}
    56  
    57  	// No recycled IDs are available, so increase the pool size.
    58  	pool.maxUsed++
    59  	return pool.maxUsed
    60  }
    61  
    62  // Put recycles an ID back into the pool for others to use. Putting back a value
    63  // or 0, or a value that is not currently "checked out", will result in a panic
    64  // because that should never happen except in the case of a programming error.
    65  func (pool *IDPool) Put(id uint32) {
    66  	pool.Lock()
    67  	defer pool.Unlock()
    68  
    69  	if id < 1 || id > pool.maxUsed {
    70  		panic(fmt.Errorf("IDPool.Put(%v): invalid value, must be in the range [1,%v]", id, pool.maxUsed))
    71  	}
    72  
    73  	if pool.used[id] {
    74  		panic(fmt.Errorf("IDPool.Put(%v): can't put value that was already recycled", id))
    75  	}
    76  
    77  	// If we're recycling maxUsed, just shrink the pool.
    78  	if id == pool.maxUsed {
    79  		pool.maxUsed = id - 1
    80  		return
    81  	}
    82  
    83  	// Add it to the set of recycled IDs.
    84  	pool.used[id] = true
    85  }