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