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

     1  package cat
     2  
     3  import (
     4  	"sync"
     5  	"time"
     6  
     7  	"github.com/badrootd/celestia-core/types"
     8  )
     9  
    10  // simple, thread-safe in memory store for transactions
    11  type store struct {
    12  	mtx   sync.RWMutex
    13  	bytes int64
    14  	txs   map[types.TxKey]*wrappedTx
    15  }
    16  
    17  func newStore() *store {
    18  	return &store{
    19  		bytes: 0,
    20  		txs:   make(map[types.TxKey]*wrappedTx),
    21  	}
    22  }
    23  
    24  func (s *store) set(wtx *wrappedTx) bool {
    25  	if wtx == nil {
    26  		return false
    27  	}
    28  	s.mtx.Lock()
    29  	defer s.mtx.Unlock()
    30  	if tx, exists := s.txs[wtx.key]; !exists || tx.height == -1 {
    31  		s.txs[wtx.key] = wtx
    32  		s.bytes += wtx.size()
    33  		return true
    34  	}
    35  	return false
    36  }
    37  
    38  func (s *store) get(txKey types.TxKey) *wrappedTx {
    39  	s.mtx.RLock()
    40  	defer s.mtx.RUnlock()
    41  	return s.txs[txKey]
    42  }
    43  
    44  func (s *store) has(txKey types.TxKey) bool {
    45  	s.mtx.RLock()
    46  	defer s.mtx.RUnlock()
    47  	_, has := s.txs[txKey]
    48  	return has
    49  }
    50  
    51  func (s *store) remove(txKey types.TxKey) bool {
    52  	s.mtx.Lock()
    53  	defer s.mtx.Unlock()
    54  	tx, exists := s.txs[txKey]
    55  	if !exists {
    56  		return false
    57  	}
    58  	s.bytes -= tx.size()
    59  	delete(s.txs, txKey)
    60  	return true
    61  }
    62  
    63  // reserve adds an empty placeholder for the specified key to prevent
    64  // a transaction with the same key from being added
    65  func (s *store) reserve(txKey types.TxKey) bool {
    66  	s.mtx.Lock()
    67  	defer s.mtx.Unlock()
    68  	_, has := s.txs[txKey]
    69  	if !has {
    70  		s.txs[txKey] = &wrappedTx{height: -1}
    71  		return true
    72  	}
    73  	return false
    74  }
    75  
    76  // release is called when a pending transaction failed
    77  // to enter the mempool. The empty element and key is removed.
    78  func (s *store) release(txKey types.TxKey) {
    79  	s.mtx.Lock()
    80  	defer s.mtx.Unlock()
    81  	value, ok := s.txs[txKey]
    82  	if ok && value.height == -1 {
    83  		delete(s.txs, txKey)
    84  	}
    85  }
    86  
    87  func (s *store) size() int {
    88  	s.mtx.RLock()
    89  	defer s.mtx.RUnlock()
    90  	return len(s.txs)
    91  }
    92  
    93  func (s *store) totalBytes() int64 {
    94  	s.mtx.RLock()
    95  	defer s.mtx.RUnlock()
    96  	return s.bytes
    97  }
    98  
    99  func (s *store) getAllKeys() []types.TxKey {
   100  	s.mtx.RLock()
   101  	defer s.mtx.RUnlock()
   102  	keys := make([]types.TxKey, len(s.txs))
   103  	idx := 0
   104  	for key := range s.txs {
   105  		keys[idx] = key
   106  		idx++
   107  	}
   108  	return keys
   109  }
   110  
   111  func (s *store) getAllTxs() []*wrappedTx {
   112  	s.mtx.RLock()
   113  	defer s.mtx.RUnlock()
   114  	txs := make([]*wrappedTx, len(s.txs))
   115  	idx := 0
   116  	for _, tx := range s.txs {
   117  		txs[idx] = tx
   118  		idx++
   119  	}
   120  	return txs
   121  }
   122  
   123  func (s *store) getTxsBelowPriority(priority int64) ([]*wrappedTx, int64) {
   124  	s.mtx.RLock()
   125  	defer s.mtx.RUnlock()
   126  	txs := make([]*wrappedTx, 0, len(s.txs))
   127  	bytes := int64(0)
   128  	for _, tx := range s.txs {
   129  		if tx.priority < priority {
   130  			txs = append(txs, tx)
   131  			bytes += tx.size()
   132  		}
   133  	}
   134  	return txs, bytes
   135  }
   136  
   137  // purgeExpiredTxs removes all transactions that are older than the given height
   138  // and time. Returns the amount of transactions that were removed
   139  func (s *store) purgeExpiredTxs(expirationHeight int64, expirationAge time.Time) int {
   140  	s.mtx.Lock()
   141  	defer s.mtx.Unlock()
   142  	counter := 0
   143  	for key, tx := range s.txs {
   144  		if tx.height < expirationHeight || tx.timestamp.Before(expirationAge) {
   145  			s.bytes -= tx.size()
   146  			delete(s.txs, key)
   147  			counter++
   148  		}
   149  	}
   150  	return counter
   151  }
   152  
   153  func (s *store) reset() {
   154  	s.mtx.Lock()
   155  	defer s.mtx.Unlock()
   156  	s.bytes = 0
   157  	s.txs = make(map[types.TxKey]*wrappedTx)
   158  }