github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/conjoiner_test.go (about) 1 // Copyright 2019 Dolthub, Inc. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 // 15 // This file incorporates work covered by the following copyright and 16 // permission notice: 17 // 18 // Copyright 2017 Attic Labs, Inc. All rights reserved. 19 // Licensed under the Apache License, version 2.0: 20 // http://www.apache.org/licenses/LICENSE-2.0 21 22 package nbs 23 24 import ( 25 "bytes" 26 "context" 27 "encoding/binary" 28 "sort" 29 "testing" 30 31 "github.com/stretchr/testify/assert" 32 "github.com/stretchr/testify/require" 33 34 "github.com/dolthub/dolt/go/store/constants" 35 "github.com/dolthub/dolt/go/store/hash" 36 ) 37 38 type tableSpecsByAscendingCount []tableSpec 39 40 func (ts tableSpecsByAscendingCount) Len() int { return len(ts) } 41 func (ts tableSpecsByAscendingCount) Less(i, j int) bool { 42 tsI, tsJ := ts[i], ts[j] 43 if tsI.chunkCount == tsJ.chunkCount { 44 return bytes.Compare(tsI.name[:], tsJ.name[:]) < 0 45 } 46 return tsI.chunkCount < tsJ.chunkCount 47 } 48 func (ts tableSpecsByAscendingCount) Swap(i, j int) { ts[i], ts[j] = ts[j], ts[i] } 49 50 func makeTestSrcs(t *testing.T, tableSizes []uint32, p tablePersister) (srcs chunkSources) { 51 count := uint32(0) 52 nextChunk := func() (chunk []byte) { 53 chunk = make([]byte, 4) 54 binary.BigEndian.PutUint32(chunk, count) 55 count++ 56 return chunk 57 } 58 59 for _, s := range tableSizes { 60 mt := newMemTable(testMemTableSize) 61 for i := uint32(0); i < s; i++ { 62 c := nextChunk() 63 mt.addChunk(computeAddr(c), c) 64 } 65 cs, err := p.Persist(context.Background(), mt, nil, &Stats{}) 66 require.NoError(t, err) 67 srcs = append(srcs, cs.Clone()) 68 } 69 return 70 } 71 72 func TestConjoin(t *testing.T) { 73 // Makes a tableSet with len(tableSizes) upstream tables containing tableSizes[N] unique chunks 74 makeTestTableSpecs := func(tableSizes []uint32, p tablePersister) (specs []tableSpec) { 75 for _, src := range makeTestSrcs(t, tableSizes, p) { 76 specs = append(specs, tableSpec{mustAddr(src.hash()), mustUint32(src.count())}) 77 err := src.Close() 78 require.NoError(t, err) 79 } 80 return 81 } 82 83 // Returns the chunk counts of the tables in ts.compacted & ts.upstream in ascending order 84 getSortedSizes := func(specs []tableSpec) (sorted []uint32) { 85 all := append([]tableSpec{}, specs...) 86 sort.Sort(tableSpecsByAscendingCount(all)) 87 for _, ts := range all { 88 sorted = append(sorted, ts.chunkCount) 89 } 90 return 91 } 92 93 assertContainAll := func(t *testing.T, p tablePersister, expect, actual []tableSpec) { 94 open := func(specs []tableSpec) (srcs chunkReaderGroup) { 95 for _, sp := range specs { 96 cs, err := p.Open(context.Background(), sp.name, sp.chunkCount, nil) 97 98 if err != nil { 99 require.NoError(t, err) 100 } 101 102 srcs = append(srcs, cs) 103 } 104 return 105 } 106 expectSrcs, actualSrcs := open(expect), open(actual) 107 chunkChan := make(chan extractRecord, mustUint32(expectSrcs.count())) 108 err := expectSrcs.extract(context.Background(), chunkChan) 109 require.NoError(t, err) 110 close(chunkChan) 111 112 for rec := range chunkChan { 113 has, err := actualSrcs.has(rec.a) 114 require.NoError(t, err) 115 assert.True(t, has) 116 } 117 } 118 119 setup := func(lock addr, root hash.Hash, sizes []uint32) (fm *fakeManifest, p tablePersister, upstream manifestContents) { 120 p = newFakeTablePersister() 121 fm = &fakeManifest{} 122 fm.set(constants.NomsVersion, lock, root, makeTestTableSpecs(sizes, p), nil) 123 124 var err error 125 _, upstream, err = fm.ParseIfExists(context.Background(), nil, nil) 126 require.NoError(t, err) 127 128 return 129 } 130 131 // Compact some tables, interloper slips in a new table 132 makeExtra := func(p tablePersister) tableSpec { 133 mt := newMemTable(testMemTableSize) 134 data := []byte{0xde, 0xad} 135 mt.addChunk(computeAddr(data), data) 136 src, err := p.Persist(context.Background(), mt, nil, &Stats{}) 137 require.NoError(t, err) 138 return tableSpec{mustAddr(src.hash()), mustUint32(src.count())} 139 } 140 141 tc := []struct { 142 name string 143 precompact []uint32 144 postcompact []uint32 145 }{ 146 {"uniform", []uint32{1, 1, 1, 1, 1}, []uint32{5}}, 147 {"all but last", []uint32{1, 1, 1, 1, 5}, []uint32{4, 5}}, 148 {"all", []uint32{5, 5, 5}, []uint32{15}}, 149 {"first four", []uint32{5, 6, 10, 11, 35, 64}, []uint32{32, 35, 64}}, 150 {"log, first two", []uint32{1, 2, 4, 8, 16, 32, 64}, []uint32{3, 4, 8, 16, 32, 64}}, 151 {"log, all", []uint32{2, 3, 4, 8, 16, 32, 64}, []uint32{129}}, 152 } 153 154 stats := &Stats{} 155 startLock, startRoot := computeAddr([]byte("lock")), hash.Of([]byte("root")) 156 t.Run("Success", func(t *testing.T) { 157 // Compact some tables, no one interrupts 158 for _, c := range tc { 159 t.Run(c.name, func(t *testing.T) { 160 fm, p, upstream := setup(startLock, startRoot, c.precompact) 161 162 _, err := conjoin(context.Background(), upstream, fm, p, stats) 163 require.NoError(t, err) 164 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 165 require.NoError(t, err) 166 assert.True(t, exists) 167 assert.Equal(t, c.postcompact, getSortedSizes(newUpstream.specs)) 168 assertContainAll(t, p, upstream.specs, newUpstream.specs) 169 }) 170 } 171 }) 172 173 t.Run("Retry", func(t *testing.T) { 174 for _, c := range tc { 175 t.Run(c.name, func(t *testing.T) { 176 fm, p, upstream := setup(startLock, startRoot, c.precompact) 177 178 newTable := makeExtra(p) 179 u := updatePreemptManifest{fm, func() { 180 specs := append([]tableSpec{}, upstream.specs...) 181 fm.set(constants.NomsVersion, computeAddr([]byte("lock2")), startRoot, append(specs, newTable), nil) 182 }} 183 _, err := conjoin(context.Background(), upstream, u, p, stats) 184 require.NoError(t, err) 185 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 186 require.NoError(t, err) 187 assert.True(t, exists) 188 assert.Equal(t, append([]uint32{1}, c.postcompact...), getSortedSizes(newUpstream.specs)) 189 assertContainAll(t, p, append(upstream.specs, newTable), newUpstream.specs) 190 }) 191 } 192 }) 193 194 t.Run("TablesDroppedUpstream", func(t *testing.T) { 195 // Interloper drops some compactees 196 for _, c := range tc { 197 t.Run(c.name, func(t *testing.T) { 198 fm, p, upstream := setup(startLock, startRoot, c.precompact) 199 200 u := updatePreemptManifest{fm, func() { 201 fm.set(constants.NomsVersion, computeAddr([]byte("lock2")), startRoot, upstream.specs[1:], nil) 202 }} 203 _, err := conjoin(context.Background(), upstream, u, p, stats) 204 require.NoError(t, err) 205 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 206 require.NoError(t, err) 207 assert.True(t, exists) 208 assert.Equal(t, c.precompact[1:], getSortedSizes(newUpstream.specs)) 209 }) 210 } 211 }) 212 213 setupAppendix := func(lock addr, root hash.Hash, specSizes, appendixSizes []uint32) (fm *fakeManifest, p tablePersister, upstream manifestContents) { 214 p = newFakeTablePersister() 215 fm = &fakeManifest{} 216 fm.set(constants.NomsVersion, lock, root, makeTestTableSpecs(specSizes, p), makeTestTableSpecs(appendixSizes, p)) 217 218 var err error 219 _, upstream, err = fm.ParseIfExists(context.Background(), nil, nil) 220 require.NoError(t, err) 221 222 return 223 } 224 225 tca := []struct { 226 name string 227 appendix []uint32 228 precompact []uint32 229 postcompact []uint32 230 }{ 231 {"uniform", []uint32{1}, []uint32{1, 1, 1, 1, 1}, []uint32{1, 4}}, 232 {"all but last", []uint32{2}, []uint32{2, 1, 1, 1, 1, 5}, []uint32{2, 4, 5}}, 233 {"all", []uint32{1, 2, 3}, []uint32{1, 2, 3, 5, 5, 5}, []uint32{1, 2, 3, 15}}, 234 {"first four", []uint32{8, 9, 10}, []uint32{8, 9, 10, 5, 6, 10, 11, 35, 64}, []uint32{8, 9, 10, 32, 35, 64}}, 235 {"log, first two", nil, []uint32{1, 2, 4, 8, 16, 32, 64}, []uint32{3, 4, 8, 16, 32, 64}}, 236 {"log, all", []uint32{9, 10, 11, 12}, []uint32{9, 10, 11, 12, 2, 3, 4, 8, 16, 32, 64}, []uint32{9, 10, 11, 12, 129}}, 237 } 238 239 t.Run("SuccessAppendix", func(t *testing.T) { 240 // Compact some tables, no one interrupts 241 for _, c := range tca { 242 t.Run(c.name, func(t *testing.T) { 243 fm, p, upstream := setupAppendix(startLock, startRoot, c.precompact, c.appendix) 244 245 _, err := conjoin(context.Background(), upstream, fm, p, stats) 246 require.NoError(t, err) 247 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 248 require.NoError(t, err) 249 assert.True(t, exists) 250 assert.Equal(t, c.postcompact, getSortedSizes(newUpstream.specs)) 251 assert.Equal(t, c.appendix, getSortedSizes(newUpstream.appendix)) 252 assertContainAll(t, p, upstream.specs, newUpstream.specs) 253 assertContainAll(t, p, upstream.appendix, newUpstream.appendix) 254 }) 255 } 256 }) 257 258 t.Run("RetryAppendixSpecsChange", func(t *testing.T) { 259 for _, c := range tca { 260 t.Run(c.name, func(t *testing.T) { 261 fm, p, upstream := setupAppendix(startLock, startRoot, c.precompact, c.appendix) 262 263 newTable := makeExtra(p) 264 u := updatePreemptManifest{fm, func() { 265 specs := append([]tableSpec{}, upstream.specs...) 266 fm.set(constants.NomsVersion, computeAddr([]byte("lock2")), startRoot, append(specs, newTable), upstream.appendix) 267 }} 268 269 _, err := conjoin(context.Background(), upstream, u, p, stats) 270 require.NoError(t, err) 271 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 272 require.NoError(t, err) 273 assert.True(t, exists) 274 assert.Equal(t, append([]uint32{1}, c.postcompact...), getSortedSizes(newUpstream.specs)) 275 assert.Equal(t, c.appendix, getSortedSizes(newUpstream.appendix)) 276 assertContainAll(t, p, append(upstream.specs, newTable), newUpstream.specs) 277 assertContainAll(t, p, upstream.appendix, newUpstream.appendix) 278 }) 279 } 280 }) 281 282 t.Run("RetryAppendixAppendixChange", func(t *testing.T) { 283 for _, c := range tca { 284 t.Run(c.name, func(t *testing.T) { 285 fm, p, upstream := setupAppendix(startLock, startRoot, c.precompact, c.appendix) 286 287 newTable := makeExtra(p) 288 u := updatePreemptManifest{fm, func() { 289 app := append([]tableSpec{}, upstream.appendix...) 290 specs := append([]tableSpec{}, newTable) 291 fm.set(constants.NomsVersion, computeAddr([]byte("lock2")), startRoot, append(specs, upstream.specs...), append(app, newTable)) 292 }} 293 294 _, err := conjoin(context.Background(), upstream, u, p, stats) 295 require.NoError(t, err) 296 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 297 require.NoError(t, err) 298 assert.True(t, exists) 299 if newUpstream.appendix != nil { 300 assert.Equal(t, append([]uint32{1}, c.appendix...), getSortedSizes(newUpstream.appendix)) 301 assertContainAll(t, p, append(upstream.appendix, newTable), newUpstream.appendix) 302 } else { 303 assert.Equal(t, upstream.appendix, newUpstream.appendix) 304 } 305 }) 306 } 307 }) 308 309 t.Run("TablesDroppedUpstreamAppendixSpecChanges", func(t *testing.T) { 310 // Interloper drops some compactees 311 for _, c := range tca { 312 t.Run(c.name, func(t *testing.T) { 313 fm, p, upstream := setupAppendix(startLock, startRoot, c.precompact, c.appendix) 314 315 u := updatePreemptManifest{fm, func() { 316 fm.set(constants.NomsVersion, computeAddr([]byte("lock2")), startRoot, upstream.specs[len(c.appendix)+1:], upstream.appendix[:]) 317 }} 318 _, err := conjoin(context.Background(), upstream, u, p, stats) 319 require.NoError(t, err) 320 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 321 require.NoError(t, err) 322 assert.True(t, exists) 323 assert.Equal(t, c.precompact[len(c.appendix)+1:], getSortedSizes(newUpstream.specs)) 324 assert.Equal(t, c.appendix, getSortedSizes(newUpstream.appendix)) 325 }) 326 } 327 }) 328 329 t.Run("TablesDroppedUpstreamAppendixAppendixChanges", func(t *testing.T) { 330 // Interloper drops some compactees 331 for _, c := range tca { 332 t.Run(c.name, func(t *testing.T) { 333 fm, p, upstream := setupAppendix(startLock, startRoot, c.precompact, c.appendix) 334 335 newTable := makeExtra(p) 336 u := updatePreemptManifest{fm, func() { 337 specs := append([]tableSpec{}, newTable) 338 specs = append(specs, upstream.specs[len(c.appendix)+1:]...) 339 fm.set(constants.NomsVersion, computeAddr([]byte("lock2")), startRoot, specs, append([]tableSpec{}, newTable)) 340 }} 341 342 _, err := conjoin(context.Background(), upstream, u, p, stats) 343 require.NoError(t, err) 344 exists, newUpstream, err := fm.ParseIfExists(context.Background(), stats, nil) 345 require.NoError(t, err) 346 assert.True(t, exists) 347 assert.Equal(t, append([]uint32{1}, c.precompact[len(c.appendix)+1:]...), getSortedSizes(newUpstream.specs)) 348 assert.Equal(t, []uint32{1}, getSortedSizes(newUpstream.appendix)) 349 }) 350 } 351 }) 352 } 353 354 type updatePreemptManifest struct { 355 manifest 356 preUpdate func() 357 } 358 359 func (u updatePreemptManifest) Update(ctx context.Context, lastLock addr, newContents manifestContents, stats *Stats, writeHook func() error) (manifestContents, error) { 360 if u.preUpdate != nil { 361 u.preUpdate() 362 } 363 return u.manifest.Update(ctx, lastLock, newContents, stats, writeHook) 364 }