services: improve Notary service locking

1. Allow parallel handling of notary requests with different main
   transaction.
2. Move expired request removal to PostBlock.

Signed-off-by: Anna Shaleva <[email protected]>

Author: AnnaShaleva

pkg/services/notary/notary.go

@@ -52,7 +52,8 @@ type (
 		// started is a status bool to protect from double start/shutdown.
 		started atomic.Bool
 
-		// reqMtx protects requests list.
+		// reqMtx protects the request list from concurrent requests addition/removal.
+		// Use per-request locks instead of this one to perform request-changing operations.
 		reqMtx sync.RWMutex
 		// requests represents a map of main transactions which needs to be completed
 		// with the associated fallback transactions grouped by the main transaction hash
@@ -89,6 +90,7 @@ const defaultTxChannelCapacity = 100
 type (
 	// request represents Notary service request.
 	request struct {
+		lock sync.RWMutex
 		// isSent indicates whether the main transaction was successfully sent to the network.
 		isSent bool
 		main   *transaction.Transaction
@@ -117,6 +119,7 @@ type (
 )
 
 // isMainCompleted denotes whether all signatures for the main transaction were collected.
+// The caller is supposed to hold the request lock.
 func (r request) isMainCompleted() bool {
 	if r.witnessInfo == nil {
 		return false
@@ -254,12 +257,14 @@ func (n *Notary) OnNewRequest(payload *payload.P2PNotaryRequest) {
 			zap.String("verification error", validationErr.Error()))
 	}
 	n.reqMtx.Lock()
-	defer n.reqMtx.Unlock()
 	r, exists := n.requests[payload.MainTransaction.Hash()]
 	if exists {
+		r.lock.Lock() // RLock doesn't fit here since we modify r.minNotValidBefore below.
 		if slices.ContainsFunc(r.fallbacks, func(fb *transaction.Transaction) bool {
 			return fb.Hash().Equals(payload.FallbackTransaction.Hash())
 		}) {
+			r.lock.Unlock()
+			n.reqMtx.Unlock()
 			return // then we already have processed this request
 		}
 		r.minNotValidBefore = min(r.minNotValidBefore, nvbFallback)
@@ -270,8 +275,10 @@ func (n *Notary) OnNewRequest(payload *payload.P2PNotaryRequest) {
 			main:              payload.MainTransaction.Copy(),
 			minNotValidBefore: nvbFallback,
 		}
+		r.lock.Lock()
 		n.requests[payload.MainTransaction.Hash()] = r
 	}
+	n.reqMtx.Unlock()
 	if r.witnessInfo == nil && validationErr == nil {
 		r.witnessInfo = newInfo
 	}
@@ -282,8 +289,14 @@ func (n *Notary) OnNewRequest(payload *payload.P2PNotaryRequest) {
 	// the copy.
 	r.fallbacks = append(r.fallbacks, payload.FallbackTransaction.Copy())
 	if exists && r.isMainCompleted() || validationErr != nil {
+		r.lock.Unlock()
 		return
 	}
+
+	// @roman-khimov, I tried to work with r.copy() here in order not to hold r.lock during
+	// witness verification. However, the race is possible in the end of this method if main
+	// transaction will be completed by another thread, which leads to the situation when
+	// two main transactions with different set of witnesses will be submitted.
 	mainHash := hash.NetSha256(uint32(n.Network), r.main).BytesBE()
 	for i, w := range payload.MainTransaction.Scripts {
 		if len(w.InvocationScript) == 0 || // check that signature for this witness was provided
@@ -338,6 +351,7 @@ func (n *Notary) OnNewRequest(payload *payload.P2PNotaryRequest) {
 				zap.Error(err))
 		}
 	}
+	r.lock.Unlock()
 }
 
 // OnRequestRemoval is a callback which is called after fallback transaction is removed
@@ -348,20 +362,20 @@ func (n *Notary) OnRequestRemoval(pld *payload.P2PNotaryRequest) {
 	}
 
 	n.reqMtx.Lock()
-	defer n.reqMtx.Unlock()
 	r, ok := n.requests[pld.MainTransaction.Hash()]
+	n.reqMtx.Unlock()
 	if !ok {
 		return
 	}
+
+	r.lock.Lock()
 	for i, fb := range r.fallbacks {
 		if fb.Hash().Equals(pld.FallbackTransaction.Hash()) {
 			r.fallbacks = append(r.fallbacks[:i], r.fallbacks[i+1:]...)
 			break
 		}
 	}
-	if len(r.fallbacks) == 0 {
-		delete(n.requests, r.main.Hash())
-	}
+	r.lock.Unlock()
 }
 
 // PostPersist is a callback which is called after a new block event is received.
@@ -376,15 +390,24 @@ func (n *Notary) PostPersist() {
 	}
 
 	n.reqMtx.Lock()
+	// @roman-khimov, an option is to copy n.requests in order not to hold lock over n.reqMtx for the whole cycle,
+	// but we need to take a write lock over every request anyway since it's modified by n.finalize.
+	// We also need to clean the list of requests with 0 fallbacks, this requires n.reqMtx to be taken.
 	defer n.reqMtx.Unlock()
 	currHeight := n.Config.Chain.BlockHeight()
 	for h, r := range n.requests {
+		r.lock.Lock()
+		if len(r.fallbacks) == 0 {
+			delete(n.requests, r.main.Hash())
+			continue
+		}
 		if !r.isSent && r.isMainCompleted() && r.minNotValidBefore > currHeight {
 			if err := n.finalize(acc, r.main, h); err != nil {
 				n.Config.Log.Error("failed to finalize main transaction after PostPersist, waiting for the next block to retry",
 					zap.String("hash", r.main.Hash().StringLE()),
 					zap.Error(err))
 			}
+			r.lock.Unlock()
 			continue
 		}
 		if r.minNotValidBefore <= currHeight { // then at least one of the fallbacks can already be sent.
@@ -400,6 +423,7 @@ func (n *Notary) PostPersist() {
 				}
 			}
 		}
+		r.lock.Unlock()
 	}
 }
 
