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