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 }