github.com/badrootd/celestia-core@v0.0.0-20240305091328-aa4207a4b25d/mempool/cat/store_test.go (about)

     1  package cat
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/badrootd/celestia-core/types"
    11  	"github.com/stretchr/testify/require"
    12  )
    13  
    14  func TestStoreSimple(t *testing.T) {
    15  	store := newStore()
    16  
    17  	tx := types.Tx("tx1")
    18  	key := tx.Key()
    19  	wtx := newWrappedTx(tx, key, 1, 1, 1, "")
    20  
    21  	// asset zero state
    22  	require.Nil(t, store.get(key))
    23  	require.False(t, store.has(key))
    24  	require.False(t, store.remove(key))
    25  	require.Zero(t, store.size())
    26  	require.Zero(t, store.totalBytes())
    27  	require.Empty(t, store.getAllKeys())
    28  	require.Empty(t, store.getAllTxs())
    29  
    30  	// add a tx
    31  	store.set(wtx)
    32  	require.True(t, store.has(key))
    33  	require.Equal(t, wtx, store.get(key))
    34  	require.Equal(t, int(1), store.size())
    35  	require.Equal(t, wtx.size(), store.totalBytes())
    36  
    37  	// remove a tx
    38  	store.remove(key)
    39  	require.False(t, store.has(key))
    40  	require.Nil(t, store.get(key))
    41  	require.Zero(t, store.size())
    42  	require.Zero(t, store.totalBytes())
    43  }
    44  
    45  func TestStoreReservingTxs(t *testing.T) {
    46  	store := newStore()
    47  
    48  	tx := types.Tx("tx1")
    49  	key := tx.Key()
    50  	wtx := newWrappedTx(tx, key, 1, 1, 1, "")
    51  
    52  	// asset zero state
    53  	store.release(key)
    54  
    55  	// reserve a tx
    56  	store.reserve(key)
    57  	require.True(t, store.has(key))
    58  	// should not update the total bytes
    59  	require.Zero(t, store.totalBytes())
    60  
    61  	// should be able to add a tx
    62  	store.set(wtx)
    63  	require.Equal(t, tx, store.get(key).tx)
    64  	require.Equal(t, wtx.size(), store.totalBytes())
    65  
    66  	// releasing should do nothing on a set tx
    67  	store.release(key)
    68  	require.True(t, store.has(key))
    69  	require.Equal(t, tx, store.get(key).tx)
    70  
    71  	store.remove(key)
    72  	require.False(t, store.has(key))
    73  
    74  	// reserve the tx again
    75  	store.reserve(key)
    76  	require.True(t, store.has(key))
    77  
    78  	// release should remove the tx
    79  	store.release(key)
    80  	require.False(t, store.has(key))
    81  }
    82  
    83  func TestStoreConcurrentAccess(t *testing.T) {
    84  	store := newStore()
    85  
    86  	numTxs := 100
    87  
    88  	wg := &sync.WaitGroup{}
    89  	for i := 0; i < numTxs; i++ {
    90  		wg.Add(1)
    91  		go func(i int) {
    92  			defer wg.Done()
    93  			ticker := time.NewTicker(10 * time.Millisecond)
    94  			for range ticker.C {
    95  				tx := types.Tx(fmt.Sprintf("tx%d", i%(numTxs/10)))
    96  				key := tx.Key()
    97  				wtx := newWrappedTx(tx, key, 1, 1, 1, "")
    98  				existingTx := store.get(key)
    99  				if existingTx != nil && bytes.Equal(existingTx.tx, tx) {
   100  					// tx has already been added
   101  					return
   102  				}
   103  				if store.reserve(key) {
   104  					// some fail
   105  					if i%3 == 0 {
   106  						store.release(key)
   107  						return
   108  					}
   109  					store.set(wtx)
   110  					// this should be a noop
   111  					store.release(key)
   112  					return
   113  				}
   114  				// already reserved so we retry in 10 milliseconds
   115  			}
   116  		}(i)
   117  	}
   118  	wg.Wait()
   119  
   120  	require.Equal(t, numTxs/10, store.size())
   121  }
   122  
   123  func TestStoreGetTxs(t *testing.T) {
   124  	store := newStore()
   125  
   126  	numTxs := 100
   127  	for i := 0; i < numTxs; i++ {
   128  		tx := types.Tx(fmt.Sprintf("tx%d", i))
   129  		key := tx.Key()
   130  		wtx := newWrappedTx(tx, key, 1, 1, int64(i), "")
   131  		store.set(wtx)
   132  	}
   133  
   134  	require.Equal(t, numTxs, store.size())
   135  
   136  	// get all txs
   137  	txs := store.getAllTxs()
   138  	require.Equal(t, numTxs, len(txs))
   139  
   140  	// get txs by keys
   141  	keys := store.getAllKeys()
   142  	require.Equal(t, numTxs, len(keys))
   143  
   144  	// get txs below a certain priority
   145  	txs, bz := store.getTxsBelowPriority(int64(numTxs / 2))
   146  	require.Equal(t, numTxs/2, len(txs))
   147  	var actualBz int64
   148  	for _, tx := range txs {
   149  		actualBz += tx.size()
   150  	}
   151  	require.Equal(t, actualBz, bz)
   152  }
   153  
   154  func TestStoreExpiredTxs(t *testing.T) {
   155  	store := newStore()
   156  	numTxs := 100
   157  	for i := 0; i < numTxs; i++ {
   158  		tx := types.Tx(fmt.Sprintf("tx%d", i))
   159  		key := tx.Key()
   160  		wtx := newWrappedTx(tx, key, int64(i), 1, 1, "")
   161  		store.set(wtx)
   162  	}
   163  
   164  	// half of them should get purged
   165  	store.purgeExpiredTxs(int64(numTxs/2), time.Time{})
   166  
   167  	remainingTxs := store.getAllTxs()
   168  	require.Equal(t, numTxs/2, len(remainingTxs))
   169  	for _, tx := range remainingTxs {
   170  		require.GreaterOrEqual(t, tx.height, int64(numTxs/2))
   171  	}
   172  
   173  	store.purgeExpiredTxs(int64(0), time.Now().Add(time.Second))
   174  	require.Empty(t, store.getAllTxs())
   175  }