github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/external_iterator_test.go (about) 1 // Copyright 2022 The LevelDB-Go and Pebble and Bitalostored 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 bitalostable 6 7 import ( 8 "bytes" 9 "fmt" 10 "math" 11 "testing" 12 "time" 13 14 "github.com/cockroachdb/errors" 15 "github.com/stretchr/testify/require" 16 "github.com/zuoyebang/bitalostable/internal/base" 17 "github.com/zuoyebang/bitalostable/internal/cache" 18 "github.com/zuoyebang/bitalostable/internal/datadriven" 19 "github.com/zuoyebang/bitalostable/internal/testkeys" 20 "github.com/zuoyebang/bitalostable/internal/testkeys/blockprop" 21 "github.com/zuoyebang/bitalostable/sstable" 22 "github.com/zuoyebang/bitalostable/vfs" 23 "golang.org/x/exp/rand" 24 ) 25 26 func TestExternalIterator(t *testing.T) { 27 mem := vfs.NewMem() 28 o := &Options{ 29 FS: mem, 30 Comparer: testkeys.Comparer, 31 FormatMajorVersion: FormatRangeKeys, 32 } 33 o.EnsureDefaults() 34 d, err := Open("", o) 35 require.NoError(t, err) 36 defer func() { require.NoError(t, d.Close()) }() 37 38 datadriven.RunTest(t, "testdata/external_iterator", func(td *datadriven.TestData) string { 39 switch td.Cmd { 40 case "reset": 41 mem = vfs.NewMem() 42 return "" 43 case "build": 44 if err := runBuildCmd(td, d, mem); err != nil { 45 return err.Error() 46 } 47 return "" 48 case "iter": 49 opts := IterOptions{KeyTypes: IterKeyTypePointsAndRanges} 50 var externalIterOpts []ExternalIterOption 51 var files [][]sstable.ReadableFile 52 for _, arg := range td.CmdArgs { 53 switch arg.Key { 54 case "fwd-only": 55 externalIterOpts = append(externalIterOpts, ExternalIterForwardOnly{}) 56 case "mask-suffix": 57 opts.RangeKeyMasking.Suffix = []byte(arg.Vals[0]) 58 case "lower": 59 opts.LowerBound = []byte(arg.Vals[0]) 60 case "upper": 61 opts.UpperBound = []byte(arg.Vals[0]) 62 case "files": 63 for _, v := range arg.Vals { 64 f, err := mem.Open(v) 65 require.NoError(t, err) 66 files = append(files, []sstable.ReadableFile{f}) 67 } 68 } 69 } 70 it, err := NewExternalIter(o, &opts, files, externalIterOpts...) 71 require.NoError(t, err) 72 return runIterCmd(td, it, true /* close iter */) 73 default: 74 return fmt.Sprintf("unknown command: %s", td.Cmd) 75 } 76 }) 77 } 78 79 func TestSimpleLevelIter(t *testing.T) { 80 mem := vfs.NewMem() 81 o := &Options{ 82 FS: mem, 83 Comparer: testkeys.Comparer, 84 FormatMajorVersion: FormatRangeKeys, 85 } 86 o.EnsureDefaults() 87 d, err := Open("", o) 88 require.NoError(t, err) 89 defer func() { require.NoError(t, d.Close()) }() 90 91 datadriven.RunTest(t, "testdata/simple_level_iter", func(td *datadriven.TestData) string { 92 switch td.Cmd { 93 case "reset": 94 mem = vfs.NewMem() 95 return "" 96 case "build": 97 if err := runBuildCmd(td, d, mem); err != nil { 98 return err.Error() 99 } 100 return "" 101 case "iter": 102 var files []sstable.ReadableFile 103 for _, arg := range td.CmdArgs { 104 switch arg.Key { 105 case "files": 106 for _, v := range arg.Vals { 107 f, err := mem.Open(v) 108 require.NoError(t, err) 109 files = append(files, f) 110 } 111 } 112 } 113 readers, err := openExternalTables(o, files, 0, o.MakeReaderOptions()) 114 require.NoError(t, err) 115 defer func() { 116 for i := range readers { 117 _ = readers[i].Close() 118 } 119 }() 120 var internalIters []internalIterator 121 for i := range readers { 122 iter, err := readers[i].NewIter(nil, nil) 123 require.NoError(t, err) 124 internalIters = append(internalIters, iter) 125 } 126 it := &simpleLevelIter{cmp: o.Comparer.Compare, iters: internalIters} 127 it.init(IterOptions{}) 128 129 response := runInternalIterCmd(td, it) 130 require.NoError(t, it.Close()) 131 return response 132 default: 133 return fmt.Sprintf("unknown command: %s", td.Cmd) 134 } 135 }) 136 } 137 138 func TestSimpleIterError(t *testing.T) { 139 s := simpleLevelIter{cmp: DefaultComparer.Compare, iters: []internalIterator{&errorIter{err: errors.New("injected")}}} 140 s.init(IterOptions{}) 141 defer s.Close() 142 143 iterKey, _ := s.First() 144 require.Nil(t, iterKey) 145 require.Error(t, s.Error()) 146 } 147 148 func TestIterRandomizedMaybeFilteredKeys(t *testing.T) { 149 mem := vfs.NewMem() 150 151 seed := *seed 152 if seed == 0 { 153 seed = uint64(time.Now().UnixNano()) 154 fmt.Printf("seed: %d\n", seed) 155 } 156 rng := rand.New(rand.NewSource(seed)) 157 numKeys := 100 + rng.Intn(5000) 158 // The block property filter will exclude keys with suffixes [0, tsSeparator-1]. 159 // We use the first "part" of the keyspace below to write keys >= tsSeparator, 160 // and the second part to write keys < tsSeparator. Successive parts (if any) 161 // will contain keys at random before or after the separator. 162 tsSeparator := 10 + rng.Intn(5000) 163 const keyLen = 5 164 165 // We split the keyspace into logical "parts" which are disjoint slices of the 166 // keyspace. That is, the keyspace a-z could be comprised of parts {a-k, l-z}. 167 // We rely on this partitioning when generating timestamps to give us some 168 // predictable clustering of timestamps in sstable blocks, however it is not 169 // strictly necessary for this test. 170 alpha := testkeys.Alpha(keyLen) 171 numParts := rng.Intn(3) + 2 172 blockSize := 16 + rng.Intn(64) 173 174 c := cache.New(128 << 20) 175 defer c.Unref() 176 177 for fileIdx, twoLevelIndex := range []bool{false, true} { 178 t.Run(fmt.Sprintf("twoLevelIndex=%v", twoLevelIndex), func(t *testing.T) { 179 keys := make([][]byte, 0, numKeys) 180 181 filename := fmt.Sprintf("test-%d", fileIdx) 182 f0, err := mem.Create(filename) 183 require.NoError(t, err) 184 185 indexBlockSize := 4096 186 if twoLevelIndex { 187 indexBlockSize = 1 188 } 189 w := sstable.NewWriter(f0, sstable.WriterOptions{ 190 BlockSize: blockSize, 191 Comparer: testkeys.Comparer, 192 IndexBlockSize: indexBlockSize, 193 TableFormat: sstable.TableFormatPebblev2, 194 BlockPropertyCollectors: []func() BlockPropertyCollector{ 195 func() BlockPropertyCollector { 196 return blockprop.NewBlockPropertyCollector() 197 }, 198 }, 199 }) 200 buf := make([]byte, alpha.MaxLen()+testkeys.MaxSuffixLen) 201 valBuf := make([]byte, 20) 202 keyIdx := 0 203 for i := 0; i < numParts; i++ { 204 // The first two parts of the keyspace are special. The first one has 205 // all keys with timestamps greater than tsSeparator, while the second 206 // one has all keys with timestamps less than tsSeparator. Any additional 207 // keys could have timestamps at random before or after the tsSeparator. 208 maxKeysPerPart := numKeys / numParts 209 for j := 0; j < maxKeysPerPart; j++ { 210 ts := 0 211 if i == 0 { 212 ts = rng.Intn(5000) + tsSeparator 213 } else if i == 1 { 214 ts = rng.Intn(tsSeparator) 215 } else { 216 ts = rng.Intn(tsSeparator + 5000) 217 } 218 n := testkeys.WriteKeyAt(buf, alpha, keyIdx*alpha.Count()/numKeys, ts) 219 keys = append(keys, append([]byte(nil), buf[:n]...)) 220 randStr(valBuf, rng) 221 require.NoError(t, w.Set(buf[:n], valBuf)) 222 keyIdx++ 223 } 224 } 225 require.NoError(t, w.Close()) 226 227 // Re-open that filename for reading. 228 f1, err := mem.Open(filename) 229 require.NoError(t, err) 230 231 r, err := sstable.NewReader(f1, sstable.ReaderOptions{ 232 Cache: c, 233 Comparer: testkeys.Comparer, 234 }) 235 require.NoError(t, err) 236 defer r.Close() 237 238 filter := blockprop.NewBlockPropertyFilter(uint64(tsSeparator), math.MaxUint64) 239 filterer := sstable.NewBlockPropertiesFilterer([]BlockPropertyFilter{filter}, nil) 240 ok, err := filterer.IntersectsUserPropsAndFinishInit(r.Properties.UserProperties) 241 require.True(t, ok) 242 require.NoError(t, err) 243 244 var iter sstable.Iterator 245 iter, err = r.NewIterWithBlockPropertyFilters(nil, nil, filterer, false /* useFilterBlock */, nil /* stats */) 246 require.NoError(t, err) 247 defer iter.Close() 248 var lastSeekKey, lowerBound, upperBound []byte 249 narrowBoundsMode := false 250 251 for i := 0; i < 10000; i++ { 252 if rng.Intn(8) == 0 { 253 // Toggle narrow bounds mode. 254 if narrowBoundsMode { 255 // Reset bounds. 256 lowerBound, upperBound = nil, nil 257 iter.SetBounds(nil /* lower */, nil /* upper */) 258 } 259 narrowBoundsMode = !narrowBoundsMode 260 } 261 keyIdx := rng.Intn(len(keys)) 262 seekKey := keys[keyIdx] 263 if narrowBoundsMode { 264 // Case 1: We just entered narrow bounds mode, and both bounds 265 // are nil. Set a lower/upper bound. 266 // 267 // Case 2: The seek key is outside our last bounds. 268 // 269 // In either case, pick a narrow range of keys to set bounds on, 270 // let's say keys[keyIdx-5] and keys[keyIdx+5], before doing our 271 // seek operation. Picking narrow bounds increases the chance of 272 // monotonic bound changes. 273 cmp := testkeys.Comparer.Compare 274 case1 := lowerBound == nil && upperBound == nil 275 case2 := (lowerBound != nil && cmp(lowerBound, seekKey) > 0) || (upperBound != nil && cmp(upperBound, seekKey) <= 0) 276 if case1 || case2 { 277 lowerBound = nil 278 if keyIdx-5 >= 0 { 279 lowerBound = keys[keyIdx-5] 280 } 281 upperBound = nil 282 if keyIdx+5 < len(keys) { 283 upperBound = keys[keyIdx+5] 284 } 285 iter.SetBounds(lowerBound, upperBound) 286 } 287 // Case 3: The current seek key is within the previously-set bounds. 288 // No need to change bounds. 289 } 290 flags := base.SeekGEFlagsNone 291 if lastSeekKey != nil && bytes.Compare(seekKey, lastSeekKey) > 0 { 292 flags = flags.EnableTrySeekUsingNext() 293 } 294 lastSeekKey = append(lastSeekKey[:0], seekKey...) 295 296 newKey, _ := iter.SeekGE(seekKey, flags) 297 if newKey == nil || !bytes.Equal(newKey.UserKey, seekKey) { 298 // We skipped some keys. Check if maybeFilteredKeys is true. 299 formattedNewKey := "<nil>" 300 if newKey != nil { 301 formattedNewKey = fmt.Sprintf("%s", testkeys.Comparer.FormatKey(newKey.UserKey)) 302 } 303 require.True(t, iter.MaybeFilteredKeys(), "seeked for key = %s, got key = %s indicating block property filtering but MaybeFilteredKeys = false", testkeys.Comparer.FormatKey(seekKey), formattedNewKey) 304 } 305 } 306 }) 307 } 308 }