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