Line data    Source code

       1             : // Copyright 2023 The LevelDB-Go and Pebble Authors. All rights reserved. Use
       2             : // of this source code is governed by a BSD-style license that can be found in
       3             : // the LICENSE file.
       4             : 
       5             : package invalidating
       6             : 
       7             : import (
       8             :         "github.com/cockroachdb/pebble/internal/base"
       9             :         "github.com/cockroachdb/pebble/internal/fastrand"
      10             :         "github.com/cockroachdb/pebble/internal/invariants"
      11             : )
      12             : 
      13             : // MaybeWrapIfInvariants wraps some iterators with an invalidating iterator.
      14             : // MaybeWrapIfInvariants does nothing in non-invariant builds.
      15           2 : func MaybeWrapIfInvariants(iter base.InternalIterator) base.InternalIterator {
      16           2 :         if invariants.Enabled {
      17           2 :                 if fastrand.Uint32n(10) == 1 {
      18           2 :                         return NewIter(iter)
      19           2 :                 }
      20             :         }
      21           2 :         return iter
      22             : }
      23             : 
      24             : // iter tests unsafe key/value slice reuse by modifying the last
      25             : // returned key/value to all 1s.
      26             : type iter struct {
      27             :         iter        base.InternalIterator
      28             :         lastKey     *base.InternalKey
      29             :         lastValue   base.LazyValue
      30             :         ignoreKinds [base.InternalKeyKindMax + 1]bool
      31             :         err         error
      32             : }
      33             : 
      34             : // Option configures the behavior of an invalidating iterator.
      35             : type Option interface {
      36             :         apply(*iter)
      37             : }
      38             : 
      39             : type funcOpt func(*iter)
      40             : 
      41           1 : func (f funcOpt) apply(i *iter) { f(i) }
      42             : 
      43             : // IgnoreKinds constructs an Option that configures an invalidating iterator to
      44             : // skip trashing k/v pairs with the provided key kinds. Some iterators provided
      45             : // key stability guarantees for specific key kinds.
      46           1 : func IgnoreKinds(kinds ...base.InternalKeyKind) Option {
      47           1 :         return funcOpt(func(i *iter) {
      48           1 :                 for _, kind := range kinds {
      49           1 :                         i.ignoreKinds[kind] = true
      50           1 :                 }
      51             :         })
      52             : }
      53             : 
      54             : // NewIter constructs a new invalidating iterator that wraps the provided
      55             : // iterator, trashing buffers for previously returned keys.
      56           2 : func NewIter(originalIterator base.InternalIterator, opts ...Option) base.InternalIterator {
      57           2 :         i := &iter{iter: originalIterator}
      58           2 :         for _, opt := range opts {
      59           1 :                 opt.apply(i)
      60           1 :         }
      61           2 :         return i
      62             : }
      63             : 
      64             : func (i *iter) update(
      65             :         key *base.InternalKey, value base.LazyValue,
      66           2 : ) (*base.InternalKey, base.LazyValue) {
      67           2 :         i.trashLastKV()
      68           2 :         if key == nil {
      69           2 :                 i.lastKey = nil
      70           2 :                 i.lastValue = base.LazyValue{}
      71           2 :                 return nil, base.LazyValue{}
      72           2 :         }
      73             : 
      74           2 :         i.lastKey = &base.InternalKey{}
      75           2 :         *i.lastKey = key.Clone()
      76           2 :         i.lastValue = base.LazyValue{
      77           2 :                 ValueOrHandle: append(make([]byte, 0, len(value.ValueOrHandle)), value.ValueOrHandle...),
      78           2 :         }
      79           2 :         if value.Fetcher != nil {
      80           2 :                 fetcher := new(base.LazyFetcher)
      81           2 :                 *fetcher = *value.Fetcher
      82           2 :                 i.lastValue.Fetcher = fetcher
      83           2 :         }
      84           2 :         return i.lastKey, i.lastValue
      85             : }
      86             : 
      87           2 : func (i *iter) trashLastKV() {
      88           2 :         if i.lastKey == nil {
      89           2 :                 return
      90           2 :         }
      91           2 :         if i.ignoreKinds[i.lastKey.Kind()] {
      92           1 :                 return
      93           1 :         }
      94             : 
      95           2 :         if i.lastKey != nil {
      96           2 :                 for j := range i.lastKey.UserKey {
      97           2 :                         i.lastKey.UserKey[j] = 0xff
      98           2 :                 }
      99           2 :                 i.lastKey.Trailer = 0xffffffffffffffff
     100             :         }
     101           2 :         for j := range i.lastValue.ValueOrHandle {
     102           2 :                 i.lastValue.ValueOrHandle[j] = 0xff
     103           2 :         }
     104           2 :         if i.lastValue.Fetcher != nil {
     105           2 :                 // Not all the LazyFetcher fields are visible, so we zero out the last
     106           2 :                 // value's Fetcher struct entirely.
     107           2 :                 *i.lastValue.Fetcher = base.LazyFetcher{}
     108           2 :         }
     109             : }
     110             : 
     111           2 : func (i *iter) SeekGE(key []byte, flags base.SeekGEFlags) (*base.InternalKey, base.LazyValue) {
     112           2 :         return i.update(i.iter.SeekGE(key, flags))
     113           2 : }
     114             : 
     115             : func (i *iter) SeekPrefixGE(
     116             :         prefix, key []byte, flags base.SeekGEFlags,
     117           2 : ) (*base.InternalKey, base.LazyValue) {
     118           2 :         return i.update(i.iter.SeekPrefixGE(prefix, key, flags))
     119           2 : }
     120             : 
     121           2 : func (i *iter) SeekLT(key []byte, flags base.SeekLTFlags) (*base.InternalKey, base.LazyValue) {
     122           2 :         return i.update(i.iter.SeekLT(key, flags))
     123           2 : }
     124             : 
     125           2 : func (i *iter) First() (*base.InternalKey, base.LazyValue) {
     126           2 :         return i.update(i.iter.First())
     127           2 : }
     128             : 
     129           2 : func (i *iter) Last() (*base.InternalKey, base.LazyValue) {
     130           2 :         return i.update(i.iter.Last())
     131           2 : }
     132             : 
     133           2 : func (i *iter) Next() (*base.InternalKey, base.LazyValue) {
     134           2 :         return i.update(i.iter.Next())
     135           2 : }
     136             : 
     137           2 : func (i *iter) Prev() (*base.InternalKey, base.LazyValue) {
     138           2 :         return i.update(i.iter.Prev())
     139           2 : }
     140             : 
     141           2 : func (i *iter) NextPrefix(succKey []byte) (*base.InternalKey, base.LazyValue) {
     142           2 :         return i.update(i.iter.NextPrefix(succKey))
     143           2 : }
     144             : 
     145           2 : func (i *iter) Error() error {
     146           2 :         if err := i.iter.Error(); err != nil {
     147           1 :                 return err
     148           1 :         }
     149           2 :         return i.err
     150             : }
     151             : 
     152           2 : func (i *iter) Close() error {
     153           2 :         return i.iter.Close()
     154           2 : }
     155             : 
     156           2 : func (i *iter) SetBounds(lower, upper []byte) {
     157           2 :         i.iter.SetBounds(lower, upper)
     158           2 : }
     159             : 
     160           0 : func (i *iter) String() string {
     161           0 :         return i.iter.String()
     162           0 : }