github.com/cockroachdb/pebble@v1.1.2/internal/invalidating/iter.go (about)

     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  func MaybeWrapIfInvariants(iter base.InternalIterator) base.InternalIterator {
    16  	if invariants.Enabled {
    17  		if fastrand.Uint32n(10) == 1 {
    18  			return NewIter(iter)
    19  		}
    20  	}
    21  	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  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  func IgnoreKinds(kinds ...base.InternalKeyKind) Option {
    47  	return funcOpt(func(i *iter) {
    48  		for _, kind := range kinds {
    49  			i.ignoreKinds[kind] = true
    50  		}
    51  	})
    52  }
    53  
    54  // NewIter constructs a new invalidating iterator that wraps the provided
    55  // iterator, trashing buffers for previously returned keys.
    56  func NewIter(originalIterator base.InternalIterator, opts ...Option) base.InternalIterator {
    57  	i := &iter{iter: originalIterator}
    58  	for _, opt := range opts {
    59  		opt.apply(i)
    60  	}
    61  	return i
    62  }
    63  
    64  func (i *iter) update(
    65  	key *base.InternalKey, value base.LazyValue,
    66  ) (*base.InternalKey, base.LazyValue) {
    67  	i.trashLastKV()
    68  	if key == nil {
    69  		i.lastKey = nil
    70  		i.lastValue = base.LazyValue{}
    71  		return nil, base.LazyValue{}
    72  	}
    73  
    74  	i.lastKey = &base.InternalKey{}
    75  	*i.lastKey = key.Clone()
    76  	i.lastValue = base.LazyValue{
    77  		ValueOrHandle: append(make([]byte, 0, len(value.ValueOrHandle)), value.ValueOrHandle...),
    78  	}
    79  	if value.Fetcher != nil {
    80  		fetcher := new(base.LazyFetcher)
    81  		*fetcher = *value.Fetcher
    82  		i.lastValue.Fetcher = fetcher
    83  	}
    84  	return i.lastKey, i.lastValue
    85  }
    86  
    87  func (i *iter) trashLastKV() {
    88  	if i.lastKey == nil {
    89  		return
    90  	}
    91  	if i.ignoreKinds[i.lastKey.Kind()] {
    92  		return
    93  	}
    94  
    95  	if i.lastKey != nil {
    96  		for j := range i.lastKey.UserKey {
    97  			i.lastKey.UserKey[j] = 0xff
    98  		}
    99  		i.lastKey.Trailer = 0xffffffffffffffff
   100  	}
   101  	for j := range i.lastValue.ValueOrHandle {
   102  		i.lastValue.ValueOrHandle[j] = 0xff
   103  	}
   104  	if i.lastValue.Fetcher != nil {
   105  		// Not all the LazyFetcher fields are visible, so we zero out the last
   106  		// value's Fetcher struct entirely.
   107  		*i.lastValue.Fetcher = base.LazyFetcher{}
   108  	}
   109  }
   110  
   111  func (i *iter) SeekGE(key []byte, flags base.SeekGEFlags) (*base.InternalKey, base.LazyValue) {
   112  	return i.update(i.iter.SeekGE(key, flags))
   113  }
   114  
   115  func (i *iter) SeekPrefixGE(
   116  	prefix, key []byte, flags base.SeekGEFlags,
   117  ) (*base.InternalKey, base.LazyValue) {
   118  	return i.update(i.iter.SeekPrefixGE(prefix, key, flags))
   119  }
   120  
   121  func (i *iter) SeekLT(key []byte, flags base.SeekLTFlags) (*base.InternalKey, base.LazyValue) {
   122  	return i.update(i.iter.SeekLT(key, flags))
   123  }
   124  
   125  func (i *iter) First() (*base.InternalKey, base.LazyValue) {
   126  	return i.update(i.iter.First())
   127  }
   128  
   129  func (i *iter) Last() (*base.InternalKey, base.LazyValue) {
   130  	return i.update(i.iter.Last())
   131  }
   132  
   133  func (i *iter) Next() (*base.InternalKey, base.LazyValue) {
   134  	return i.update(i.iter.Next())
   135  }
   136  
   137  func (i *iter) Prev() (*base.InternalKey, base.LazyValue) {
   138  	return i.update(i.iter.Prev())
   139  }
   140  
   141  func (i *iter) NextPrefix(succKey []byte) (*base.InternalKey, base.LazyValue) {
   142  	return i.update(i.iter.NextPrefix(succKey))
   143  }
   144  
   145  func (i *iter) Error() error {
   146  	if err := i.iter.Error(); err != nil {
   147  		return err
   148  	}
   149  	return i.err
   150  }
   151  
   152  func (i *iter) Close() error {
   153  	return i.iter.Close()
   154  }
   155  
   156  func (i *iter) SetBounds(lower, upper []byte) {
   157  	i.iter.SetBounds(lower, upper)
   158  }
   159  
   160  func (i *iter) String() string {
   161  	return i.iter.String()
   162  }