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