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  }