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