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