@@ -448,33 +472,37 @@ func (n *Notary) newTxCallbackLoop() {
 		case tx := <-n.newTxs:
 			isMain := tx.tx.Hash() == tx.mainHash
 
-			n.reqMtx.Lock()
+			n.reqMtx.RLock()
 			r, ok := n.requests[tx.mainHash]
-			if !ok || isMain && (r.isSent || r.minNotValidBefore <= n.Config.Chain.BlockHeight()) {
-				n.reqMtx.Unlock()
+			n.reqMtx.RUnlock()
+			if !ok {
 				continue
 			}
+			r.lock.Lock()
+			if isMain && (r.isSent || r.minNotValidBefore <= n.Config.Chain.BlockHeight()) {
+				r.lock.Unlock()
+				continue
+			}
+
 			if !isMain {
 				// Ensure that fallback was not already completed.
 				var isPending = slices.ContainsFunc(r.fallbacks, func(fb *transaction.Transaction) bool {
 					return fb.Hash() == tx.tx.Hash()
 				})
 				if !isPending {
-					n.reqMtx.Unlock()
+					r.lock.Unlock()
 					continue
 				}
 			}
 
-			n.reqMtx.Unlock()
 			err := n.onTransaction(tx.tx)
 			if err != nil {
 				n.Config.Log.Error("new transaction callback finished with error",
 					zap.Error(err),
 					zap.Bool("is main", isMain))
+				r.lock.Unlock()
 				continue
 			}
-
-			n.reqMtx.Lock()
 			if isMain {
 				r.isSent = true
 			} else {
@@ -484,11 +512,8 @@ func (n *Notary) newTxCallbackLoop() {
 						break
 					}
 				}
-				if len(r.fallbacks) == 0 {
-					delete(n.requests, tx.mainHash)
-				}
 			}
-			n.reqMtx.Unlock()
+			r.lock.Unlock()
 		case <-n.stopCh:
 			return
 		}