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 }