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