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  }