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