github.com/aclements/go-misc@v0.0.0-20240129233631-2f6ede80790c/split/example_id_test.go (about)

     1  // Copyright 2018 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package split
     6  
     7  import (
     8  	"fmt"
     9  	"sync"
    10  	"sync/atomic"
    11  )
    12  
    13  // UIDGenerator generates unique, reasonably dense integer IDs.
    14  //
    15  // The implementation supports efficient concurrent generation of IDs.
    16  // It works by retrieving batches of 256 IDs at a time from a central
    17  // ID source, and sub-allocating IDs within those batches.
    18  type UIDGenerator struct {
    19  	v        *Value
    20  	base     uint64
    21  	baseLock sync.Mutex
    22  }
    23  
    24  const batchSize = 256
    25  
    26  type uidShard struct {
    27  	next, limit uint64
    28  }
    29  
    30  // NewUIDGenerator returns a new generator for unique IDs.
    31  func NewUIDGenerator() *UIDGenerator {
    32  	g := &UIDGenerator{}
    33  	g.v = New(func(s *uidShard) {
    34  		g.baseLock.Lock()
    35  		defer g.baseLock.Unlock()
    36  		*s = uidShard{g.base, g.base + batchSize}
    37  		g.base += batchSize
    38  	})
    39  	return g
    40  }
    41  
    42  // GenUID returns a uint64 that is distinct from the uint64 returned
    43  // by every other call to GenUID on g.
    44  func (g *UIDGenerator) GenUID() uint64 {
    45  retry:
    46  	shard := g.v.Get().(*uidShard)
    47  	limit := atomic.LoadUint64(&shard.limit)
    48  	id := atomic.AddUint64(&shard.next, 1)
    49  	if id < limit {
    50  		// Fast path: we got an ID in the batch.
    51  		return id
    52  	}
    53  	// Slow path: the batch ran out. Get a new batch. This
    54  	// is tricky because multiple genUIDs could enter the
    55  	// slow path for the same shard.
    56  	g.baseLock.Lock()
    57  	// Check if another genUID already got a new batch for
    58  	// this shard.
    59  	if atomic.LoadUint64(&shard.limit) != limit {
    60  		g.baseLock.Unlock()
    61  		goto retry
    62  	}
    63  	// This genUID won the race to get a new shard for
    64  	// this batch.
    65  	base := g.base
    66  	g.base += batchSize
    67  	// Store the next first so another genUID on this
    68  	// shard will continue to fail the limit check.
    69  	atomic.StoreUint64(&shard.next, base+1)
    70  	// Now store to limit, which commits this batch.
    71  	atomic.StoreUint64(&shard.limit, base+batchSize)
    72  	g.baseLock.Unlock()
    73  	return base
    74  }
    75  
    76  func Example_idGenerator() {
    77  	ids := NewUIDGenerator()
    78  
    79  	// Generate a bunch of UIDs in parallel.
    80  	const nGoroutines = 64
    81  	const nIDs = 500
    82  	generatedIDs := make([]uint64, nGoroutines*nIDs)
    83  	var wg sync.WaitGroup
    84  	for i := 0; i < nGoroutines; i++ {
    85  		wg.Add(1)
    86  		go func(i int) {
    87  			// Generate 500 unique IDs.
    88  			for j := 0; j < nIDs; j++ {
    89  				generatedIDs[i*nIDs+j] = ids.GenUID()
    90  			}
    91  			wg.Done()
    92  		}(i)
    93  	}
    94  	wg.Wait()
    95  
    96  	// Check that all UIDs were distinct.
    97  	idMap := make(map[uint64]bool)
    98  	for _, id := range generatedIDs {
    99  		if idMap[id] {
   100  			fmt.Printf("ID %d generated more than once\n", id)
   101  			return
   102  		}
   103  		idMap[id] = true
   104  	}
   105  	// Output: All IDs were unique
   106  	fmt.Println("All IDs were unique")
   107  }