1 /*
2 * SPDX-FileCopyrightText: © Hypermode Inc. <hello@hypermode.com>
3 * SPDX-License-Identifier: Apache-2.0
4 */
5
6 package z
7
8 import (
9 "context"
10 "sync"
11
12 "github.com/cespare/xxhash/v2"
13 )
14
15 type Key interface {
16 uint64 | string | []byte | byte | int | uint | int32 | uint32 | int64
17 }
18
19 // TODO: Figure out a way to re-use memhash for the second uint64 hash,
20 // we already know that appending bytes isn't reliable for generating a
21 // second hash (see Ristretto PR #88).
22 // We also know that while the Go runtime has a runtime memhash128
23 // function, it's not possible to use it to generate [2]uint64 or
24 // anything resembling a 128bit hash, even though that's exactly what
25 // we need in this situation.
26 func KeyToHash[K Key](key K) (uint64, uint64) {
27 keyAsAny := any(key)
28 switch k := keyAsAny.(type) {
29 case uint64:
30 return k, 0
31 case string:
32 return MemHashString(k), xxhash.Sum64String(k)
33 case []byte:
34 return MemHash(k), xxhash.Sum64(k)
35 case byte:
36 return uint64(k), 0
37 case uint:
38 return uint64(k), 0
39 case int:
40 return uint64(k), 0
41 case int32:
42 return uint64(k), 0
43 case uint32:
44 return uint64(k), 0
45 case int64:
46 return uint64(k), 0
47 default:
48 panic("Key type not supported")
49 }
50 }
51
52 var (
53 dummyCloserChan <-chan struct{}
54 tmpDir string
55 )
56
57 // Closer holds the two things we need to close a goroutine and wait for it to
58 // finish: a chan to tell the goroutine to shut down, and a WaitGroup with
59 // which to wait for it to finish shutting down.
60 type Closer struct {
61 waiting sync.WaitGroup
62
63 ctx context.Context
64 cancel context.CancelFunc
65 }
66
67 // SetTmpDir sets the temporary directory for the temporary buffers.
68 func SetTmpDir(dir string) {
69 tmpDir = dir
70 }
71
72 // NewCloser constructs a new Closer, with an initial count on the WaitGroup.
73 func NewCloser(initial int) *Closer {
74 ret := &Closer{}
75 ret.ctx, ret.cancel = context.WithCancel(context.Background())
76 ret.waiting.Add(initial)
77 return ret
78 }
79
80 // AddRunning Add()'s delta to the WaitGroup.
81 func (lc *Closer) AddRunning(delta int) {
82 lc.waiting.Add(delta)
83 }
84
85 // Ctx can be used to get a context, which would automatically get cancelled when Signal is called.
86 func (lc *Closer) Ctx() context.Context {
87 if lc == nil {
88 return context.Background()
89 }
90 return lc.ctx
91 }
92
93 // Signal signals the HasBeenClosed signal.
94 func (lc *Closer) Signal() {
95 // Todo(ibrahim): Change Signal to return error on next badger breaking change.
96 lc.cancel()
97 }
98
99 // HasBeenClosed gets signaled when Signal() is called.
100 func (lc *Closer) HasBeenClosed() <-chan struct{} {
101 if lc == nil {
102 return dummyCloserChan
103 }
104 return lc.ctx.Done()
105 }
106
107 // Done calls Done() on the WaitGroup.
108 func (lc *Closer) Done() {
109 if lc == nil {
110 return
111 }
112 lc.waiting.Done()
113 }
114
115 // Wait waits on the WaitGroup. (It waits for NewCloser's initial value, AddRunning, and Done
116 // calls to balance out.)
117 func (lc *Closer) Wait() {
118 lc.waiting.Wait()
119 }
120
121 // SignalAndWait calls Signal(), then Wait().
122 func (lc *Closer) SignalAndWait() {
123 lc.Signal()
124 lc.Wait()
125 }
126
127 // ZeroOut zeroes out all the bytes in the range [start, end).
128 func ZeroOut(dst []byte, start, end int) {
129 if start < 0 || start >= len(dst) {
130 return // BAD
131 }
132 if end >= len(dst) {
133 end = len(dst)
134 }
135 if end-start <= 0 {
136 return
137 }
138 Memclr(dst[start:end])
139 // b := dst[start:end]
140 // for i := range b {
141 // b[i] = 0x0
142 // }
143 }
144