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