github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/storage/memdb_test.go (about)

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