github.com/badrootd/celestia-core@v0.0.0-20240305091328-aa4207a4b25d/mempool/cat/requests.go (about) 1 package cat 2 3 import ( 4 "sync" 5 "time" 6 7 "github.com/badrootd/celestia-core/types" 8 ) 9 10 const defaultGlobalRequestTimeout = 1 * time.Hour 11 12 // requestScheduler tracks the lifecycle of outbound transaction requests. 13 type requestScheduler struct { 14 mtx sync.Mutex 15 16 // responseTime is the time the scheduler 17 // waits for a response from a peer before 18 // invoking the callback 19 responseTime time.Duration 20 21 // globalTimeout represents the longest duration 22 // to wait for any late response (after the reponseTime). 23 // After this period the request is garbage collected. 24 globalTimeout time.Duration 25 26 // requestsByPeer is a lookup table of requests by peer. 27 // Multiple transactions can be requested by a single peer at one 28 requestsByPeer map[uint16]requestSet 29 30 // requestsByTx is a lookup table for requested txs. 31 // There can only be one request per tx. 32 requestsByTx map[types.TxKey]uint16 33 } 34 35 type requestSet map[types.TxKey]*time.Timer 36 37 func newRequestScheduler(responseTime, globalTimeout time.Duration) *requestScheduler { 38 return &requestScheduler{ 39 responseTime: responseTime, 40 globalTimeout: globalTimeout, 41 requestsByPeer: make(map[uint16]requestSet), 42 requestsByTx: make(map[types.TxKey]uint16), 43 } 44 } 45 46 func (r *requestScheduler) Add(key types.TxKey, peer uint16, onTimeout func(key types.TxKey)) bool { 47 if peer == 0 { 48 return false 49 } 50 r.mtx.Lock() 51 defer r.mtx.Unlock() 52 53 // not allowed to have more than one outgoing transaction at once 54 if _, ok := r.requestsByTx[key]; ok { 55 return false 56 } 57 58 timer := time.AfterFunc(r.responseTime, func() { 59 r.mtx.Lock() 60 delete(r.requestsByTx, key) 61 r.mtx.Unlock() 62 63 // trigger callback. Callback can `Add` the tx back to the scheduler 64 if onTimeout != nil { 65 onTimeout(key) 66 } 67 68 // We set another timeout because the peer could still send 69 // a late response after the first timeout and it's important 70 // to recognise that it is a transaction in response to a 71 // request and not a new transaction being broadcasted to the entire 72 // network. This timer cannot be stopped and is used to ensure 73 // garbage collection. 74 time.AfterFunc(r.globalTimeout, func() { 75 r.mtx.Lock() 76 defer r.mtx.Unlock() 77 delete(r.requestsByPeer[peer], key) 78 }) 79 }) 80 if _, ok := r.requestsByPeer[peer]; !ok { 81 r.requestsByPeer[peer] = requestSet{key: timer} 82 } else { 83 r.requestsByPeer[peer][key] = timer 84 } 85 r.requestsByTx[key] = peer 86 return true 87 } 88 89 func (r *requestScheduler) ForTx(key types.TxKey) uint16 { 90 r.mtx.Lock() 91 defer r.mtx.Unlock() 92 93 return r.requestsByTx[key] 94 } 95 96 func (r *requestScheduler) Has(peer uint16, key types.TxKey) bool { 97 r.mtx.Lock() 98 defer r.mtx.Unlock() 99 100 requestSet, ok := r.requestsByPeer[peer] 101 if !ok { 102 return false 103 } 104 _, ok = requestSet[key] 105 return ok 106 } 107 108 func (r *requestScheduler) ClearAllRequestsFrom(peer uint16) requestSet { 109 r.mtx.Lock() 110 defer r.mtx.Unlock() 111 112 requests, ok := r.requestsByPeer[peer] 113 if !ok { 114 return requestSet{} 115 } 116 for tx, timer := range requests { 117 timer.Stop() 118 delete(r.requestsByTx, tx) 119 } 120 delete(r.requestsByPeer, peer) 121 return requests 122 } 123 124 func (r *requestScheduler) MarkReceived(peer uint16, key types.TxKey) bool { 125 r.mtx.Lock() 126 defer r.mtx.Unlock() 127 128 if _, ok := r.requestsByPeer[peer]; !ok { 129 return false 130 } 131 132 if timer, ok := r.requestsByPeer[peer][key]; ok { 133 timer.Stop() 134 } else { 135 return false 136 } 137 138 delete(r.requestsByPeer[peer], key) 139 delete(r.requestsByTx, key) 140 return true 141 } 142 143 // Close stops all timers and clears all requests. 144 // Add should never be called after `Close`. 145 func (r *requestScheduler) Close() { 146 r.mtx.Lock() 147 defer r.mtx.Unlock() 148 149 for _, requestSet := range r.requestsByPeer { 150 for _, timer := range requestSet { 151 timer.Stop() 152 } 153 } 154 }