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 itertest provides facilities for testing internal iterators.
6 : package itertest
7 :
8 : import (
9 : "bytes"
10 : "fmt"
11 : "io"
12 : "strconv"
13 : "strings"
14 : "testing"
15 :
16 : "github.com/cockroachdb/crlib/crstrings"
17 : "github.com/cockroachdb/datadriven"
18 : "github.com/cockroachdb/pebble/internal/base"
19 : "github.com/cockroachdb/pebble/internal/keyspan"
20 : "github.com/cockroachdb/pebble/internal/testkeys"
21 : "github.com/stretchr/testify/require"
22 : )
23 :
24 : type iterCmdOpts struct {
25 : fmtKV formatKV
26 : showCommands bool
27 : withoutNewlines bool
28 : stats *base.InternalIteratorStats
29 : }
30 :
31 : // An IterOpt configures the behavior of RunInternalIterCmd.
32 : type IterOpt func(*iterCmdOpts)
33 :
34 : // A formatKV configures the formatting to use when presenting key-value pairs.
35 : type formatKV func(w io.Writer, key *base.InternalKey, v []byte, iter base.InternalIterator)
36 :
37 : // Condensed configures RunInternalIterCmd to output condensed results without
38 : // values, collapsed onto a single line.
39 1 : func Condensed(opts *iterCmdOpts) {
40 1 : opts.fmtKV = condensedFormatKV
41 1 : opts.withoutNewlines = true
42 1 : }
43 :
44 : // ShowCommands configures RunInternalIterCmd to show the command in each output
45 : // line (so you don't have to visually match the line to the command).
46 1 : func ShowCommands(opts *iterCmdOpts) {
47 1 : opts.showCommands = true
48 1 : }
49 :
50 : // Verbose configures RunInternalIterCmd to output verbose results.
51 1 : func Verbose(opts *iterCmdOpts) { opts.fmtKV = verboseFormatKV }
52 :
53 : // WithSpan configures RunInternalIterCmd to print the span returned by spanFunc
54 : // after each iteration operation.
55 1 : func WithSpan(spanFunc func() *keyspan.Span) IterOpt {
56 1 : return func(opts *iterCmdOpts) {
57 1 : prevFmtKV := opts.fmtKV
58 1 : opts.fmtKV = func(w io.Writer, key *base.InternalKey, v []byte, iter base.InternalIterator) {
59 1 : prevFmtKV(w, key, v, iter)
60 1 : if s := spanFunc(); s != nil {
61 1 : fmt.Fprintf(w, " / %s", s)
62 1 : }
63 : }
64 : }
65 : }
66 :
67 : // WithStats configures RunInternalIterCmd to collect iterator stats in the
68 : // struct pointed to by stats.
69 1 : func WithStats(stats *base.InternalIteratorStats) IterOpt {
70 1 : return func(opts *iterCmdOpts) { opts.stats = stats }
71 : }
72 :
73 1 : func defaultFormatKV(w io.Writer, key *base.InternalKey, v []byte, iter base.InternalIterator) {
74 1 : if key != nil {
75 1 : fmt.Fprintf(w, "%s:%s", key.UserKey, v)
76 1 : } else if err := iter.Error(); err != nil {
77 1 : fmt.Fprintf(w, "err=%v", err)
78 1 : } else {
79 1 : fmt.Fprintf(w, ".")
80 1 : }
81 : }
82 :
83 : // condensedFormatKV is a FormatKV that outputs condensed results.
84 1 : func condensedFormatKV(w io.Writer, key *base.InternalKey, v []byte, iter base.InternalIterator) {
85 1 : if key != nil {
86 1 : fmt.Fprintf(w, "<%s:%d>", key.UserKey, key.SeqNum())
87 1 : } else if err := iter.Error(); err != nil {
88 0 : fmt.Fprintf(w, "err=%v", err)
89 1 : } else {
90 1 : fmt.Fprint(w, ".")
91 1 : }
92 : }
93 :
94 : // verboseFormatKV is a FormatKV that outputs verbose results.
95 1 : func verboseFormatKV(w io.Writer, key *base.InternalKey, v []byte, iter base.InternalIterator) {
96 1 : if key != nil {
97 1 : fmt.Fprintf(w, "%s:%s", key, v)
98 1 : return
99 1 : }
100 1 : defaultFormatKV(w, key, v, iter)
101 : }
102 :
103 : // RunInternalIterCmd evaluates a datadriven command controlling an internal
104 : // iterator, returning a string with the results of the iterator operations.
105 : func RunInternalIterCmd(
106 : t *testing.T, d *datadriven.TestData, iter base.InternalIterator, opts ...IterOpt,
107 1 : ) string {
108 1 : var buf bytes.Buffer
109 1 : RunInternalIterCmdWriter(t, &buf, d, iter, opts...)
110 1 : return buf.String()
111 1 : }
112 :
113 : // RunInternalIterCmdWriter evaluates a datadriven command controlling an
114 : // internal iterator, writing the results of the iterator operations to the
115 : // provided Writer.
116 : func RunInternalIterCmdWriter(
117 : t *testing.T, w io.Writer, d *datadriven.TestData, iter base.InternalIterator, opts ...IterOpt,
118 1 : ) {
119 1 : o := iterCmdOpts{fmtKV: defaultFormatKV}
120 1 : for _, opt := range opts {
121 1 : opt(&o)
122 1 : }
123 :
124 1 : var prefix []byte
125 1 : var prevKey []byte
126 1 : getKV := func(kv *base.InternalKV) (*base.InternalKey, []byte) {
127 1 : if kv == nil {
128 1 : prevKey = nil
129 1 : return nil, nil
130 1 : }
131 1 : prevKey = kv.K.UserKey
132 1 : v, _, err := kv.Value(nil)
133 1 : require.NoError(t, err)
134 1 : return &kv.K, v
135 : }
136 1 : lines := crstrings.Lines(d.Input)
137 1 : maxCmdLen := 1
138 1 : for _, line := range lines {
139 1 : maxCmdLen = max(maxCmdLen, len(line))
140 1 : }
141 1 : for _, line := range lines {
142 1 : parts := strings.Fields(line)
143 1 : var key *base.InternalKey
144 1 : var value []byte
145 1 : if o.showCommands {
146 1 : fmt.Fprintf(w, "%*s: ", min(maxCmdLen, 40), line)
147 1 : }
148 1 : switch parts[0] {
149 1 : case "seek-ge":
150 1 : if len(parts) < 2 || len(parts) > 3 {
151 0 : fmt.Fprint(w, "seek-ge <key> [<try-seek-using-next>]\n")
152 0 : return
153 0 : }
154 1 : prefix = nil
155 1 : var flags base.SeekGEFlags
156 1 : if len(parts) == 3 {
157 0 : if trySeekUsingNext, err := strconv.ParseBool(parts[2]); err != nil {
158 0 : fmt.Fprintf(w, "%s", err.Error())
159 0 : return
160 0 : } else if trySeekUsingNext {
161 0 : flags = flags.EnableTrySeekUsingNext()
162 0 : }
163 : }
164 1 : key, value = getKV(iter.SeekGE([]byte(strings.TrimSpace(parts[1])), flags))
165 1 : case "seek-prefix-ge":
166 1 : if len(parts) != 2 && len(parts) != 3 {
167 0 : fmt.Fprint(w, "seek-prefix-ge <key> [<try-seek-using-next>]\n")
168 0 : return
169 0 : }
170 1 : prefix = []byte(strings.TrimSpace(parts[1]))
171 1 : var flags base.SeekGEFlags
172 1 : if len(parts) == 3 {
173 1 : if trySeekUsingNext, err := strconv.ParseBool(parts[2]); err != nil {
174 0 : fmt.Fprintf(w, "%s", err.Error())
175 0 : return
176 1 : } else if trySeekUsingNext {
177 1 : flags = flags.EnableTrySeekUsingNext()
178 1 : }
179 : }
180 1 : key, value = getKV(iter.SeekPrefixGE(prefix, prefix /* key */, flags))
181 1 : case "seek-lt":
182 1 : if len(parts) != 2 {
183 0 : fmt.Fprint(w, "seek-lt <key>\n")
184 0 : return
185 0 : }
186 1 : prefix = nil
187 1 : key, value = getKV(iter.SeekLT([]byte(strings.TrimSpace(parts[1])), base.SeekLTFlagsNone))
188 1 : case "first":
189 1 : prefix = nil
190 1 : key, value = getKV(iter.First())
191 1 : case "last":
192 1 : prefix = nil
193 1 : key, value = getKV(iter.Last())
194 1 : case "next":
195 1 : key, value = getKV(iter.Next())
196 1 : case "next-prefix":
197 1 : succKey := testkeys.Comparer.ImmediateSuccessor(prevKey[:testkeys.Comparer.Split(prevKey)], nil)
198 1 : key, value = getKV(iter.NextPrefix(succKey))
199 1 : case "prev":
200 1 : key, value = getKV(iter.Prev())
201 1 : case "set-bounds":
202 1 : if len(parts) <= 1 || len(parts) > 3 {
203 0 : fmt.Fprint(w, "set-bounds lower=<lower> upper=<upper>\n")
204 0 : return
205 0 : }
206 1 : var lower []byte
207 1 : var upper []byte
208 1 : for _, part := range parts[1:] {
209 1 : arg := strings.Split(strings.TrimSpace(part), "=")
210 1 : switch arg[0] {
211 1 : case "lower":
212 1 : lower = []byte(arg[1])
213 1 : case "upper":
214 1 : upper = []byte(arg[1])
215 0 : default:
216 0 : fmt.Fprintf(w, "set-bounds: unknown arg: %s", arg)
217 0 : return
218 : }
219 : }
220 1 : iter.SetBounds(lower, upper)
221 1 : continue
222 1 : case "stats":
223 1 : if o.stats != nil {
224 1 : // The timing is non-deterministic, so set to 0.
225 1 : o.stats.BlockReadDuration = 0
226 1 : fmt.Fprintf(w, "%+v\n", *o.stats)
227 1 : }
228 1 : continue
229 1 : case "reset-stats":
230 1 : if o.stats != nil {
231 1 : *o.stats = base.InternalIteratorStats{}
232 1 : }
233 1 : continue
234 1 : case "is-lower-bound":
235 1 : // This command is specific to colblk.DataBlockIter.
236 1 : if len(parts) != 2 {
237 0 : fmt.Fprint(w, "is-lower-bound <key>\n")
238 0 : return
239 0 : }
240 1 : i := iter.(interface{ IsLowerBound(key []byte) bool })
241 1 : fmt.Fprintf(w, "%v\n", i.IsLowerBound([]byte(parts[1])))
242 1 : continue
243 0 : default:
244 0 : fmt.Fprintf(w, "unknown op: %s", parts[0])
245 0 : return
246 : }
247 1 : o.fmtKV(w, key, value, iter)
248 1 : if !o.withoutNewlines {
249 1 : fmt.Fprintln(w)
250 1 : }
251 : }
252 : }
|
