github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/store/nbs/table_set_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 2016 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  	"context"
    26  	"testing"
    27  
    28  	lru "github.com/hashicorp/golang-lru/v2"
    29  	"github.com/stretchr/testify/assert"
    30  	"github.com/stretchr/testify/require"
    31  
    32  	"github.com/dolthub/dolt/go/store/hash"
    33  )
    34  
    35  var testChunks = [][]byte{[]byte("hello2"), []byte("goodbye2"), []byte("badbye2")}
    36  
    37  var hasManyHasAll = func([]hasRecord) (hash.HashSet, error) {
    38  	return hash.HashSet{}, nil
    39  }
    40  
    41  func TestTableSetPrependEmpty(t *testing.T) {
    42  	hasCache, err := lru.New2Q[hash.Hash, struct{}](1024)
    43  	require.NoError(t, err)
    44  	ts, err := newFakeTableSet(&UnlimitedQuotaProvider{}).append(context.Background(), newMemTable(testMemTableSize), hasManyHasAll, hasCache, &Stats{})
    45  	require.NoError(t, err)
    46  	specs, err := ts.toSpecs()
    47  	require.NoError(t, err)
    48  	assert.Empty(t, specs)
    49  }
    50  
    51  func TestTableSetPrepend(t *testing.T) {
    52  	assert := assert.New(t)
    53  	ts := newFakeTableSet(&UnlimitedQuotaProvider{})
    54  	specs, err := ts.toSpecs()
    55  	defer func() {
    56  		ts.close()
    57  	}()
    58  	require.NoError(t, err)
    59  	assert.Empty(specs)
    60  	mt := newMemTable(testMemTableSize)
    61  	mt.addChunk(computeAddr(testChunks[0]), testChunks[0])
    62  	hasCache, err := lru.New2Q[hash.Hash, struct{}](1024)
    63  	require.NoError(t, err)
    64  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
    65  	require.NoError(t, err)
    66  
    67  	firstSpecs, err := ts.toSpecs()
    68  	require.NoError(t, err)
    69  	assert.Len(firstSpecs, 1)
    70  
    71  	mt = newMemTable(testMemTableSize)
    72  	mt.addChunk(computeAddr(testChunks[1]), testChunks[1])
    73  	mt.addChunk(computeAddr(testChunks[2]), testChunks[2])
    74  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
    75  	require.NoError(t, err)
    76  
    77  	secondSpecs, err := ts.toSpecs()
    78  	require.NoError(t, err)
    79  	assert.Len(secondSpecs, 2)
    80  	assert.Equal(firstSpecs[0], secondSpecs[0])
    81  }
    82  
    83  func TestTableSetToSpecsExcludesEmptyTable(t *testing.T) {
    84  	assert := assert.New(t)
    85  	ts := newFakeTableSet(&UnlimitedQuotaProvider{})
    86  	defer func() {
    87  		ts.close()
    88  	}()
    89  	specs, err := ts.toSpecs()
    90  	require.NoError(t, err)
    91  	assert.Empty(specs)
    92  	mt := newMemTable(testMemTableSize)
    93  	mt.addChunk(computeAddr(testChunks[0]), testChunks[0])
    94  	hasCache, err := lru.New2Q[hash.Hash, struct{}](1024)
    95  	require.NoError(t, err)
    96  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
    97  	require.NoError(t, err)
    98  
    99  	mt = newMemTable(testMemTableSize)
   100  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   101  	require.NoError(t, err)
   102  
   103  	mt = newMemTable(testMemTableSize)
   104  	mt.addChunk(computeAddr(testChunks[1]), testChunks[1])
   105  	mt.addChunk(computeAddr(testChunks[2]), testChunks[2])
   106  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   107  	require.NoError(t, err)
   108  
   109  	specs, err = ts.toSpecs()
   110  	require.NoError(t, err)
   111  	assert.Len(specs, 2)
   112  }
   113  
   114  func TestTableSetFlattenExcludesEmptyTable(t *testing.T) {
   115  	assert := assert.New(t)
   116  	ts := newFakeTableSet(&UnlimitedQuotaProvider{})
   117  	defer func() {
   118  		ts.close()
   119  	}()
   120  	specs, err := ts.toSpecs()
   121  	require.NoError(t, err)
   122  	assert.Empty(specs)
   123  	mt := newMemTable(testMemTableSize)
   124  	mt.addChunk(computeAddr(testChunks[0]), testChunks[0])
   125  	hasCache, err := lru.New2Q[hash.Hash, struct{}](1024)
   126  	require.NoError(t, err)
   127  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   128  	require.NoError(t, err)
   129  
   130  	mt = newMemTable(testMemTableSize)
   131  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   132  	require.NoError(t, err)
   133  
   134  	mt = newMemTable(testMemTableSize)
   135  	mt.addChunk(computeAddr(testChunks[1]), testChunks[1])
   136  	mt.addChunk(computeAddr(testChunks[2]), testChunks[2])
   137  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   138  	require.NoError(t, err)
   139  
   140  	ts, err = ts.flatten(context.Background())
   141  	require.NoError(t, err)
   142  	assert.EqualValues(ts.Size(), 2)
   143  }
   144  
   145  func persist(t *testing.T, p tablePersister, chunks ...[]byte) {
   146  	for _, c := range chunks {
   147  		mt := newMemTable(testMemTableSize)
   148  		mt.addChunk(computeAddr(c), c)
   149  		cs, err := p.Persist(context.Background(), mt, nil, &Stats{})
   150  		require.NoError(t, err)
   151  		require.NoError(t, cs.close())
   152  	}
   153  }
   154  
   155  func TestTableSetRebase(t *testing.T) {
   156  	assert := assert.New(t)
   157  	q := NewUnlimitedMemQuotaProvider()
   158  	persister := newFakeTablePersister(q)
   159  	hasCache, err := lru.New2Q[hash.Hash, struct{}](1024)
   160  	require.NoError(t, err)
   161  
   162  	insert := func(ts tableSet, chunks ...[]byte) tableSet {
   163  		var err error
   164  		for _, c := range chunks {
   165  			mt := newMemTable(testMemTableSize)
   166  			mt.addChunk(computeAddr(c), c)
   167  			ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   168  			require.NoError(t, err)
   169  		}
   170  		return ts
   171  	}
   172  
   173  	fullTS := newTableSet(persister, q)
   174  	defer func() {
   175  		require.NoError(t, fullTS.close())
   176  	}()
   177  	specs, err := fullTS.toSpecs()
   178  	require.NoError(t, err)
   179  	assert.Empty(specs)
   180  	fullTS = insert(fullTS, testChunks...)
   181  	fullTS, err = fullTS.flatten(context.Background())
   182  	require.NoError(t, err)
   183  
   184  	ts := newTableSet(persister, q)
   185  	ts = insert(ts, testChunks[0])
   186  	assert.Equal(1, ts.Size())
   187  	ts, err = ts.flatten(context.Background())
   188  	require.NoError(t, err)
   189  	ts = insert(ts, []byte("novel"))
   190  
   191  	specs, err = fullTS.toSpecs()
   192  	require.NoError(t, err)
   193  	ts2, err := ts.rebase(context.Background(), specs, nil)
   194  	require.NoError(t, err)
   195  	defer func() {
   196  		require.NoError(t, ts2.close())
   197  	}()
   198  	err = ts.close()
   199  	require.NoError(t, err)
   200  	assert.Equal(4, ts2.Size())
   201  }
   202  
   203  func TestTableSetPhysicalLen(t *testing.T) {
   204  	assert := assert.New(t)
   205  	ts := newFakeTableSet(&UnlimitedQuotaProvider{})
   206  	defer func() {
   207  		ts.close()
   208  	}()
   209  	specs, err := ts.toSpecs()
   210  	require.NoError(t, err)
   211  	assert.Empty(specs)
   212  	mt := newMemTable(testMemTableSize)
   213  	mt.addChunk(computeAddr(testChunks[0]), testChunks[0])
   214  	hasCache, err := lru.New2Q[hash.Hash, struct{}](1024)
   215  	require.NoError(t, err)
   216  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   217  	require.NoError(t, err)
   218  
   219  	mt = newMemTable(testMemTableSize)
   220  	mt.addChunk(computeAddr(testChunks[1]), testChunks[1])
   221  	mt.addChunk(computeAddr(testChunks[2]), testChunks[2])
   222  	ts, err = ts.append(context.Background(), mt, hasManyHasAll, hasCache, &Stats{})
   223  	require.NoError(t, err)
   224  
   225  	assert.True(mustUint64(ts.physicalLen()) > indexSize(mustUint32(ts.count())))
   226  }
   227  
   228  func TestTableSetClosesOpenedChunkSourcesOnErr(t *testing.T) {
   229  	q := NewUnlimitedMemQuotaProvider()
   230  	p := newFakeTablePersister(q)
   231  	persist(t, p, testChunks...)
   232  
   233  	once := true
   234  	var specs []tableSpec
   235  	for a := range p.sources {
   236  		if once {
   237  			// map iteration is randomized
   238  			p.sourcesToFail[a] = true
   239  		}
   240  		once = false
   241  		specs = append(specs, tableSpec{a, 1})
   242  	}
   243  
   244  	ts := newTableSet(p, q)
   245  	ts2, err := ts.rebase(context.Background(), specs, &Stats{})
   246  	require.Error(t, err)
   247  
   248  	assert.NoError(t, ts.close())
   249  	assert.NoError(t, ts2.close())
   250  	assert.Equal(t, 0, int(q.Usage()))
   251  }