github.com/fiatjaf/generic-ristretto@v0.0.1/z/z.go (about) 1 /* 2 * Copyright 2019 Dgraph Labs, Inc. and Contributors 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 z 18 19 import ( 20 "context" 21 "github.com/cespare/xxhash/v2" 22 "log" 23 "sync" 24 ) 25 26 // GetKeyToHash will get the default KeyToHash function for the primitive types. 27 func GetKeyToHash[K any](emptyKey K) func(any) (uint64, uint64) { 28 keyAsAny := any(emptyKey) 29 30 if keyAsAny == nil { 31 log.Fatal("must provide custom KeyToHash function for pointer type") 32 } 33 switch keyAsAny.(type) { 34 case uint64: 35 case string: 36 case []byte: 37 case byte: 38 case int: 39 case int32: 40 case uint32: 41 case int64: 42 default: 43 log.Fatal("must provide custom KeyToHash function for type") 44 } 45 46 return KeyToHash 47 } 48 49 // TODO: Figure out a way to re-use memhash for the second uint64 hash, we 50 // 51 // already know that appending bytes isn't reliable for generating a 52 // second hash (see Ristretto PR #88). 53 // 54 // We also know that while the Go runtime has a runtime memhash128 55 // function, it's not possible to use it to generate [2]uint64 or 56 // anything resembling a 128bit hash, even though that's exactly what 57 // we need in this situation. 58 59 func KeyToHash(key any) (uint64, uint64) { 60 switch k := key.(type) { 61 case uint64: 62 return k, 0 63 case string: 64 return MemHashString(k), xxhash.Sum64String(k) 65 case []byte: 66 return MemHash(k), xxhash.Sum64(k) 67 case byte: 68 return uint64(k), 0 69 case int: 70 return uint64(k), 0 71 case int32: 72 return uint64(k), 0 73 case uint32: 74 return uint64(k), 0 75 case int64: 76 return uint64(k), 0 77 default: 78 panic("Key type not supported") 79 } 80 81 return 0, 0 82 } 83 84 var ( 85 dummyCloserChan <-chan struct{} 86 tmpDir string 87 ) 88 89 // Closer holds the two things we need to close a goroutine and wait for it to 90 // finish: a chan to tell the goroutine to shut down, and a WaitGroup with 91 // which to wait for it to finish shutting down. 92 type Closer struct { 93 waiting sync.WaitGroup 94 95 ctx context.Context 96 cancel context.CancelFunc 97 } 98 99 // SetTmpDir sets the temporary directory for the temporary buffers. 100 func SetTmpDir(dir string) { 101 tmpDir = dir 102 } 103 104 // NewCloser constructs a new Closer, with an initial count on the WaitGroup. 105 func NewCloser(initial int) *Closer { 106 ret := &Closer{} 107 ret.ctx, ret.cancel = context.WithCancel(context.Background()) 108 ret.waiting.Add(initial) 109 return ret 110 } 111 112 // AddRunning Add()'s delta to the WaitGroup. 113 func (lc *Closer) AddRunning(delta int) { 114 lc.waiting.Add(delta) 115 } 116 117 // Ctx can be used to get a context, which would automatically get cancelled when Signal is called. 118 func (lc *Closer) Ctx() context.Context { 119 if lc == nil { 120 return context.Background() 121 } 122 return lc.ctx 123 } 124 125 // Signal signals the HasBeenClosed signal. 126 func (lc *Closer) Signal() { 127 // Todo(ibrahim): Change Signal to return error on next badger breaking change. 128 lc.cancel() 129 } 130 131 // HasBeenClosed gets signaled when Signal() is called. 132 func (lc *Closer) HasBeenClosed() <-chan struct{} { 133 if lc == nil { 134 return dummyCloserChan 135 } 136 return lc.ctx.Done() 137 } 138 139 // Done calls Done() on the WaitGroup. 140 func (lc *Closer) Done() { 141 if lc == nil { 142 return 143 } 144 lc.waiting.Done() 145 } 146 147 // Wait waits on the WaitGroup. (It waits for NewCloser's initial value, AddRunning, and Done 148 // calls to balance out.) 149 func (lc *Closer) Wait() { 150 lc.waiting.Wait() 151 } 152 153 // SignalAndWait calls Signal(), then Wait(). 154 func (lc *Closer) SignalAndWait() { 155 lc.Signal() 156 lc.Wait() 157 } 158 159 // ZeroOut zeroes out all the bytes in the range [start, end). 160 func ZeroOut(dst []byte, start, end int) { 161 if start < 0 || start >= len(dst) { 162 return // BAD 163 } 164 if end >= len(dst) { 165 end = len(dst) 166 } 167 if end-start <= 0 { 168 return 169 } 170 Memclr(dst[start:end]) 171 // b := dst[start:end] 172 // for i := range b { 173 // b[i] = 0x0 174 // } 175 }