github.com/unicornultrafoundation/go-u2u@v1.0.0-rc1.0.20240205080301-e74a83d3fadc/utils/dbutil/switchable/snapshot_test.go (about)

     1  package switchable
     2  
     3  import (
     4  	"math/rand"
     5  	"sync"
     6  	"sync/atomic"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/unicornultrafoundation/go-helios/common/bigendian"
    13  	"github.com/unicornultrafoundation/go-helios/u2udb"
    14  	"github.com/unicornultrafoundation/go-helios/u2udb/memorydb"
    15  
    16  	"github.com/unicornultrafoundation/go-u2u/utils/dbutil/dbcounter"
    17  )
    18  
    19  func decodePair(b []byte) (uint32, uint32) {
    20  	v1 := bigendian.BytesToUint32(b[:4])
    21  	v2 := bigendian.BytesToUint32(b[4:])
    22  	return v1, v2
    23  }
    24  
    25  type UncallableAfterRelease struct {
    26  	u2udb.Snapshot
    27  	iterators []*uncallableAfterReleaseIterator
    28  	mu        sync.Mutex
    29  }
    30  
    31  type uncallableAfterReleaseIterator struct {
    32  	u2udb.Iterator
    33  }
    34  
    35  func (db *UncallableAfterRelease) NewIterator(prefix []byte, start []byte) u2udb.Iterator {
    36  	it := db.Snapshot.NewIterator(prefix, start)
    37  	wrapped := &uncallableAfterReleaseIterator{it}
    38  
    39  	db.mu.Lock()
    40  	defer db.mu.Unlock()
    41  	db.iterators = append(db.iterators, wrapped)
    42  
    43  	return wrapped
    44  }
    45  
    46  func (db *UncallableAfterRelease) Release() {
    47  	db.Snapshot.Release()
    48  	// ensure nil pointer exception on any next call
    49  	db.Snapshot = nil
    50  
    51  	db.mu.Lock()
    52  	defer db.mu.Unlock()
    53  	for _, it := range db.iterators {
    54  		it.Iterator = nil
    55  	}
    56  }
    57  
    58  func TestSnapshot_SwitchTo(t *testing.T) {
    59  	require := require.New(t)
    60  
    61  	const prefixes = 100
    62  	const keys = 100
    63  	const checkers = 5
    64  	const switchers = 5
    65  	const duration = time.Millisecond * 400
    66  
    67  	// fill DB with data
    68  	memdb := dbcounter.WrapStore(memorydb.New(), "", false)
    69  	for i := uint32(0); i < prefixes; i++ {
    70  		for j := uint32(0); j < keys; j++ {
    71  			key := append(bigendian.Uint32ToBytes(i), bigendian.Uint32ToBytes(j)...)
    72  			val := append(bigendian.Uint32ToBytes(i*i), bigendian.Uint32ToBytes(j*j)...)
    73  			require.NoError(memdb.Put(key, val))
    74  		}
    75  	}
    76  
    77  	// 4 readers, one interrupter
    78  	snap, err := memdb.GetSnapshot()
    79  	require.NoError(err)
    80  	switchable := Wrap(&UncallableAfterRelease{
    81  		Snapshot: snap,
    82  	})
    83  
    84  	stop := uint32(0)
    85  	wg := sync.WaitGroup{}
    86  	wg.Add(checkers + switchers)
    87  	for worker := 0; worker < checkers; worker++ {
    88  		go func() {
    89  			defer wg.Done()
    90  			for atomic.LoadUint32(&stop) == 0 {
    91  				var prevPrefix uint32
    92  				var prevKey uint32
    93  				prefixN := rand.Uint32() % prefixes
    94  				prefix := bigendian.Uint32ToBytes(prefixN)
    95  				keyN := rand.Uint32() % prefixes
    96  				start := bigendian.Uint32ToBytes(keyN)
    97  				prevKey = keyN - 1
    98  				if rand.Intn(10) == 0 {
    99  					start = nil
   100  					prevKey = 0
   101  					prevKey--
   102  					if rand.Intn(2) == 0 {
   103  						prefix = nil
   104  						prevPrefix = 0
   105  					}
   106  				}
   107  				it := switchable.NewIterator(prefix, start)
   108  				require.NoError(it.Error())
   109  				for it.Next() {
   110  					require.NoError(it.Error())
   111  					require.Equal(8, len(it.Key()))
   112  					require.Equal(8, len(it.Value()))
   113  					p, k := decodePair(it.Key())
   114  					sp, sk := decodePair(it.Value())
   115  					require.Equal(p*p, sp)
   116  					require.Equal(k*k, sk)
   117  					if prefix != nil {
   118  						require.Equal(prefixN, p)
   119  					} else if p != prevPrefix {
   120  						require.Equal(prevPrefix+1, p)
   121  						prevPrefix = p
   122  						prevKey = 0
   123  						prevKey--
   124  					}
   125  
   126  					require.Equal(prevKey+1, k, prefix)
   127  					prevKey = k
   128  				}
   129  				require.NoError(it.Error())
   130  				it.Release()
   131  			}
   132  		}()
   133  	}
   134  	for worker := 0; worker < switchers; worker++ {
   135  		go func() {
   136  			defer wg.Done()
   137  
   138  			for atomic.LoadUint32(&stop) == 0 {
   139  				snap, err := memdb.GetSnapshot()
   140  				require.NoError(err)
   141  				old := switchable.SwitchTo(&UncallableAfterRelease{
   142  					Snapshot: snap,
   143  				})
   144  				old.Release()
   145  			}
   146  		}()
   147  	}
   148  	time.Sleep(duration)
   149  	atomic.StoreUint32(&stop, 1)
   150  	wg.Wait()
   151  	switchable.Release()
   152  	require.NoError(memdb.Close())
   153  }