github.com/KinWaiYuen/client-go/v2@v2.5.4/internal/unionstore/memdb_test.go (about)

     1  // Copyright 2021 TiKV Authors
     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  // NOTE: The code in this file is based on code from the
    16  // TiDB project, licensed under the Apache License v 2.0
    17  //
    18  // https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/unionstore/memdb_test.go
    19  //
    20  
    21  // Copyright 2020 PingCAP, Inc.
    22  //
    23  // Copyright 2015 Wenbin Xiao
    24  //
    25  // Licensed under the Apache License, Version 2.0 (the "License");
    26  // you may not use this file except in compliance with the License.
    27  // You may obtain a copy of the License at
    28  //
    29  //     http://www.apache.org/licenses/LICENSE-2.0
    30  //
    31  // Unless required by applicable law or agreed to in writing, software
    32  // distributed under the License is distributed on an "AS IS" BASIS,
    33  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    34  // See the License for the specific language governing permissions and
    35  // limitations under the License.
    36  
    37  package unionstore
    38  
    39  import (
    40  	"encoding/binary"
    41  	"fmt"
    42  	"testing"
    43  
    44  	"github.com/KinWaiYuen/client-go/v2/kv"
    45  	leveldb "github.com/pingcap/goleveldb/leveldb/memdb"
    46  	"github.com/stretchr/testify/assert"
    47  	"github.com/stretchr/testify/require"
    48  )
    49  
    50  type KeyFlags = kv.KeyFlags
    51  
    52  func init() {
    53  	testMode = true
    54  }
    55  
    56  // DeleteKey is used in test to verify the `deleteNode` used in `vlog.revertToCheckpoint`.
    57  func (db *MemDB) DeleteKey(key []byte) {
    58  	x := db.traverse(key, false)
    59  	if x.isNull() {
    60  		return
    61  	}
    62  	db.size -= len(db.vlog.getValue(x.vptr))
    63  	db.deleteNode(x)
    64  }
    65  
    66  func TestGetSet(t *testing.T) {
    67  	require := require.New(t)
    68  
    69  	const cnt = 10000
    70  	p := fillDB(cnt)
    71  
    72  	var buf [4]byte
    73  	for i := 0; i < cnt; i++ {
    74  		binary.BigEndian.PutUint32(buf[:], uint32(i))
    75  		v, err := p.Get(buf[:])
    76  		require.Nil(err)
    77  		require.Equal(v, buf[:])
    78  	}
    79  }
    80  
    81  func TestBigKV(t *testing.T) {
    82  	assert := assert.New(t)
    83  	db := newMemDB()
    84  	db.Set([]byte{1}, make([]byte, 80<<20))
    85  	assert.Equal(db.vlog.blockSize, maxBlockSize)
    86  	assert.Equal(len(db.vlog.blocks), 1)
    87  	h := db.Staging()
    88  	db.Set([]byte{2}, make([]byte, 127<<20))
    89  	db.Release(h)
    90  	assert.Equal(db.vlog.blockSize, maxBlockSize)
    91  	assert.Equal(len(db.vlog.blocks), 2)
    92  	assert.PanicsWithValue("alloc size is larger than max block size", func() { db.Set([]byte{3}, make([]byte, maxBlockSize+1)) })
    93  }
    94  
    95  func TestIterator(t *testing.T) {
    96  	assert := assert.New(t)
    97  	const cnt = 10000
    98  	db := fillDB(cnt)
    99  
   100  	var buf [4]byte
   101  	var i int
   102  
   103  	for it, _ := db.Iter(nil, nil); it.Valid(); it.Next() {
   104  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   105  		assert.Equal(it.Key(), buf[:])
   106  		assert.Equal(it.Value(), buf[:])
   107  		i++
   108  	}
   109  	assert.Equal(i, cnt)
   110  
   111  	i--
   112  	for it, _ := db.IterReverse(nil); it.Valid(); it.Next() {
   113  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   114  		assert.Equal(it.Key(), buf[:])
   115  		assert.Equal(it.Value(), buf[:])
   116  		i--
   117  	}
   118  	assert.Equal(i, -1)
   119  }
   120  
   121  func TestDiscard(t *testing.T) {
   122  	assert := assert.New(t)
   123  
   124  	const cnt = 10000
   125  	db := newMemDB()
   126  	base := deriveAndFill(0, cnt, 0, db)
   127  	sz := db.Size()
   128  
   129  	db.Cleanup(deriveAndFill(0, cnt, 1, db))
   130  	assert.Equal(db.Len(), cnt)
   131  	assert.Equal(db.Size(), sz)
   132  
   133  	var buf [4]byte
   134  
   135  	for i := 0; i < cnt; i++ {
   136  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   137  		v, err := db.Get(buf[:])
   138  		assert.Nil(err)
   139  		assert.Equal(v, buf[:])
   140  	}
   141  
   142  	var i int
   143  	for it, _ := db.Iter(nil, nil); it.Valid(); it.Next() {
   144  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   145  		assert.Equal(it.Key(), buf[:])
   146  		assert.Equal(it.Value(), buf[:])
   147  		i++
   148  	}
   149  	assert.Equal(i, cnt)
   150  
   151  	i--
   152  	for it, _ := db.IterReverse(nil); it.Valid(); it.Next() {
   153  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   154  		assert.Equal(it.Key(), buf[:])
   155  		assert.Equal(it.Value(), buf[:])
   156  		i--
   157  	}
   158  	assert.Equal(i, -1)
   159  
   160  	db.Cleanup(base)
   161  	for i := 0; i < cnt; i++ {
   162  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   163  		_, err := db.Get(buf[:])
   164  		assert.NotNil(err)
   165  	}
   166  	it1, _ := db.Iter(nil, nil)
   167  	it := it1.(*MemdbIterator)
   168  	it.seekToFirst()
   169  	assert.False(it.Valid())
   170  	it.seekToLast()
   171  	assert.False(it.Valid())
   172  	it.seek([]byte{0xff})
   173  	assert.False(it.Valid())
   174  }
   175  
   176  func TestFlushOverwrite(t *testing.T) {
   177  	assert := assert.New(t)
   178  
   179  	const cnt = 10000
   180  	db := newMemDB()
   181  	db.Release(deriveAndFill(0, cnt, 0, db))
   182  	sz := db.Size()
   183  
   184  	db.Release(deriveAndFill(0, cnt, 1, db))
   185  
   186  	assert.Equal(db.Len(), cnt)
   187  	assert.Equal(db.Size(), sz)
   188  
   189  	var kbuf, vbuf [4]byte
   190  
   191  	for i := 0; i < cnt; i++ {
   192  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   193  		binary.BigEndian.PutUint32(vbuf[:], uint32(i+1))
   194  		v, err := db.Get(kbuf[:])
   195  		assert.Nil(err)
   196  		assert.Equal(v, vbuf[:])
   197  	}
   198  
   199  	var i int
   200  	for it, _ := db.Iter(nil, nil); it.Valid(); it.Next() {
   201  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   202  		binary.BigEndian.PutUint32(vbuf[:], uint32(i+1))
   203  		assert.Equal(it.Key(), kbuf[:])
   204  		assert.Equal(it.Value(), vbuf[:])
   205  		i++
   206  	}
   207  	assert.Equal(i, cnt)
   208  
   209  	i--
   210  	for it, _ := db.IterReverse(nil); it.Valid(); it.Next() {
   211  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   212  		binary.BigEndian.PutUint32(vbuf[:], uint32(i+1))
   213  		assert.Equal(it.Key(), kbuf[:])
   214  		assert.Equal(it.Value(), vbuf[:])
   215  		i--
   216  	}
   217  	assert.Equal(i, -1)
   218  }
   219  
   220  func TestComplexUpdate(t *testing.T) {
   221  	assert := assert.New(t)
   222  
   223  	const (
   224  		keep      = 3000
   225  		overwrite = 6000
   226  		insert    = 9000
   227  	)
   228  
   229  	db := newMemDB()
   230  	db.Release(deriveAndFill(0, overwrite, 0, db))
   231  	assert.Equal(db.Len(), overwrite)
   232  	db.Release(deriveAndFill(keep, insert, 1, db))
   233  	assert.Equal(db.Len(), insert)
   234  
   235  	var kbuf, vbuf [4]byte
   236  
   237  	for i := 0; i < insert; i++ {
   238  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   239  		binary.BigEndian.PutUint32(vbuf[:], uint32(i))
   240  		if i >= keep {
   241  			binary.BigEndian.PutUint32(vbuf[:], uint32(i+1))
   242  		}
   243  		v, err := db.Get(kbuf[:])
   244  		assert.Nil(err)
   245  		assert.Equal(v, vbuf[:])
   246  	}
   247  }
   248  
   249  func TestNestedSandbox(t *testing.T) {
   250  	assert := assert.New(t)
   251  	db := newMemDB()
   252  	h0 := deriveAndFill(0, 200, 0, db)
   253  	h1 := deriveAndFill(0, 100, 1, db)
   254  	h2 := deriveAndFill(50, 150, 2, db)
   255  	h3 := deriveAndFill(100, 120, 3, db)
   256  	h4 := deriveAndFill(0, 150, 4, db)
   257  	db.Cleanup(h4) // Discard (0..150 -> 4)
   258  	db.Release(h3) // Flush (100..120 -> 3)
   259  	db.Cleanup(h2) // Discard (100..120 -> 3) & (50..150 -> 2)
   260  	db.Release(h1) // Flush (0..100 -> 1)
   261  	db.Release(h0) // Flush (0..100 -> 1) & (0..200 -> 0)
   262  	// The final result should be (0..100 -> 1) & (101..200 -> 0)
   263  
   264  	var kbuf, vbuf [4]byte
   265  
   266  	for i := 0; i < 200; i++ {
   267  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   268  		binary.BigEndian.PutUint32(vbuf[:], uint32(i))
   269  		if i < 100 {
   270  			binary.BigEndian.PutUint32(vbuf[:], uint32(i+1))
   271  		}
   272  		v, err := db.Get(kbuf[:])
   273  		assert.Nil(err)
   274  		assert.Equal(v, vbuf[:])
   275  	}
   276  
   277  	var i int
   278  
   279  	for it, _ := db.Iter(nil, nil); it.Valid(); it.Next() {
   280  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   281  		binary.BigEndian.PutUint32(vbuf[:], uint32(i))
   282  		if i < 100 {
   283  			binary.BigEndian.PutUint32(vbuf[:], uint32(i+1))
   284  		}
   285  		assert.Equal(it.Key(), kbuf[:])
   286  		assert.Equal(it.Value(), vbuf[:])
   287  		i++
   288  	}
   289  	assert.Equal(i, 200)
   290  
   291  	i--
   292  	for it, _ := db.IterReverse(nil); it.Valid(); it.Next() {
   293  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   294  		binary.BigEndian.PutUint32(vbuf[:], uint32(i))
   295  		if i < 100 {
   296  			binary.BigEndian.PutUint32(vbuf[:], uint32(i+1))
   297  		}
   298  		assert.Equal(it.Key(), kbuf[:])
   299  		assert.Equal(it.Value(), vbuf[:])
   300  		i--
   301  	}
   302  	assert.Equal(i, -1)
   303  }
   304  
   305  func TestOverwrite(t *testing.T) {
   306  	assert := assert.New(t)
   307  
   308  	const cnt = 10000
   309  	db := fillDB(cnt)
   310  	var buf [4]byte
   311  
   312  	sz := db.Size()
   313  	for i := 0; i < cnt; i += 3 {
   314  		var newBuf [4]byte
   315  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   316  		binary.BigEndian.PutUint32(newBuf[:], uint32(i*10))
   317  		db.Set(buf[:], newBuf[:])
   318  	}
   319  	assert.Equal(db.Len(), cnt)
   320  	assert.Equal(db.Size(), sz)
   321  
   322  	for i := 0; i < cnt; i++ {
   323  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   324  		val, _ := db.Get(buf[:])
   325  		v := binary.BigEndian.Uint32(val)
   326  		if i%3 == 0 {
   327  			assert.Equal(v, uint32(i*10))
   328  		} else {
   329  			assert.Equal(v, uint32(i))
   330  		}
   331  	}
   332  
   333  	var i int
   334  
   335  	for it, _ := db.Iter(nil, nil); it.Valid(); it.Next() {
   336  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   337  		assert.Equal(it.Key(), buf[:])
   338  		v := binary.BigEndian.Uint32(it.Value())
   339  		if i%3 == 0 {
   340  			assert.Equal(v, uint32(i*10))
   341  		} else {
   342  			assert.Equal(v, uint32(i))
   343  		}
   344  		i++
   345  	}
   346  	assert.Equal(i, cnt)
   347  
   348  	i--
   349  	for it, _ := db.IterReverse(nil); it.Valid(); it.Next() {
   350  		binary.BigEndian.PutUint32(buf[:], uint32(i))
   351  		assert.Equal(it.Key(), buf[:])
   352  		v := binary.BigEndian.Uint32(it.Value())
   353  		if i%3 == 0 {
   354  			assert.Equal(v, uint32(i*10))
   355  		} else {
   356  			assert.Equal(v, uint32(i))
   357  		}
   358  		i--
   359  	}
   360  	assert.Equal(i, -1)
   361  }
   362  
   363  func TestKVLargeThanBlock(t *testing.T) {
   364  	assert := assert.New(t)
   365  	db := newMemDB()
   366  	db.Set([]byte{1}, make([]byte, 1))
   367  	db.Set([]byte{2}, make([]byte, 4096))
   368  	assert.Equal(len(db.vlog.blocks), 2)
   369  	db.Set([]byte{3}, make([]byte, 3000))
   370  	assert.Equal(len(db.vlog.blocks), 2)
   371  	val, err := db.Get([]byte{3})
   372  	assert.Nil(err)
   373  	assert.Equal(len(val), 3000)
   374  }
   375  
   376  func TestEmptyDB(t *testing.T) {
   377  	assert := assert.New(t)
   378  	db := newMemDB()
   379  	_, err := db.Get([]byte{0})
   380  	assert.NotNil(err)
   381  	it1, _ := db.Iter(nil, nil)
   382  	it := it1.(*MemdbIterator)
   383  	it.seekToFirst()
   384  	assert.False(it.Valid())
   385  	it.seekToLast()
   386  	assert.False(it.Valid())
   387  	it.seek([]byte{0xff})
   388  	assert.False(it.Valid())
   389  }
   390  
   391  func TestReset(t *testing.T) {
   392  	assert := assert.New(t)
   393  	db := fillDB(1000)
   394  	db.Reset()
   395  	_, err := db.Get([]byte{0, 0, 0, 0})
   396  	assert.NotNil(err)
   397  	it1, _ := db.Iter(nil, nil)
   398  	it := it1.(*MemdbIterator)
   399  	it.seekToFirst()
   400  	assert.False(it.Valid())
   401  	it.seekToLast()
   402  	assert.False(it.Valid())
   403  	it.seek([]byte{0xff})
   404  	assert.False(it.Valid())
   405  }
   406  
   407  func TestInspectStage(t *testing.T) {
   408  	assert := assert.New(t)
   409  
   410  	db := newMemDB()
   411  	h1 := deriveAndFill(0, 1000, 0, db)
   412  	h2 := deriveAndFill(500, 1000, 1, db)
   413  	for i := 500; i < 1500; i++ {
   414  		var kbuf [4]byte
   415  		// don't update in place
   416  		var vbuf [5]byte
   417  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   418  		binary.BigEndian.PutUint32(vbuf[:], uint32(i+2))
   419  		db.Set(kbuf[:], vbuf[:])
   420  	}
   421  	h3 := deriveAndFill(1000, 2000, 3, db)
   422  
   423  	db.InspectStage(h3, func(key []byte, _ KeyFlags, val []byte) {
   424  		k := int(binary.BigEndian.Uint32(key))
   425  		v := int(binary.BigEndian.Uint32(val))
   426  
   427  		assert.True(k >= 1000 && k < 2000)
   428  		assert.Equal(v-k, 3)
   429  	})
   430  
   431  	db.InspectStage(h2, func(key []byte, _ KeyFlags, val []byte) {
   432  		k := int(binary.BigEndian.Uint32(key))
   433  		v := int(binary.BigEndian.Uint32(val))
   434  
   435  		assert.True(k >= 500 && k < 2000)
   436  		if k < 1000 {
   437  			assert.Equal(v-k, 2)
   438  		} else {
   439  			assert.Equal(v-k, 3)
   440  		}
   441  	})
   442  
   443  	db.Cleanup(h3)
   444  	db.Release(h2)
   445  
   446  	db.InspectStage(h1, func(key []byte, _ KeyFlags, val []byte) {
   447  		k := int(binary.BigEndian.Uint32(key))
   448  		v := int(binary.BigEndian.Uint32(val))
   449  
   450  		assert.True(k >= 0 && k < 1500)
   451  		if k < 500 {
   452  			assert.Equal(v-k, 0)
   453  		} else {
   454  			assert.Equal(v-k, 2)
   455  		}
   456  	})
   457  
   458  	db.Release(h1)
   459  }
   460  
   461  func TestDirty(t *testing.T) {
   462  	assert := assert.New(t)
   463  
   464  	db := newMemDB()
   465  	db.Set([]byte{1}, []byte{1})
   466  	assert.True(db.Dirty())
   467  
   468  	db = newMemDB()
   469  	h := db.Staging()
   470  	db.Set([]byte{1}, []byte{1})
   471  	db.Cleanup(h)
   472  	assert.False(db.Dirty())
   473  
   474  	h = db.Staging()
   475  	db.Set([]byte{1}, []byte{1})
   476  	db.Release(h)
   477  	assert.True(db.Dirty())
   478  
   479  	// persistent flags will make memdb dirty.
   480  	db = newMemDB()
   481  	h = db.Staging()
   482  	db.SetWithFlags([]byte{1}, []byte{1}, kv.SetKeyLocked)
   483  	db.Cleanup(h)
   484  	assert.True(db.Dirty())
   485  
   486  	// non-persistent flags will not make memdb dirty.
   487  	db = newMemDB()
   488  	h = db.Staging()
   489  	db.SetWithFlags([]byte{1}, []byte{1}, kv.SetPresumeKeyNotExists)
   490  	db.Cleanup(h)
   491  	assert.False(db.Dirty())
   492  }
   493  
   494  func TestFlags(t *testing.T) {
   495  	assert := assert.New(t)
   496  
   497  	const cnt = 10000
   498  	db := newMemDB()
   499  	h := db.Staging()
   500  	for i := uint32(0); i < cnt; i++ {
   501  		var buf [4]byte
   502  		binary.BigEndian.PutUint32(buf[:], i)
   503  		if i%2 == 0 {
   504  			db.SetWithFlags(buf[:], buf[:], kv.SetPresumeKeyNotExists, kv.SetKeyLocked)
   505  		} else {
   506  			db.SetWithFlags(buf[:], buf[:], kv.SetPresumeKeyNotExists)
   507  		}
   508  	}
   509  	db.Cleanup(h)
   510  
   511  	for i := uint32(0); i < cnt; i++ {
   512  		var buf [4]byte
   513  		binary.BigEndian.PutUint32(buf[:], i)
   514  		_, err := db.Get(buf[:])
   515  		assert.NotNil(err)
   516  		flags, err := db.GetFlags(buf[:])
   517  		if i%2 == 0 {
   518  			assert.Nil(err)
   519  			assert.True(flags.HasLocked())
   520  			assert.False(flags.HasPresumeKeyNotExists())
   521  		} else {
   522  			assert.NotNil(err)
   523  		}
   524  	}
   525  
   526  	assert.Equal(db.Len(), 5000)
   527  	assert.Equal(db.Size(), 20000)
   528  
   529  	it1, _ := db.Iter(nil, nil)
   530  	it := it1.(*MemdbIterator)
   531  	assert.False(it.Valid())
   532  
   533  	it.includeFlags = true
   534  	it.init()
   535  
   536  	for ; it.Valid(); it.Next() {
   537  		k := binary.BigEndian.Uint32(it.Key())
   538  		assert.True(k%2 == 0)
   539  	}
   540  
   541  	for i := uint32(0); i < cnt; i++ {
   542  		var buf [4]byte
   543  		binary.BigEndian.PutUint32(buf[:], i)
   544  		db.UpdateFlags(buf[:], kv.DelKeyLocked)
   545  	}
   546  	for i := uint32(0); i < cnt; i++ {
   547  		var buf [4]byte
   548  		binary.BigEndian.PutUint32(buf[:], i)
   549  		_, err := db.Get(buf[:])
   550  		assert.NotNil(err)
   551  
   552  		// UpdateFlags will create missing node.
   553  		flags, err := db.GetFlags(buf[:])
   554  		assert.Nil(err)
   555  		assert.False(flags.HasLocked())
   556  	}
   557  }
   558  
   559  func checkConsist(t *testing.T, p1 *MemDB, p2 *leveldb.DB) {
   560  	assert := assert.New(t)
   561  
   562  	assert.Equal(p1.Len(), p2.Len())
   563  	assert.Equal(p1.Size(), p2.Size())
   564  
   565  	it1, _ := p1.Iter(nil, nil)
   566  	it2 := p2.NewIterator(nil)
   567  
   568  	var prevKey, prevVal []byte
   569  	for it2.First(); it2.Valid(); it2.Next() {
   570  		v, err := p1.Get(it2.Key())
   571  		assert.Nil(err)
   572  		assert.Equal(v, it2.Value())
   573  
   574  		assert.Equal(it1.Key(), it2.Key())
   575  		assert.Equal(it1.Value(), it2.Value())
   576  
   577  		it, _ := p1.Iter(it2.Key(), nil)
   578  		assert.Equal(it.Key(), it2.Key())
   579  		assert.Equal(it.Value(), it2.Value())
   580  
   581  		if prevKey != nil {
   582  			it, _ = p1.IterReverse(it2.Key())
   583  			assert.Equal(it.Key(), prevKey)
   584  			assert.Equal(it.Value(), prevVal)
   585  		}
   586  
   587  		it1.Next()
   588  		prevKey = it2.Key()
   589  		prevVal = it2.Value()
   590  	}
   591  
   592  	it1, _ = p1.IterReverse(nil)
   593  	for it2.Last(); it2.Valid(); it2.Prev() {
   594  		assert.Equal(it1.Key(), it2.Key())
   595  		assert.Equal(it1.Value(), it2.Value())
   596  		it1.Next()
   597  	}
   598  }
   599  
   600  func fillDB(cnt int) *MemDB {
   601  	db := newMemDB()
   602  	h := deriveAndFill(0, cnt, 0, db)
   603  	db.Release(h)
   604  	return db
   605  }
   606  
   607  func deriveAndFill(start, end, valueBase int, db *MemDB) int {
   608  	h := db.Staging()
   609  	var kbuf, vbuf [4]byte
   610  	for i := start; i < end; i++ {
   611  		binary.BigEndian.PutUint32(kbuf[:], uint32(i))
   612  		binary.BigEndian.PutUint32(vbuf[:], uint32(i+valueBase))
   613  		db.Set(kbuf[:], vbuf[:])
   614  	}
   615  	return h
   616  }
   617  
   618  const (
   619  	startIndex = 0
   620  	testCount  = 2
   621  	indexStep  = 2
   622  )
   623  
   624  func insertData(t *testing.T, buffer *MemDB) {
   625  	for i := startIndex; i < testCount; i++ {
   626  		val := encodeInt(i * indexStep)
   627  		err := buffer.Set(val, val)
   628  		assert.Nil(t, err)
   629  	}
   630  }
   631  
   632  func encodeInt(n int) []byte {
   633  	return []byte(fmt.Sprintf("%010d", n))
   634  }
   635  
   636  func decodeInt(s []byte) int {
   637  	var n int
   638  	fmt.Sscanf(string(s), "%010d", &n)
   639  	return n
   640  }
   641  
   642  func valToStr(iter Iterator) string {
   643  	val := iter.Value()
   644  	return string(val)
   645  }
   646  
   647  func checkNewIterator(t *testing.T, buffer *MemDB) {
   648  	assert := assert.New(t)
   649  	for i := startIndex; i < testCount; i++ {
   650  		val := encodeInt(i * indexStep)
   651  		iter, err := buffer.Iter(val, nil)
   652  		assert.Nil(err)
   653  		assert.Equal(iter.Key(), val)
   654  		assert.Equal(decodeInt([]byte(valToStr(iter))), i*indexStep)
   655  		iter.Close()
   656  	}
   657  
   658  	// Test iterator Next()
   659  	for i := startIndex; i < testCount-1; i++ {
   660  		val := encodeInt(i * indexStep)
   661  		iter, err := buffer.Iter(val, nil)
   662  		assert.Nil(err)
   663  		assert.Equal(iter.Key(), val)
   664  		assert.Equal(valToStr(iter), string(val))
   665  
   666  		err = iter.Next()
   667  		assert.Nil(err)
   668  		assert.True(iter.Valid())
   669  
   670  		val = encodeInt((i + 1) * indexStep)
   671  		assert.Equal(iter.Key(), val)
   672  		assert.Equal(valToStr(iter), string(val))
   673  		iter.Close()
   674  	}
   675  
   676  	// Non exist and beyond maximum seek test
   677  	iter, err := buffer.Iter(encodeInt(testCount*indexStep), nil)
   678  	assert.Nil(err)
   679  	assert.False(iter.Valid())
   680  
   681  	// Non exist but between existing keys seek test,
   682  	// it returns the smallest key that larger than the one we are seeking
   683  	inBetween := encodeInt((testCount-1)*indexStep - 1)
   684  	last := encodeInt((testCount - 1) * indexStep)
   685  	iter, err = buffer.Iter(inBetween, nil)
   686  	assert.Nil(err)
   687  	assert.True(iter.Valid())
   688  	assert.NotEqual(iter.Key(), inBetween)
   689  	assert.Equal(iter.Key(), last)
   690  	iter.Close()
   691  }
   692  
   693  func mustGet(t *testing.T, buffer *MemDB) {
   694  	for i := startIndex; i < testCount; i++ {
   695  		s := encodeInt(i * indexStep)
   696  		val, err := buffer.Get(s)
   697  		assert.Nil(t, err)
   698  		assert.Equal(t, string(val), string(s))
   699  	}
   700  }
   701  
   702  func TestKVGetSet(t *testing.T) {
   703  	buffer := newMemDB()
   704  	insertData(t, buffer)
   705  	mustGet(t, buffer)
   706  }
   707  
   708  func TestNewIterator(t *testing.T) {
   709  	assert := assert.New(t)
   710  	buffer := newMemDB()
   711  	// should be invalid
   712  	iter, err := buffer.Iter(nil, nil)
   713  	assert.Nil(err)
   714  	assert.False(iter.Valid())
   715  
   716  	insertData(t, buffer)
   717  	checkNewIterator(t, buffer)
   718  }
   719  
   720  // FnKeyCmp is the function for iterator the keys
   721  type FnKeyCmp func(key []byte) bool
   722  
   723  // NextUntil applies FnKeyCmp to each entry of the iterator until meets some condition.
   724  // It will stop when fn returns true, or iterator is invalid or an error occurs.
   725  func NextUntil(it Iterator, fn FnKeyCmp) error {
   726  	var err error
   727  	for it.Valid() && !fn(it.Key()) {
   728  		err = it.Next()
   729  		if err != nil {
   730  			return err
   731  		}
   732  	}
   733  	return nil
   734  }
   735  
   736  func TestIterNextUntil(t *testing.T) {
   737  	assert := assert.New(t)
   738  	buffer := newMemDB()
   739  	insertData(t, buffer)
   740  
   741  	iter, err := buffer.Iter(nil, nil)
   742  	assert.Nil(err)
   743  
   744  	err = NextUntil(iter, func(k []byte) bool {
   745  		return false
   746  	})
   747  	assert.Nil(err)
   748  	assert.False(iter.Valid())
   749  }
   750  
   751  func TestBasicNewIterator(t *testing.T) {
   752  	assert := assert.New(t)
   753  	buffer := newMemDB()
   754  	it, err := buffer.Iter([]byte("2"), nil)
   755  	assert.Nil(err)
   756  	assert.False(it.Valid())
   757  }
   758  
   759  func TestNewIteratorMin(t *testing.T) {
   760  	assert := assert.New(t)
   761  	kvs := []struct {
   762  		key   string
   763  		value string
   764  	}{
   765  		{"DATA_test_main_db_tbl_tbl_test_record__00000000000000000001", "lock-version"},
   766  		{"DATA_test_main_db_tbl_tbl_test_record__00000000000000000001_0002", "1"},
   767  		{"DATA_test_main_db_tbl_tbl_test_record__00000000000000000001_0003", "hello"},
   768  		{"DATA_test_main_db_tbl_tbl_test_record__00000000000000000002", "lock-version"},
   769  		{"DATA_test_main_db_tbl_tbl_test_record__00000000000000000002_0002", "2"},
   770  		{"DATA_test_main_db_tbl_tbl_test_record__00000000000000000002_0003", "hello"},
   771  	}
   772  	buffer := newMemDB()
   773  	for _, kv := range kvs {
   774  		err := buffer.Set([]byte(kv.key), []byte(kv.value))
   775  		assert.Nil(err)
   776  	}
   777  
   778  	cnt := 0
   779  	it, err := buffer.Iter(nil, nil)
   780  	assert.Nil(err)
   781  	for it.Valid() {
   782  		cnt++
   783  		err := it.Next()
   784  		assert.Nil(err)
   785  	}
   786  	assert.Equal(cnt, 6)
   787  
   788  	it, err = buffer.Iter([]byte("DATA_test_main_db_tbl_tbl_test_record__00000000000000000000"), nil)
   789  	assert.Nil(err)
   790  	assert.Equal(string(it.Key()), "DATA_test_main_db_tbl_tbl_test_record__00000000000000000001")
   791  }
   792  
   793  func TestMemDBStaging(t *testing.T) {
   794  	assert := assert.New(t)
   795  	buffer := newMemDB()
   796  	err := buffer.Set([]byte("x"), make([]byte, 2))
   797  	assert.Nil(err)
   798  
   799  	h1 := buffer.Staging()
   800  	err = buffer.Set([]byte("x"), make([]byte, 3))
   801  	assert.Nil(err)
   802  
   803  	h2 := buffer.Staging()
   804  	err = buffer.Set([]byte("yz"), make([]byte, 1))
   805  	assert.Nil(err)
   806  
   807  	v, _ := buffer.Get([]byte("x"))
   808  	assert.Equal(len(v), 3)
   809  
   810  	buffer.Release(h2)
   811  
   812  	v, _ = buffer.Get([]byte("yz"))
   813  	assert.Equal(len(v), 1)
   814  
   815  	buffer.Cleanup(h1)
   816  
   817  	v, _ = buffer.Get([]byte("x"))
   818  	assert.Equal(len(v), 2)
   819  }
   820  
   821  func TestBufferLimit(t *testing.T) {
   822  	assert := assert.New(t)
   823  	buffer := newMemDB()
   824  	buffer.bufferSizeLimit = 1000
   825  	buffer.entrySizeLimit = 500
   826  
   827  	err := buffer.Set([]byte("x"), make([]byte, 500))
   828  	assert.NotNil(err) // entry size limit
   829  
   830  	err = buffer.Set([]byte("x"), make([]byte, 499))
   831  	assert.Nil(err)
   832  	err = buffer.Set([]byte("yz"), make([]byte, 499))
   833  	assert.NotNil(err) // buffer size limit
   834  
   835  	err = buffer.Delete(make([]byte, 499))
   836  	assert.Nil(err)
   837  
   838  	err = buffer.Delete(make([]byte, 500))
   839  	assert.NotNil(err)
   840  }