From 384d899eff0b958b78bdb9bc3ec855342977d942 Mon Sep 17 00:00:00 2001
From: Solovyov1796 <xinyu@0g.ai>
Date: Mon, 10 Mar 2025 20:34:02 +0800
Subject: [PATCH] add mempool tx replace callback

---
 app/priority_nonce.go | 17 ++++++++++++++++-
 cmd/0gchaind/app.go   |  4 ++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/app/priority_nonce.go b/app/priority_nonce.go
index dbbd0027..f083fccd 100644
--- a/app/priority_nonce.go
+++ b/app/priority_nonce.go
@@ -38,6 +38,8 @@ type PriorityNonceMempool struct {
 	maxTx           int
 	counterBySender map[string]int
 	txRecord        map[txMeta]struct{}
+
+	txReplacedCallback func(ctx context.Context, oldTx, newTx sdk.Tx)
 }
 
 type PriorityNonceIterator struct {
@@ -124,6 +126,12 @@ func PriorityNonceWithMaxTx(maxTx int) PriorityNonceMempoolOption {
 	}
 }
 
+func PriorityNonceWithTxReplacedCallback(cb func(ctx context.Context, oldTx, newTx sdk.Tx)) PriorityNonceMempoolOption {
+	return func(mp *PriorityNonceMempool) {
+		mp.txReplacedCallback = cb
+	}
+}
+
 // DefaultPriorityMempool returns a priorityNonceMempool with no options.
 func DefaultPriorityMempool() mempool.Mempool {
 	return NewPriorityMempool()
@@ -184,6 +192,7 @@ func (mp *PriorityNonceMempool) Insert(ctx context.Context, tx sdk.Tx) error {
 	sdkContext := sdk.UnwrapSDKContext(ctx)
 	priority := sdkContext.Priority()
 
+	var replacedTx sdk.Tx
 	var sender string
 	var nonce uint64
 
@@ -244,6 +253,7 @@ func (mp *PriorityNonceMempool) Insert(ctx context.Context, tx sdk.Tx) error {
 	//
 	// This O(log n) remove operation is rare and only happens when a tx's priority
 	// changes.
+
 	sk := txMeta{nonce: nonce, sender: sender}
 	if oldScore, txExists := mp.scores[sk]; txExists {
 		if mp.txReplacement != nil && !mp.txReplacement(oldScore.priority, priority, senderIndex.Get(key).Value.(sdk.Tx), tx) {
@@ -256,12 +266,13 @@ func (mp *PriorityNonceMempool) Insert(ctx context.Context, tx sdk.Tx) error {
 			)
 		}
 
-		mp.priorityIndex.Remove(txMeta{
+		e := mp.priorityIndex.Remove(txMeta{
 			nonce:    nonce,
 			sender:   sender,
 			priority: oldScore.priority,
 			weight:   oldScore.weight,
 		})
+		replacedTx = e.Value.(sdk.Tx)
 		mp.priorityCounts[oldScore.priority]--
 	}
 
@@ -274,6 +285,10 @@ func (mp *PriorityNonceMempool) Insert(ctx context.Context, tx sdk.Tx) error {
 	mp.scores[sk] = txMeta{priority: priority}
 	mp.priorityIndex.Set(key, tx)
 
+	if mp.txReplacedCallback != nil && replacedTx != nil {
+		mp.txReplacedCallback(ctx, replacedTx, tx)
+	}
+
 	return nil
 }
 
diff --git a/cmd/0gchaind/app.go b/cmd/0gchaind/app.go
index 96c54f64..03abd93a 100644
--- a/cmd/0gchaind/app.go
+++ b/cmd/0gchaind/app.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -129,6 +130,9 @@ func (ac appCreator) newApp(
 
 	mempool := app.NewPriorityMempool(
 		app.PriorityNonceWithMaxTx(cast.ToInt(appOpts.Get(server.FlagMempoolMaxTxs))),
+		app.PriorityNonceWithTxReplacedCallback(func(ctx context.Context, oldTx, newTx sdk.Tx) {
+			bApp.RegisterMempoolTxReplacedEvent(ctx, oldTx, newTx)
+		}),
 	)
 	bApp.SetMempool(mempool)